initial commit

wangrong

2025-01-03 2c4b8838f89bebca97615c5f0bac2c58e0349771

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
 
#include "vp_mask_rcnn_detector_node.h" 
 
 
namespace vp_nodes { 
    vp_mask_rcnn_detector_node::vp_mask_rcnn_detector_node(std::string node_name,  
                            std::string model_path,  
                            std::string model_config_path,  
                            std::string labels_path,  
                            int input_width,  
                            int input_height,  
                            int batch_size, 
                            int class_id_offset, 
                            float score_threshold, 
                            float scale, 
                            cv::Scalar mean, 
                            cv::Scalar std, 
                            bool swap_rb): 
                            vp_primary_infer_node(node_name, model_path, model_config_path, labels_path, input_width, input_height, batch_size, class_id_offset, scale, mean, std, swap_rb), 
                            score_threshold(score_threshold) { 
        this->initialized(); 
    } 
     
    vp_mask_rcnn_detector_node::~vp_mask_rcnn_detector_node() { 
        deinitialized(); 
    } 
 
 
    void vp_mask_rcnn_detector_node::infer(const cv::Mat& blob_to_infer, std::vector<cv::Mat>& raw_outputs) { 
        // blob_to_infer is a 4D matrix 
        // the first dim is number of batch, MUST be 1 
        assert(blob_to_infer.dims == 4); 
        assert(blob_to_infer.size[0] == 1); 
        assert(!net.empty()); 
 
        net.setInput(blob_to_infer); 
        net.forward(raw_outputs, out_names); 
    } 
 
 
    void vp_mask_rcnn_detector_node::preprocess(const std::vector<cv::Mat>& mats_to_infer, cv::Mat& blob_to_infer) { 
        // ignore preprocess logic in base class 
        cv::dnn::blobFromImages(mats_to_infer, blob_to_infer, 1.0, cv::Size(), cv::Scalar(), true, false); 
    } 
 
    void vp_mask_rcnn_detector_node::postprocess(const std::vector<cv::Mat>& raw_outputs, const std::vector<std::shared_ptr<vp_objects::vp_frame_meta>>& frame_meta_with_batch) { 
        auto& frame_meta = frame_meta_with_batch[0]; 
         
        auto outDetections = raw_outputs[0]; 
        auto outMasks = raw_outputs[1]; 
     
        // Output size of masks is NxCxHxW where 
        // N - number of detected boxes 
        // C - number of classes (excluding background) 
        // HxW - segmentation shape 
        const int numDetections = outDetections.size[2]; 
        const int numClasses = outMasks.size[1]; 
     
        auto outDetections_ = outDetections.reshape(1, outDetections.total() / 7); 
        auto& frame = frame_meta->frame; 
        for (int i = 0; i < numDetections; ++i) { 
            float score = outDetections_.at<float>(i, 2); 
            if (score > score_threshold) { 
                // Extract the bounding box 
                int classId = static_cast<int>(outDetections_.at<float>(i, 1)); 
                int left = static_cast<int>(frame.cols * outDetections_.at<float>(i, 3)); 
                int top = static_cast<int>(frame.rows * outDetections_.at<float>(i, 4)); 
                int right = static_cast<int>(frame.cols * outDetections_.at<float>(i, 5)); 
                int bottom = static_cast<int>(frame.rows * outDetections_.at<float>(i, 6)); 
     
                left = max(0, min(left, frame.cols - 1)); 
                top = max(0, min(top, frame.rows - 1)); 
                right = max(0, min(right, frame.cols - 1)); 
                bottom = max(0, min(bottom, frame.rows - 1)); 
 
                auto label = (labels.size() < classId + 1) ? "" : labels[classId]; 
 
                // Extract the mask for the object 
                cv::Mat objectMask(outMasks.size[2], outMasks.size[3], CV_32F, outMasks.ptr<float>(i, classId)); 
                 
                classId += class_id_offset; 
                auto target = std::make_shared<vp_objects::vp_frame_target>(left, top, right - left, bottom - top, classId, score, frame_meta->frame_index, frame_meta->channel_index, label); 
                target->mask = objectMask; 
 
                frame_meta->targets.push_back(target); 
            } 
        } 
    } 
}