example_face_identify.cpp

Go to the documentation of this file.

 
#include "sfe_toolkit/sfe_core.h"
#include "sfe_toolkit/sfe_face.h"
#include <regex>
#include <vector>
 
#include <iomanip>
#include <sstream>
 
#include "annotate.hpp"
#include "solvers.h"
#include "utils.hpp"
#include <unordered_map>
 
std::string image_gallery = "./assets/face/entities/";
 
std::string image_probe = "assets/face/images/obiwan0.png";
 
std::string solver_face_detect = SOLVER_FACE_DETECT;
std::string solver_face_landmarks = SOLVER_FACE_LANDMARKS;
std::string solver_face_template = SOLVER_FACE_EXTRACTION;
 
const auto SOLVER_PARAMETERS = std::vector<SFESolverParameter>{};
 
size_t face_size_min = 28;
size_t face_size_max = 170;
 
float detection_threshold = 0.1f;
float identification_threshold = 0.6f;
 
void getRecommendedImageSize(const std::string &solver_face_detect,
                             SFEImage &image, const size_t face_size_min,
                             const size_t face_size_max,
                             size_t &recommended_width,
                             size_t &recommended_height) {
 
  // If the face detection solver requires static input (filename contains width
  // and height), we need to prepare the input image accordingly Typically the
  // format is: face_detect_accurate_mask_w1920h1080_op11.onnxrt.solver NOTE:
  // This can be hardcoded into the application, or we can parse the width and
  // height from the solver filename
 
  // Try to use regex capture to extract width and height from solver name
  std::smatch width_height;
  if (std::regex_match(solver_face_detect, width_height,
                       std::regex(".*w([0-9]+)h([0-9]+).*"))) {
    recommended_width = std::stoi(width_height[1]);
    recommended_height = std::stoi(width_height[2]);
  } else {
    // Solvers without width and height in name can accept dynamic input
    // image size. For best results we recommend to scale the input image to
    // a resolution recommended by the sfeFaceDetectInputSize function
 
    // Use sfeFaceDetectInputSize to get recommended input image dimensions
    // for given min_face_size/max_face_size
    SFEDetectionInputSize input_size{};
    SFEError error = sfeFaceDetectInputSize(
        image,
        SFEFaceDetectionAccuracyType::SFE_FACE_DETECT_ACCURACY_TYPE_ACCURATE,
        face_size_min, face_size_max, &input_size);
    utils::checkError(error);
    recommended_width = input_size.width;
    recommended_height = input_size.height;
  };
}
 
void printHelp() {
  std::cout << "Help: Usage of the program." << std::endl;
  std::cout << "Options:" << std::endl;
  std::cout << "-h: Display help." << std::endl;
  std::cout << "-p: Probe image file." << std::endl;
  std::cout << "-g: Gallery folder." << std::endl;
  std::cout << "-d: Path to detector solver." << std::endl;
  std::cout << "-l: Path to landmarks solver." << std::endl;
  std::cout << "-e: Path to extraction solver." << std::endl;
  std::cout << "-t: Detection threshold. <0,1>" << std::endl;
  std::cout << "-i: Identification threshold. <0,1>" << std::endl;
  std::cout << "-m: Minimal face size in pixels to detect." << std::endl;
  std::cout << "-x: Max face size in pixels to detect." << std::endl;
}
 
void detectFace(SFEImage &image, SFESolver &detector_solver,
                SFEDetection &detected_face) {
 
  SFEImage resized_image{};
  DEFER(sfeImageFree(resized_image));
  SFEError error{};
  { // STAGE 1 Load image
    size_t recommended_width{};
    size_t recommended_height{};
 
    // Calculate optimal input image size for face detection. Image will be
    // resized to recommended size for optimal performance.
    getRecommendedImageSize(solver_face_detect, image, face_size_min,
                            face_size_max, recommended_width,
                            recommended_height);
 
    // Resize image to recommended size
    // NOTE: This function will resize the image without preserving the aspect
    // ratio of the image.
 
    error = sfeImageResize(image, recommended_width, recommended_height,
                           &resized_image);
    utils::checkError(error);
  }
 
  size_t detection_count = 1;
  { // STAGE 2: Detect face in the image
 
    // Detect face in the image
    error = sfeDetect(detector_solver, resized_image, detection_threshold,
                          &detected_face, &detection_count);
    utils::checkError(error);
 
    if (detection_count == 0) {
      std::ostringstream oss;
      oss << "Error: No face detected in the image ";
      throw std::runtime_error(oss.str());
 
    } else if (detection_count > 1) {
      std::ostringstream oss;
      oss << "Error: Found " << detection_count
          << " faces. Please provide an image with only one face.";
 
      throw std::runtime_error(oss.str());
    }
 
    std::cout << "Detected face in the image." << std::endl;
  }
}
 
int main(int argc, char *argv[]) {
 
  { // STAGE 0: Parse command line arguments
    // Map to store argument values
    std::unordered_map<std::string, std::string> args;
 
    // Parse command line arguments
    for (int i = 1; i < argc; ++i) {
      std::string arg = argv[i];
      if (arg[0] == '-') {
        if (arg == "-h") {
          printHelp();
          return 0;
        }
        // Check if there's a next argument and it isn't another option
        if (i + 1 < argc && argv[i + 1][0] != '-') {
          args[arg] = argv[++i];
        } else {
          std::cerr << "Option " << arg << " requires a value." << std::endl;
          return 1;
        }
      } else {
        std::cerr << "Unknown option: " << arg << std::endl;
        return 1;
      }
    }
 
    // Assign values to variables based on parsed arguments
    if (args.count("-p"))
      image_probe = args["-p"];
    if (args.count("-g"))
      image_gallery = args["-g"];
    if (args.count("-d"))
      solver_face_detect = args["-d"];
    if (args.count("-l"))
      solver_face_landmarks = args["-l"];
    if (args.count("-e"))
      solver_face_template = args["-e"];
    if (args.count("-t"))
      detection_threshold = std::stof(args["-t"]);
    if (args.count("-i"))
      identification_threshold = std::stof(args["-i"]);
    if (args.count("-m"))
      face_size_min = std::stoi(args["-m"]);
    if (args.count("-x"))
      face_size_max = std::stoi(args["-x"]);
 
    utils::printFormatted("EXAMPLE PARAMETERS");
    // Print resource names
    std::cout << "Probe image: " << image_probe << std::endl;
    std::cout << "Gallery image: " << image_gallery << std::endl;
 
    // Print solver names
    std::cout << "Face detection solver: " << solver_face_detect << std::endl;
    std::cout << "Face landmarks solver: " << solver_face_landmarks
              << std::endl;
    std::cout << "Face template solver: " << solver_face_template << std::endl;
 
    // Print other options
    std::cout << "Min face size [px]: " << face_size_min << std::endl;
    std::cout << "Max face size [px]: " << face_size_max << std::endl;
    std::cout << "Detection threshold: " << detection_threshold << std::endl;
    std::cout << "Identification threshold: " << identification_threshold
              << std::endl;
  }
 
  utils::printToolkitInfo();
 
  SFEError error{};
 
  SFESolver detector_solver{};
  DEFER(sfeSolverFree(detector_solver));
  SFESolver landmarks_solver{};
  DEFER(sfeSolverFree(landmarks_solver));
  SFESolver template_solver{};
  DEFER(sfeSolverFree(template_solver));
  { // STAGE 1.1: loading solvers
    utils::printFormatted("LOADING SOLVERS");
    // Initialize face detection solver
    error = sfeSolverCreate(
        solver_face_detect.c_str(), SOLVER_PARAMETERS.data(),
        SOLVER_PARAMETERS.size(), &detector_solver);
    utils::checkError(error);
 
    // Initialize landmarks detection solver
    error = sfeSolverCreate(
        solver_face_landmarks.c_str(), SOLVER_PARAMETERS.data(),
        SOLVER_PARAMETERS.size(), &landmarks_solver);
    utils::checkError(error);
 
    // Initialize template extraction solver
    error = sfeSolverCreate(
        solver_face_template.c_str(), SOLVER_PARAMETERS.data(),
        SOLVER_PARAMETERS.size(), &template_solver);
    utils::checkError(error);
  }
 
  SFEFaceTemplate probe_face_template;
  { // STAGE 1: Extract probe face template
    utils::printFormatted("PROBE TEMPLATE EXTRACTION");
 
    SFEImage image{};
    DEFER(sfeImageFree(image));
    { // STAGE 1.1: Load image
      // Load image data from file
      auto image_data = utils::readFile(image_probe);
 
      // Decode image from data
      error = sfeImageDecode(image_data.data(), image_data.size(), &image);
      utils::checkError(error);
    }
 
    SFEDetection detected_face = {};
    { // STAGE 1.2: Detect face in the image
      // Detect a face in the image
      detectFace(image, detector_solver, detected_face);
    }
 
    std::vector<SFEFaceLandmarks> landmarks(SFE_FACE_LANDMARK_COUNT);
    { // STAGE 1.3: Get face landmarks
      // Use landmarks detection solver to get relevant landmarks for
      // the detected face
      error = sfeFaceLandmarks(landmarks_solver, image, &detected_face,
                               landmarks.data());
      utils::checkError(error);
    }
 
    { // STAGE 1.4: Extract probe face template
      // Use template extraction solver to create new face
      // template
      error = sfeFaceTemplateExtract(template_solver, image, &detected_face,
                                     landmarks.data(), &probe_face_template);
      utils::checkError(error);
    }
 
    { // STAGE 1.5: (optional) Print template attributes
      utils::printFormatted("PROBE TEMPLATE ATTRIBUTES");
 
      SFEFaceTemplateVersion version = {};
      error = sfeFaceTemplateVersion(&probe_face_template, &version);
      utils::checkError(error);
 
      auto version_minor = (int)version.version_minor - (int)'0';
      std::cout << "Version of the probe template: " << version.version_major
                << '.' << version_minor << std::endl;
    }
 
    { // STAGE 1.6: (optional) Export and re-import template
      std::vector<uint8_t> buf(1024);
      size_t size = buf.size();
      error = sfeFaceTemplateExport(&probe_face_template, buf.data(), &size);
      if (error && size > buf.size()) {
        buf.resize(size);
        size = buf.size();
        error = sfeFaceTemplateExport(&probe_face_template, buf.data(), &size);
      }
      utils::checkError(error);
 
      SFEFaceTemplate imported_template = {};
      error = sfeFaceTemplateImport(buf.data(), size, &imported_template);
      utils::checkError(error);
 
      float roundtrip_score = 0.f;
      error = sfeFaceTemplateMatch(&probe_face_template, &imported_template,
                                  &roundtrip_score);
      utils::checkError(error);
      std::cout << "Export/import round-trip match score: " << roundtrip_score
                << std::endl;
    }
  }
 
  std::vector<std::string> image_paths{};
  { // STAGE 2: Get gallery image paths
    // Get folder names from the face image gallery
    auto folder_names = utils::getFiles(image_gallery);
 
    for (auto folder_name : folder_names) {
      auto gallery_folder = image_gallery + folder_name + "/";
 
      // Get image names from folder
      auto image_names = utils::getFiles(gallery_folder);
 
      // Extract template for each image in the folder
      for (auto image_name : image_names) {
 
        auto image_path = gallery_folder + image_name;
 
        image_paths.push_back(image_path);
      }
    }
  }
 
  std::vector<SFEFaceTemplate> gallery_face_templates(image_paths.size());
  std::vector<SFEDetection> detected_faces(image_paths.size());
  std::vector<std::array<SFEFaceLandmarks, SFE_FACE_LANDMARK_COUNT>> landmarks(
      image_paths.size());
  { // STAGE 2: Load gallery image and extract face templates for each image in
    // the gallery
    utils::printFormatted("GALLERY TEMPLATE EXTRACTION");
 
    for (size_t i = 0; i < image_paths.size(); i++) {
      std::cout << "Processing image #" << i << ", path: " << image_paths[i]
                << std::endl;
      SFEImage image{};
      DEFER(sfeImageFree(image));
      { // STAGE 2.1: Load image
        // Load image data from file
        auto image_data = utils::readFile(image_paths[i]);
 
        // Decode image from data
        error = sfeImageDecode(image_data.data(), image_data.size(), &image);
        utils::checkError(error);
      }
 
      { // STAGE 2.2: Detect face in the image
        // Detect a face in the image
        detectFace(image, detector_solver, detected_faces[i]);
      }
 
      { // STAGE 2.3: Get face landmarks
        // Use landmarks detection solver to get relevant landmarks for
        // the detected face
        error = sfeFaceLandmarks(landmarks_solver, image, &detected_faces[i],
                                 landmarks[i].data());
        utils::checkError(error);
      }
 
      { // STAGE 2.3: Extract face template
        // Use template extraction solver to get face_template solver
        error = sfeFaceTemplateExtract(
            template_solver, image, &detected_faces[i], landmarks[i].data(),
            &gallery_face_templates[i]);
        utils::checkError(error);
        std::cout << "Extracted face template." << std::endl;
      }
      std::cout << std::endl;
    }
  }
 
  size_t candidate_count = 1;
  int best_candidate_index = -1;
  std::vector<SFETemplateIdentificationCandidate> identification_results(
      image_paths.size());
  { // STAGE 3: Identification
    utils::printFormatted("1:N IDENTIFICATION");
 
    error = sfeFaceTemplateIdentify(
        &probe_face_template, gallery_face_templates.data(),
        gallery_face_templates.size(), identification_threshold,
        identification_results.data(), &candidate_count, 4);
    utils::checkError(error);
 
    // Resize the results vector to the actual number of candidates found, in
    // the case there is less than candidate_count
    identification_results.resize(candidate_count);
 
    std::cout << "Found " << identification_results.size()
              << " candidates above identification threshold "
              << identification_threshold << std::endl;
    for (auto &result : identification_results)
      std::cout << "Template index: #" << result.index
                << ", score: " << result.score << std::endl;
 
    // Since it's sorted vector by score, the best candidate is the first one
    best_candidate_index = identification_results[0].index;
  }
 
  { // STAGE 4: (optional) 1:1 matching with top candidate to showcase 1:1
    // matching
    utils::printFormatted("1:1 MATCHING WITH TOP CANDIDATE");
 
    if (identification_results.size() == 0) {
      std::cout << "No candidates found above identification threshold "
                << identification_threshold << std::endl;
      return 0;
    }
 
    auto best_candidate_template = gallery_face_templates[best_candidate_index];
 
    float match_confidence;
    error = sfeFaceTemplateMatch(&probe_face_template, &best_candidate_template,
                                 &match_confidence);
    utils::checkError(error);
 
    std::cout << "Matching score of probe template with template index #"
              << best_candidate_index << ", score: " << match_confidence
              << std::endl;
  }
 
  { // STAGE: 5 Optional, get face crop of best face and its bounding box and
    // save it to file
    utils::printFormatted("SAVE FACE CROP AND ANNOTATED IMAGE");
 
    auto best_face_index = identification_results[0].index;
 
    SFEImage image{};
    DEFER(sfeImageFree(image));
    { // STAGE 5.1: Load image
      // Load image data from file
      auto image_data = utils::readFile(image_paths[best_face_index]);
 
      // Decode image from data
      error = sfeImageDecode(image_data.data(), image_data.size(), &image);
      utils::checkError(error);
    }
 
    SFEFaceCrop crop_data = {};
    DEFER(sfeImageFree(crop_data.crop_image));
    { // STAGE 5.2: Get face crop
      //  Defines the size of the crop as an extension of
      // the detection bounding box.
      float face_size_extension = 2.0f;
      SFEError error =
          sfeFaceCrop(image, landmarks[best_face_index].data(),
                      face_size_extension, (float)(face_size_max), &crop_data);
      utils::checkError(error);
 
      std::cout << "Face crop extension: " << crop_data.face_size_extension
                << ", width of cropped image: " << crop_data.crop_image.width
                << "px, height of cropped image: "
                << crop_data.crop_image.height << "px." << std::endl;
    }
 
    { // STAGE 5.3: Save face crop to file
      auto png_file = std::vector<unsigned char>();
      size_t size =
          crop_data.crop_image.width * crop_data.crop_image.height * 3;
      png_file.resize(size);
      SFEError error2 = sfeImageEncode(crop_data.crop_image, SFE_IMAGE_FORMAT_PNG,
                                     png_file.data(), &size);
      utils::checkError(error2);
 
      std::cout << "Face crop saved as crop_identified_person.png" << std::endl;
      utils::saveFile("crop_identified_person.png", png_file);
    }
 
    { // STAGE 5.4: Annotate the image
      // Render bounding box with detection confidence and identification score
      std::stringstream label;
      label << " D:" << std::fixed << std::setprecision(2)
            << detected_faces[best_face_index].confidence
            << " M:" << identification_results[0].score;
 
      auto color = annotate::RED;
 
      annotate::labelBox(image, label.str(),
                         detected_faces[best_candidate_index].bounding_box,
                         annotate::WHITE, color);
 
      // Render landmarks
      for (auto &landmark : landmarks[best_face_index]) {
        auto x = landmark.x * image.width;
        auto y = landmark.y * image.height;
        annotate::circle(image, x, y, 3, annotate::YELLOW);
      }
 
      // Save annotated image
      size_t size = image.width * image.height * 3;
      auto png_file = std::vector<unsigned char>(size);
      error = sfeImageEncode(image, SFE_IMAGE_FORMAT_PNG, png_file.data(), &size);
      utils::checkError(error);
      png_file.resize(size);
      utils::saveFile("face_identify.png", png_file);
 
      std::cout << std::endl;
      std::cout << "Annotated image saved as face_identify.png" << std::endl;
    }
  }
 
  utils::printFormatted("FINISHED");
}