BPI-CM9-Linux/aml-comp/multimedia/gst-aml-utils/detect-sample/main.cpp

#include <iostream>
#include <fstream>
#include <string.h>

#include "nn_detect.h"
#include "nn_detect_utils.h"

using namespace std;

#define _CHECK_STATUS_(status, stat, lbl) do {\
    if (status != stat) \
    { \
        cout << "_CHECK_STATUS_ File" << __FUNCTION__ << __LINE__ <<endl; \
    }\
    goto lbl; \
}while(0)

static void *read_file(const char *file_path, int *file_size)
{
    FILE *fp = NULL;
    int size = 0;
    void *buf = NULL;
    fp = fopen(file_path, "rb");
    if (NULL == fp)
    {
        printf("open file fail!\n");
        return 0;
    }

    fseek(fp, 0, SEEK_END);
    size = ftell(fp);
    rewind(fp);

    buf = malloc(sizeof(unsigned char) * size);

    size_t bytes_read = fread(buf, 1, size, fp);
    if (bytes_read != size) {
        fclose(fp);
        free(buf);
        return NULL;
    }

    fclose(fp);

    *file_size = size;
    return buf;
}

string create_result_filename(const string& picture_path) {
    size_t last_slash = picture_path.find_last_of("/\\");
    string dir_path = (last_slash == string::npos) ? "" : picture_path.substr(0, last_slash + 1);
    string base_name = (last_slash == string::npos) ? picture_path : picture_path.substr(last_slash + 1);

    size_t last_dot = base_name.find_last_of('.');
    string name_part = (last_dot == string::npos) ? base_name : base_name.substr(0, last_dot);

    string result_filename = dir_path + name_part + "_result.txt";
    return result_filename;
}

static void write_labeled_float(FILE *file, const char *label, const void* value, const type_info& type) {

    fprintf(file, "%s ", label);
    if (type == typeid(float)) {
        fprintf(file, "%.6f\n", *reinterpret_cast<const float*>(value));
    } else if (type == typeid(int)) {
        fprintf(file, "%d\n", *reinterpret_cast<const int*>(value));
    } else {
        fprintf(file, "Unknown type\n");
    }

}


int run_detect_model(int argc, char** argv)
{
    int ret = 0;
    int nn_height, nn_width, nn_channel;
    det_model_type type = DET_YOLOFACE_V2;
    pDetResult resultData;

    if (argc !=3) {
        cout << "input param error" <<endl;
        cout << "Usage: " << argv[0] << " type  picture_path"<<endl;
        return -1;
    }
    type = (det_model_type)atoi(argv[1]);

    char* picture_path = argv[2];
    cout << "det_set_log_config Debug" <<endl;

    //get the filename
    string picture_path_str(picture_path);
    string result_filename = create_result_filename(picture_path_str);
    cout << "Result filename: " << result_filename << endl;

    //prepare model
    ret = det_set_model(type);
    if (ret) {
        cout << "det_set_model fail. ret=" << ret <<endl;
        return ret;
    }
    cout << "det_set_model success!!" << endl;

    ret = det_get_model_size(type, &nn_width, &nn_height, &nn_channel);
    if (ret) {
        cout << "det_get_model_size fail" <<endl;
        return ret;
    }

    cout << "\nmodel.width:" << nn_width <<endl;
    cout << "model.height:" << nn_height <<endl;
    cout << "model.channel:" << nn_channel << "\n" <<endl;

    int input_size;
    input_image_t image;
    image.data      = (unsigned char *)read_file(picture_path, &input_size);
    image.width     = nn_width;
    image.height    = nn_height;
    image.channel   = nn_channel;
    image.pixel_format = PIX_FMT_RGB888;

    cout << "Det_set_input START" << endl;
    ret = det_set_input(image, type);
    if (ret) {
        cout << "det_set_input fail. ret=" << ret << endl;
        det_release_model(type);
        return ret;
    }
    cout << "Det_set_input END" << endl;

    cout << "Det_get_result START" << endl;
    resultData = (pDetResult)malloc(sizeof(*resultData));
    ret = det_get_result(resultData, type);
    if (ret) {
        cout << "det_get_result fail. ret=" << ret << endl;
        det_release_model(type);
        return ret;
    }

    //wirite result to file
    FILE *file = fopen(result_filename.c_str(), "wb");
    if (file == NULL) {
        perror("Failed to open file");
        return ret;
    }

    for ( int i = 0; i < resultData->result.det_result.detect_num; i++) {

        cout << "num" << i+1 << endl;
        cout << "left" << resultData->result.det_result.point[i].point.rectPoint.left << endl;
        cout << "top" << resultData->result.det_result.point[i].point.rectPoint.top << endl;
        cout << "right" << resultData->result.det_result.point[i].point.rectPoint.right << endl;
        cout << "bottom" << resultData->result.det_result.point[i].point.rectPoint.bottom << endl;
        cout << "score" << resultData->result.det_result.point[i].point.rectPoint.score << endl;

        //write_labeled_float(file, "num:", i);
        write_labeled_float(file, "num", &(resultData->result.det_result.detect_num),typeid(resultData->result.det_result.detect_num));
        write_labeled_float(file, "left", &(resultData->result.det_result.point[i].point.rectPoint.left),typeid(resultData->result.det_result.point[i].point.rectPoint.left));
        write_labeled_float(file, "top", &(resultData->result.det_result.point[i].point.rectPoint.top),typeid(resultData->result.det_result.point[i].point.rectPoint.top));
        write_labeled_float(file, "right", &(resultData->result.det_result.point[i].point.rectPoint.right),typeid(resultData->result.det_result.point[i].point.rectPoint.right));
        write_labeled_float(file, "bottom", &(resultData->result.det_result.point[i].point.rectPoint.bottom),typeid(resultData->result.det_result.point[i].point.rectPoint.bottom));
        write_labeled_float(file, "score", &(resultData->result.det_result.point[i].point.rectPoint.score),typeid(resultData->result.det_result.point[i].point.rectPoint.score));
    }
    fclose(file);

    cout << "Det_get_result END" << endl;

    det_release_model(type);
    return ret;
}

int main(int argc, char** argv)
{
    det_model_type type;
    if (argc < 3) {
        cout << "input param error" <<endl;
        cout << "Usage: " << argv[0] << " type  picture_path"<<endl;
        return -1;
    }

    type = (det_model_type)atoi(argv[1]);
    switch (type) {
        case DET_YOLO_V3:
            run_detect_model(argc, argv);
            break;
        case DET_AML_FACE_DETECTION:
            run_detect_model(argc, argv);
            break;
        default:
            cerr << "not support type=" << type <<endl;
            break;
    }

    return 0;
}