example_person_attributes.cpp

Go to the documentation of this file.

#include <sstream>
#include <iomanip>
#include <vector>
#include <algorithm>
 
#include <sfe_toolkit/sfe_core.h>
#include <sfe_toolkit/sfe_object.h>
#include <sfe_toolkit/sfe_person.h>
 
#include "utils.hpp"
#include "solvers.h"
#include "annotate.hpp"
#include "strings.hpp"
 
// Solver parameters, we are using ONNX Runtime CPU backend
const auto SOLVER_PARAMETERS = std::vector<SFESolverParameter>{
    SFESolverParameter { "intra_threads", "4" },
};
 
std::string image_path = "assets/person/person_attributes.jpg";
std::string solver_object_detect = SOLVER_OBJECT_DETECT;
std::string solver_person_attributes = SOLVER_PERSON_ATTRIBUTES;
std::string solver_person_pose = SOLVER_PERSON_POSE;
 
int main() {
    SFEError error;
 
    SFEImage image{};
    DEFER(sfeImageFree(image));
    { // STAGE 1: Load Image
        auto encoded_image = utils::readFile(image_path);
        error = sfeImageDecode(encoded_image.data(), encoded_image.size(), &image);
        utils::checkError(error);
 
        std::cout << "Image size: " << image.width << "x" << image.height << std::endl;
    }
 
    std::vector<SFEDetection> detections(10);
    { // STAGE 2: Run object detection
        // Initialize detection solver
        SFESolver detection_solver{};
        DEFER(sfeSolverFree(detection_solver));
        error = sfeSolverCreate(solver_object_detect.c_str(), SOLVER_PARAMETERS.data(), SOLVER_PARAMETERS.size(), &detection_solver);
        utils::checkError(error);
 
        // Run detection for up to 10 objects
        size_t detection_count = detections.size();
        float detection_threshold = 0.3;
        error = sfeDetect(detection_solver,
            image,
            detection_threshold,
            detections.data(), &detection_count);
        utils::checkError(error);
        detections.resize(detection_count);
 
        std::cout << "Detected " << detection_count << " objects" << std::endl;
    }
 
    std::vector<size_t> detection_classes;
    { // STAGE 3: Check for best detection classes
        for (size_t i = 0; i < detections.size(); i++) {
            auto& detection = detections[i];
            size_t best_class_index = 0;
            for (size_t c = 0; c < SFE_OBJECT_DETECT; c++) {
                auto object = detection.modality_data.object.objects[c];
                if (object.confidence < 0.01) continue;
                if (object.confidence > detection.modality_data.object.objects[best_class_index].confidence) {
                    best_class_index = c;
                }
                
                std::cout << "Detection " << i << " : " << OBJECT_CLASS_NAMES[c] << " confidence: " << std::fixed << std::setprecision(2) << detection.modality_data.object.objects[c].confidence << std::endl;
            }
            detection_classes.push_back(best_class_index);
        }
    }
 
    size_t person_index;
    { // STAGE 4: Find first person in detections
        auto person_iter = std::find(detection_classes.begin(), detection_classes.end(), SFE_OBJECT_TYPE_PERSON);
        if (person_iter == detection_classes.end()) {
            std::cout << "No person detected" << std::endl;
            return 0;
        }
        person_index = std::distance(detection_classes.begin(), person_iter);
 
        std::cout << "Person detected at index " << person_index << std::endl;
    }
 
    std::array<SFEPersonAttribute, SFE_PERSON_ATTRIBUTE_COUNT> person_attributes;
    { // STAGE 5: Extract person attributes
        SFESolver attribute_solver{};
        DEFER(sfeSolverFree(attribute_solver));
        error = sfeSolverCreate(solver_person_attributes.c_str(), SOLVER_PARAMETERS.data(), SOLVER_PARAMETERS.size(), &attribute_solver);
        utils::checkError(error);
 
        // Extract attributes for the detected person
        error = sfePersonAttributes(attribute_solver, image, &detections[person_index], person_attributes.data());
        utils::checkError(error);
 
        for (auto& attribute : person_attributes) {
            if (attribute.confidence < 0.1) continue;
            std::cout << "Attribute " << PERSON_ATTRIBUTE_TYPE[attribute.attribute_type] << " confidence: " << std::fixed << std::setprecision(2) << attribute.confidence << std::endl;
        }
    }
 
    std::array<SFEPersonPoseKeypoint, SFE_PERSON_POSE_KEYPOINT_COUNT> person_pose;
    { // STAGE 6: Person pose estimation
        SFESolver pose_solver;
        DEFER(sfeSolverFree(pose_solver));
        error = sfeSolverCreate(solver_person_pose.c_str(), SOLVER_PARAMETERS.data(), SOLVER_PARAMETERS.size(), &pose_solver);
        utils::checkError(error);
 
        // Estimate pose for the detected person
        error = sfePersonPose(pose_solver, image, &detections[person_index], person_pose.data());
        utils::checkError(error);
 
        for (auto& keypoint : person_pose) {
            if (keypoint.confidence < 0.1) continue;
            std::cout << "Keypoint " << keypoint.keypoint_type << " confidence: " << std::fixed << std::setprecision(2) << keypoint.confidence << std::endl;
        }
    }
 
    { // END STAGE: Display results
        for (size_t i = 0; i < detections.size(); i++) {
            auto& detection = detections[i];
            auto& class_index = detection_classes[i];
            auto color = i == person_index ? annotate::GREEN : annotate::RED;
            // Render bounding box
            std::stringstream label;
            label << i << ": " << OBJECT_CLASS_NAMES[class_index] << std::fixed << std::setprecision(2)
                << " c=" << detection.modality_data.object.objects[class_index].confidence
                << " d=" << detection.confidence;
 
            annotate::labelBox(image, label.str(), detection.bounding_box, annotate::WHITE, color);
        }
 
        // Render pose points
        for (auto& keypoint : person_pose) {
            annotate::circle(image, keypoint.x * image.width, keypoint.y * image.height, 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("person_attributes.png", png_file);
    }
}