face-clustering-rk3588./src/face_scan.cc

#include <cmath>
#include <cstddef>
#include <iomanip>
#include <iostream>
#include <string>
#include <vector>

#include "face_engine.h"
#include "face_config.h"

static constexpr std::size_t EMBEDDING_SIZE = 128;

static void write_json_string(const std::string& value)
{
    std::cout << '"';

    for (unsigned char ch : value) {
        switch (ch) {
        case '"':
            std::cout << "\\\"";
            break;

        case '\\':
            std::cout << "\\\\";
            break;

        case '\b':
            std::cout << "\\b";
            break;

        case '\f':
            std::cout << "\\f";
            break;

        case '\n':
            std::cout << "\\n";
            break;

        case '\r':
            std::cout << "\\r";
            break;

        case '\t':
            std::cout << "\\t";
            break;

        default:
            if (ch < 0x20) {
                const char* hex = "0123456789abcdef";

                std::cout
                    << "\\u00"
                    << hex[(ch >> 4) & 0x0f]
                    << hex[ch & 0x0f];
            } else {
                std::cout << static_cast<char>(ch);
            }
        }
    }

    std::cout << '"';
}

static bool finite_face(const Face& face)
{
    if (!std::isfinite(face.score) ||
        !std::isfinite(face.x1) ||
        !std::isfinite(face.y1) ||
        !std::isfinite(face.x2) ||
        !std::isfinite(face.y2)) {
        return false;
    }

    for (const Point& point : face.kps) {
        if (!std::isfinite(point.x) ||
            !std::isfinite(point.y)) {
            return false;
        }
    }

    if (face.embedding.size() != EMBEDDING_SIZE)
        return false;

    for (float value : face.embedding) {
        if (!std::isfinite(value))
            return false;
    }

    return true;
}

static void write_face_json(
    const Face& face,
    std::size_t face_index)
{
    const float width = face.x2 - face.x1;
    const float height = face.y2 - face.y1;

    std::cout << "{";

    std::cout
        << "\"face_index\":"
        << face_index
        << ",";

    std::cout
        << "\"score\":"
        << face.score
        << ",";

    std::cout << "\"bbox\":{";

    std::cout
        << "\"x\":"
        << face.x1
        << ",";

    std::cout
        << "\"y\":"
        << face.y1
        << ",";

    std::cout
        << "\"width\":"
        << width
        << ",";

    std::cout
        << "\"height\":"
        << height;

    std::cout << "},";

    std::cout << "\"landmarks\":[";

    for (int i = 0; i < 5; ++i) {
        if (i != 0)
            std::cout << ",";

        std::cout
            << "["
            << face.kps[i].x
            << ","
            << face.kps[i].y
            << "]";
    }

    std::cout << "],";

    std::cout << "\"embedding\":[";

    for (std::size_t i = 0;
         i < face.embedding.size();
         ++i) {
        if (i != 0)
            std::cout << ",";

        std::cout << face.embedding[i];
    }

    std::cout << "]";

    std::cout << "}";
}

int main(int argc, char** argv)
{
    if (argc != 2) {
        std::cerr
            << "Uso: "
            << argv[0]
            << " image.jpg\n";

        return 2;
    }

    const std::string root =
        FACE_CLUSTERING_ROOT;

    const std::string scrfd_model =
        root + "/models/SCRFD_500M_KPS_640.rknn";

    const std::string sface_model =
        root + "/models/face_recognition_sface_2021dec.rknn";

    FaceEngine engine;

    if (!engine.init(scrfd_model, sface_model)) {
        std::cerr
            << "Impossibile inizializzare FaceEngine\n";

        return 3;
    }

    std::vector<Face> faces;
    int image_width = 0;
    int image_height = 0;

    if (!engine.process_image(
            argv[1],
            faces,
            image_width,
            image_height)) {
        std::cerr
            << "Scansione fallita: "
            << argv[1]
            << "\n";

        return 4;
    }

    for (std::size_t i = 0; i < faces.size(); ++i) {
        if (!finite_face(faces[i])) {
            std::cerr
                << "Risultato non valido per face_index="
                << i
                << "\n";

            return 5;
        }
    }

    std::cout << std::setprecision(9);

    std::cout << "{";

    std::cout << "\"image\":{";

    std::cout << "\"path\":";
    write_json_string(argv[1]);

    std::cout
        << ",\"width\":"
        << image_width
        << ",\"height\":"
        << image_height;

    std::cout << "},";

    std::cout << "\"faces\":[";

    for (std::size_t i = 0; i < faces.size(); ++i) {
        if (i != 0)
            std::cout << ",";

        write_face_json(faces[i], i);
    }

    std::cout << "]";

    std::cout << "}\n";

    return 0;
}