first commit

nodes/infers/vp_yolo_detector_node.cpp

@@ -0,0 +1,118 @@
+
+#include "vp_yolo_detector_node.h"
+
+namespace vp_nodes {
+    
+    vp_yolo_detector_node::vp_yolo_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 confidence_threshold,
+                            float nms_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),
+                            confidence_threshold(confidence_threshold),
+                            nms_threshold(nms_threshold) {
+        this->initialized();
+    }
+    
+    vp_yolo_detector_node::~vp_yolo_detector_node() {
+        deinitialized();
+    }
+    
+    void vp_yolo_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) {
+        // make sure heads of output are not zero
+        assert(raw_outputs.size() > 0);
+
+        // check dims of each output
+        auto dim_offset = raw_outputs[0].dims == 2 ? 0 : 1;
+        auto batch = dim_offset == 0 ? 1 : raw_outputs[0].size[0];
+
+        std::vector<int> class_ids;
+        std::vector<float> confidences;
+        std::vector<cv::Rect> boxes;
+
+        for (int k = 0; k < batch; k++) {  // scan batch
+            auto& frame_meta = frame_meta_with_batch[k];
+            for (int i = 0; i < raw_outputs.size(); ++i) {  // scan heads of output
+                cv::Mat output = cv::Mat(raw_outputs[i].size[0 + dim_offset], raw_outputs[i].size[1 + dim_offset], CV_32F, const_cast<uchar*>(raw_outputs[i].ptr(k)));
+
+                auto data = (float*)output.data;
+                for (int j = 0; j < output.rows; ++j, data += output.cols) {
+                    
+                    float confidence = data[4];
+                    // check confidence threshold
+                    if (confidence < confidence_threshold) {
+                        continue;
+                    }
+
+                    cv::Mat scores = output.row(j).colRange(5, output.cols);
+                    cv::Point class_id;
+                    double max_score;
+                    // Get the value and location of the maximum score
+                    cv::minMaxLoc(scores, 0, &max_score, 0, &class_id);
+
+                    // check score threshold
+                    if (max_score >= score_threshold) {
+                        // smaller than 1.0 means relative value 
+                        // greater than 1.0 means absolute value which need to be converted to relative value
+                        int centerX = (int)((data[0] <= 1 ? data[0] : data[0] / input_width) * frame_meta->frame.cols);
+                        int centerY = (int)((data[1] <= 1 ? data[1] : data[1] / input_height) * frame_meta->frame.rows);
+                        int width = (int)((data[2] <= 1 ? data[2] : data[2] / input_width) * frame_meta->frame.cols);
+                        int height = (int)((data[3] <= 1 ? data[3] : data[3] / input_height) * frame_meta->frame.rows);
+                        int left = centerX - width / 2;
+                        int top = centerY - height / 2;
+
+                        class_ids.push_back(class_id.x);
+                        confidences.push_back((float)confidence);
+                        boxes.push_back(cv::Rect(left, top, width, height));
+                    }
+                }
+            }
+
+            // nms
+            std::vector<int> indices;
+            cv::dnn::NMSBoxes(boxes, confidences, confidence_threshold, nms_threshold, indices);
+
+            // create target
+            for (int i = 0; i < indices.size(); ++i) {
+                int idx = indices[i];
+                auto box = boxes[idx];
+                auto confidence = confidences[idx];
+                auto class_id = class_ids[idx];
+                auto label = (labels.size() < class_id + 1) ? "" : labels[class_id];
+
+                // check value range
+                box.x = std::max(box.x, 0);
+                box.y = std::max(box.y, 0);
+                box.width = std::min(box.width, frame_meta->frame.cols - box.x);
+                box.height = std::min(box.height, frame_meta->frame.rows - box.y);
+                if (box.width <= 0 || box.height <= 0) {
+                    continue;
+                }
+                
+                // apply offset to class id since multi detectors can exist at front of this one.
+                // later we MUST use the new class id (applied offset) instead of orignal one.
+                class_id += class_id_offset;
+
+                auto target = std::make_shared<vp_objects::vp_frame_target>(box.x, box.y, box.width, box.height, class_id, confidence, frame_meta->frame_index, frame_meta->channel_index, label);
+
+                // insert target back to frame meta
+                frame_meta->targets.push_back(target);
+            }
+
+            class_ids.clear();
+            confidences.clear();
+            boxes.clear();
+        }
+    }
+}