example_face_track.cpp

Go to the documentation of this file.

 
#include "sfe_toolkit/sfe_core.h"
#include "sfe_toolkit/sfe_face.h"
#include <ostream>
#include <regex>
#include <vector>
 
#include <iomanip>
#include <sstream>
 
#include "annotate.hpp"
#include "solvers.h"
#include "utils.hpp"
#include <unordered_map>
 
std::string image_probe = "assets/face/images/obiwan0.png";
 
std::string solver_face_detect = SOLVER_FACE_DETECT;
 
const auto SOLVER_PARAMETERS = std::vector<SFESolverParameter>{};
 
size_t face_size_min = 28;
size_t face_size_max = 170;
 
float detection_threshold = 0.1f;
 
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 << "-d: Path to detector solver." << std::endl;
  std::cout << "-t: Detection 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;
}
 
std::vector<SFEDetection> detectFaces(SFEImage &image,
                                       SFESolver &detector_solver) {
 
  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 = 10;
  std::vector<SFEDetection> detected_faces(detection_count);
  { // STAGE 2: Detect face in the image
 
    // Detect faces in the image
    error = sfeDetect(detector_solver, resized_image, detection_threshold,
                          detected_faces.data(), &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());
    }
    detected_faces.resize(detection_count);
 
    std::cout << "Detected " << detection_count << " faces in the image."
              << std::endl;
  }
  return detected_faces;
}
 
std::ostream &operator<<(std::ostream &os, const SFEEntity &entity) {
  for (size_t i = 0; i < 16; ++i) {
    os << std::hex << std::setw(2) << std::setfill('0')
       << static_cast<int>(entity.uuid[i]);
    if (i == 3 || i == 5 || i == 7 || i == 9) {
      os << '-';
    }
  }
  os << std::dec;
  return os;
}
 
std::ostream &operator<<(std::ostream &os, const SFETrackedState &state) {
  switch (state) {
  case SFE_TRACKED_STATE_NEW:
    os << "NEW";
    break;
  case SFE_TRACKED_STATE_TRACKED:
    os << "TRACKED";
    break;
  case SFE_TRACKED_STATE_LOST:
    os << "LOST";
    break;
  case SFE_TRACKED_STATE_REMOVED:
    os << "REMOVED";
    break;
  default:
    os << "UNKNOWN";
    break;
  }
  return os;
}
 
std::ostream &operator<<(std::ostream &os, const SFETracked &tracked) {
  os << "Face ID: " << tracked.id << " UUID: " << tracked.uuid
     << " State: " << tracked.state;
  return os;
}
 
template <typename T>
std::ostream &operator<<(std::ostream &os, const std::vector<T> &vec) {
  for (auto &item : vec) {
    os << item << std::endl;
  }
  return os;
}
 
void frame(std::vector<SFEDetection> &detections, SFETracker tracker) {
  size_t reserved_size = 10;
  size_t tracked_faces_count = reserved_size;
  size_t lost_faces_count = reserved_size;
  size_t removed_faces_count = reserved_size;
 
  std::vector<SFETracked> tracked_faces(tracked_faces_count);
  std::vector<SFEEntity> lost_faces(lost_faces_count);
  std::vector<SFEEntity> removed_faces(removed_faces_count);
 
  // Update tracker with detected faces
  SFEError error =
      sfeDetectionTrackerUpdate(tracker, detections.data(), detections.size(),
                           tracked_faces.data(), &tracked_faces_count);
  utils::checkError(error);
  tracked_faces.resize(tracked_faces_count);
 
  // Get lost faces
  error = sfeDetectionTrackerLost(tracker, lost_faces.data(), &lost_faces_count);
  utils::checkError(error);
  lost_faces.resize(lost_faces_count);
 
  // Get removed faces
  error = sfeDetectionTrackerRemoved(tracker, removed_faces.data(),
                                &removed_faces_count);
  utils::checkError(error);
  removed_faces.resize(removed_faces_count);
 
  // Print out frame results
  std::cout << "Tracked" << std::endl;
  std::cout << tracked_faces;
  std::cout << "Lost" << std::endl;
  std::cout << lost_faces;
  std::cout << "Removed" << std::endl;
  std::cout << removed_faces;
}
 
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("-d"))
      solver_face_detect = args["-d"];
    if (args.count("-t"))
      detection_threshold = std::stof(args["-t"]);
    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;
 
    // Print solver names
    std::cout << "Face detection solver: " << solver_face_detect << 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;
  }
 
  utils::printToolkitInfo();
 
  SFEError error{};
 
  SFESolver detector_solver{};
  DEFER(sfeSolverFree(detector_solver));
  { // STAGE 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);
  }
 
  std::vector<SFEDetection> detected_faces;
  { // STAGE 2: Detect faces
    utils::printFormatted("DETECT FACES");
 
    SFEImage image{};
    DEFER(sfeImageFree(image));
    { // STAGE 2.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 2.2: Detect face in the image
      // Detect a face in the image
      detected_faces = detectFaces(image, detector_solver);
    }
  }
 
  SFETracker tracker{};
  DEFER(sfeDetectionTrackerFree(tracker));
  { // STAGE 3: Track faces
    utils::printFormatted("TRACK FACES");
    error = sfeDetectionTrackerCreate(0.3f, 0.1f, 0.3f, 0.1f, 1, &tracker);
    utils::checkError(error);
 
    // We will be using the current detected faces and simulate detection across
    // 4 frames In a real application, you would call detectFaces() for each
    // frame and pass the detected faces to the tracker
 
    // First frame with detected faces, all should track as NEW
    utils::printFormatted("FRAME 1");
    frame(detected_faces, tracker);
 
    // Second frame with the same UUIDs, all should track as TRACKED
    utils::printFormatted("FRAME 2");
    frame(detected_faces, tracker);
 
    // Simulate a lost detection
    detected_faces.clear();
 
    // We should see the lost UUIDs
    utils::printFormatted("FRAME 3");
    frame(detected_faces, tracker);
 
    // We should see the removed UUIDs
    utils::printFormatted("FRAME 4");
    frame(detected_faces, tracker);
  }
  utils::printFormatted("FINISHED");
}