mmdeploy的object_detection优化

mmdeploy编译后，bin目录下有c++编译的object_detection程序，用于目标检测和实例分割，该源码位于：

mmdeploy/demo/csrc/c/object_detection.cpp

#include <fstream>
#include <opencv2/imgcodecs/imgcodecs.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <string>
#include <iostream>
#include <ctime>
#include "mmdeploy/detector.h"

std::string TrimEnd(const std::string& str) {
    size_t endpos = str.find_last_not_of(" \t\r\n");
    return (endpos != std::string::npos) ? str.substr(0, endpos + 1) : str;
}

// 在加载类别时使用TrimEnd()函数
std::vector<std::string> LoadCategories(const std::string& coco_class_file) {
    std::ifstream file(coco_class_file);
    std::vector<std::string> categories;
    if (file.is_open()) {
        std::string line;
        while (getline(file, line)) {
            categories.push_back(TrimEnd(line)); // 显式去除末尾空白和换行符
        }
        file.close();
    } else {
        std::cerr << "Error opening class file: " << coco_class_file << std::endl;
        exit(1);
    }
    return categories;
}

void SaveResultsToJson(const std::vector<std::vector<float>>& bboxes, 
                       const std::vector<int>& labels, 
                       const std::vector<float>& scores, 
                       const std::string& json_save_path) {
    std::ofstream json_file(json_save_path);
    if (json_file.is_open()) {
        json_file << "{\n";
        json_file << "  \"bboxes\": [\n";
        for (size_t i = 0; i < bboxes.size(); ++i) {
            json_file << "    [";
            for (size_t j = 0; j < bboxes[i].size(); ++j) {
                json_file << bboxes[i][j];
                if (j < bboxes[i].size() - 1) json_file << ", ";
            }
            json_file << "]";
            if (i < bboxes.size() - 1) json_file << ",";
            json_file << "\n";
        }
        json_file << "  ],\n";
        json_file << "  \"labels\": [\n    ";
        for (size_t i = 0; i < labels.size(); ++i) {
            json_file << labels[i];
            if (i < labels.size() - 1) json_file << ", ";
        }
        json_file << "\n  ],\n";
        json_file << "  \"scores\": [\n    ";
        for (size_t i = 0; i < scores.size(); ++i) {
            json_file << scores[i];
            if (i < scores.size() - 1) json_file << ", ";
        }
        json_file << "\n  ]\n";
        json_file << "}\n";
        json_file.close();
    } else {
        std::cerr << "Error opening JSON file for writing: " << json_save_path << std::endl;
    }
}
int main(int argc, char* argv[]) {
    if (argc != 8) {
        fprintf(stderr, "usage:\n  object_detection device_name model_path image_path save_path coco_txt_path score\n");
        return 1;
    }

    auto device_name = argv[1];
    auto model_path = argv[2];
    auto image_path = argv[3];
    auto save_path = argv[4];
    auto json_save_path = argv[5];
    auto coco_txt_path = argv[6];
    float score_threshold = std::stof(argv[7]);

    // 读取分类名称
    auto categories = LoadCategories(coco_txt_path);
    cv::Mat img = cv::imread(image_path);
    if (!img.data) {
        fprintf(stderr, "failed to load image: %s\n", image_path);
        return 1;
    }
    cv::Size img_size = img.size();
    mmdeploy_detector_t detector{};
    int status{};
    status = mmdeploy_detector_create_by_path(model_path, device_name, 0, &detector);
    if (status != MMDEPLOY_SUCCESS) {
        fprintf(stderr, "failed to create detector, code: %d\n", (int)status);
        return 1;
    }

    mmdeploy_mat_t mat{
        img.data, img.rows, img.cols, 3, MMDEPLOY_PIXEL_FORMAT_BGR, MMDEPLOY_DATA_TYPE_UINT8
    };

    mmdeploy_detection_t* bboxes{};
    int* res_count{};

    status = mmdeploy_detector_apply(detector, &mat, 1, &bboxes, &res_count);
    if (status != MMDEPLOY_SUCCESS) {
        fprintf(stderr, "failed to apply detector, code: %d\n", (int)status);
        return 1;
    }

    std::vector<std::vector<float>> json_bboxes;
    std::vector<int> json_labels;
    std::vector<float> json_scores;

    for (int i = 0; i < *res_count; ++i) {
        const auto& box = bboxes[i].bbox;
        const auto& mask = bboxes[i].mask;

        // 跳过无效的检测框（框的高度或宽度 < 1）
        if ((box.right - box.left) < 1 || (box.bottom - box.top) < 1) {
            continue;
        }

        // 跳过得分低于指定阈值的检测
        if (bboxes[i].score < score_threshold) {
            continue;
        }

        std::vector<float> bbox = {box.left, box.top, box.right, box.bottom};
        json_bboxes.push_back(bbox);
        json_labels.push_back(bboxes[i].label_id);
        json_scores.push_back(bboxes[i].score);
        
        // 生成遮罩叠加，如果模型导出遮罩
        if (mask != nullptr) {
          //fprintf(stdout, "mask %d, height=%d, width=%d\n", i, mask->height, mask->width);
          // 分离RGB通道，叠加遮罩到指定颜色通道
          cv::Mat ch[3], mask_img;
          int col = 0;  // int col = i % 3;
          split(img, ch);
          cv::Mat imgMask(mask->height, mask->width, CV_8UC1, &mask->data[0]);
          // rtmdet-inst
          if (img_size.height == mask->height && img_size.width == mask->width) {
            mask_img = ch[col];
          }
          else {
            auto x0 = std::max(std::floor(box.left) - 1, 0.f);
            auto y0 = std::max(std::floor(box.top) - 1, 0.f);
            cv::Rect roi((int)x0, (int)y0, mask->width, mask->height);
            mask_img = ch[col](roi);
          }
          cv::bitwise_or(imgMask, mask_img, mask_img);
          merge(ch, 3, img);
        }

        // 在图像上绘制矩形框
        cv::rectangle(img, cv::Point{(int)box.left, (int)box.top},
                      cv::Point{(int)box.right, (int)box.bottom}, cv::Scalar{0, 255, 0});

        // 在矩形框的右上角绘制分类名称
        std::string class_name = categories[bboxes[i].label_id];
        int baseLine;
        cv::Size labelSize = cv::getTextSize(class_name, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
        int top = std::max((int)box.top, labelSize.height);
        cv::putText(img, class_name, cv::Point((int)box.left, top), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255), 1);
    }

    cv::imwrite(save_path, img);

    // 将检测结果保存为JSON文件
    SaveResultsToJson(json_bboxes, json_labels, json_scores, json_save_path);

    mmdeploy_detector_release_result(bboxes, res_count, 1);
    mmdeploy_detector_destroy(detector);

    return 0;
}

调用方法：

./bin/object_detection cuda yolox_x ../mmdetection3/demo/demo.jpg demo.jpg 1.json coco.txt 0.5

其中coco.txt为coco数据集标签文件，内容如下：

person
bicycle
car
motorbike
aeroplane
bus
train
truck
boat
traffic light
fire hydrant
stop sign
parking meter
bench
bird
cat
dog
horse
sheep
cow
elephant
bear
zebra
giraffe
backpack
umbrella
handbag
tie
suitcase
frisbee
skis
snowboard
sports ball
kite
baseball bat
baseball glove
skateboard
surfboard
tennis racket
bottle
wine glass
cup
fork
knife
spoon
bowl
banana
apple
sandwich
orange
broccoli
carrot
hot dog
pizza
donut
cake
chair
sofa
pottedplant
bed
diningtable
toilet
tvmonitor
laptop
mouse
remote
keyboard
cell phone
microwave
oven
toaster
sink
refrigerator
book
clock
vase
scissors
teddy bear
hair drier
toothbrush
motorcycle
potted plant
dining table
tv
couch
airplane