update mask postprocess

Li-Hongda · Li-Hongda · commit 9f0c165a0828 · 2023-05-19T22:47:39.000+08:00
diff --git a/configs/rtdetr.yaml b/configs/rtdetr.yaml
@@ -1,14 +1,12 @@
 rtdetr:
-    onnx_file:     "../weights/rtdetr/rtdetr_r50vd_6x_coco.onnx"
-    engine_file:   "../weights/rtdetr/rtdetr_r50vd.trt"
+    onnx_file:     "../weights/rtdetr/rtdetr_hgnetv2_l_6x_coco.onnx"
+    engine_file:   "../weights/rtdetr/rtdetr_hgnetl.trt"
     type:          "coco80"
     mode:          "fp32"
     batchSize:     1
-    inputChannel:  3
     imageWidth:    640
     imageHeight:   640
-    conf_thr:      0.25
-    nms_thr:       0.45
-    strides:       [8, 16, 32]
+    conf_thr:      0.5
+    num_queries:   300
     imgMean:       [ 0, 0, 0 ]
     imgStd:        [ 1, 1, 1 ]
diff --git a/configs/yolov5-seg.yaml b/configs/yolov5-seg.yaml
@@ -4,7 +4,6 @@ yolov5-seg:
     type:          "coco80"
     mode:          "fp32"
     batchSize:     1
-    inputChannel:  3
     imageWidth:    640
     imageHeight:   640
     conf_thr:      0.25
diff --git a/configs/yolov5.yaml b/configs/yolov5.yaml
@@ -4,7 +4,6 @@ yolov5:
     type:          "coco80"
     mode:          "fp32"
     batchSize:     1
-    inputChannel:  3
     imageWidth:    640
     imageHeight:   640
     conf_thr:      0.25
diff --git a/configs/yolov7-p6.yaml b/configs/yolov7-p6.yaml
@@ -4,7 +4,6 @@ yolov7:
     type:          "coco80"
     mode:          "fp32"    
     batchSize:     1
-    inputChannel:  3
     imageWidth:    1280
     imageHeight:   1280
     conf_thr:      0.4
diff --git a/configs/yolov7.yaml b/configs/yolov7.yaml
@@ -4,7 +4,6 @@ yolov7:
     type:          "coco80"
     mode:          "fp32"
     batchSize:     1
-    inputChannel:  3
     imageWidth:    640
     imageHeight:   640
     conf_thr:      0.25
diff --git a/configs/yolov8-seg.yaml b/configs/yolov8-seg.yaml
@@ -4,7 +4,6 @@ yolov8-seg:
     type:          "coco80"
     mode:          "fp32"
     batchSize:     1
-    inputChannel:  3
     imageWidth:    640
     imageHeight:   640
     conf_thr:      0.25
diff --git a/configs/yolov8.yaml b/configs/yolov8.yaml
@@ -1,10 +1,9 @@
 yolov8:
-    onnx_file:     "../weights/yolov8/yolov8m.onnx"
-    engine_file:   "../weights/yolov8/yolov8m.trt"
+    onnx_file:     "../weights/yolov8/yolov8s.onnx"
+    engine_file:   "../weights/yolov8/yolov8s.trt"
     type:          "coco80"
     mode:          "fp32"
     batchSize:     1
-    inputChannel:  3
     imageWidth:    640
     imageHeight:   640
     conf_thr:      0.25
diff --git a/include/basemodel.h b/include/basemodel.h
@@ -22,7 +22,6 @@ class Model
     std::string mode;
     AffineMatrix dst2src;
     int batchSize;
-    int inputChannel;
     int imageWidth;
     int imageHeight;
     std::string names[10];
diff --git a/include/cuda_function.h b/include/cuda_function.h
@@ -20,12 +20,15 @@ void preprocess(uint8_t* src, AffineMatrix d2s, int src_width, int src_height,
                      cudaStream_t stream);					 
 
 
-void postprocess_box(float* predict, int num_bboxes, int num_out, int num_classes, float conf_thr,
+void postprocess_box(float* predict, int num_bboxes, int num_classes, int num_out, float conf_thr,
 				float nms_thr, cudaStream_t stream, float* dst);
 
-void yolov8_postprocess_box(float* predict, int num_bboxes, int num_out, int num_classes, float conf_thr,
+void yolov8_postprocess_box(float* predict, int num_bboxes, int num_classes, int num_out, float conf_thr,
 				float nms_thr, cudaStream_t stream, float* dst);
 
+void rtdetr_postprocess_box(float* predict_box, float* predict_cls, int num_bboxes,  int num_classes, int num_out,
+							float conf_thr, int imageWidth, int imageHeight, cudaStream_t stream, float* dst);
+
 void postprocess_box_mask(float* predict, int num_bboxes, int num_classes, int num_out, 
 						  float conf_thr, float nms_thr, cudaStream_t stream, float* dst);
 
diff --git a/include/detection.h b/include/detection.h
@@ -81,27 +81,22 @@ class Detection : public Model
 public:
 
     explicit Detection(const YAML::Node &config);
-    std::vector<Detections> InferenceImages(std::vector<cv::Mat> &imgBatch) noexcept;
+    virtual std::vector<Detections> InferenceImages(std::vector<cv::Mat> &imgBatch) = 0;
     void Inference(const std::string &input_path, const std::string &save_path, const bool video) override;
-    virtual void Inference(const std::string &input_path, const std::string &save_path) override;
+    void Inference(const std::string &input_path, const std::string &save_path) override;
     void Visualize(const std::vector<Detections> &detections, std::vector<cv::Mat> &imgBatch,
                      std::vector<std::string> image_names);
     void Visualize(const std::vector<Detections> &detections, std::vector<cv::Mat> &imgBatch,
                      cv::String save_name, int fps, cv::Size size); 
-    cv::Rect get_rect(cv::Mat& img, float bbox[4]);
-    static float DIoU(const Box &det_a, const Box &det_b);
 
 protected:
-    virtual std::vector<Detections> PostProcess(const std::vector<cv::Mat> &vec_Mat, float* output)=0;
-    void NMS(std::vector<Box> &detections);    
+    virtual std::vector<Detections> PostProcess(const std::vector<cv::Mat> &vec_Mat, float* output);
     int num_classes;
     float conf_thr;
-    float nms_thr;
     std::string type;
     std::vector<std::string> class_labels;
     std::vector<cv::Scalar> class_colors;
-    std::vector<int> strides;
-    int num_rows = 0;
+    int num_bboxes = 0;
 };
 
 #endif
diff --git a/include/instance_segmentation.h b/include/instance_segmentation.h
@@ -22,28 +22,25 @@ class InstanceSegmentation : public Model
 {
 public:
     explicit InstanceSegmentation(const YAML::Node &config);
-    std::vector<Segmentations> InferenceImages(std::vector<cv::Mat> &imgBatch);
+    virtual std::vector<Segmentations> InferenceImages(std::vector<cv::Mat> &imgBatch) = 0;
     void Inference(const std::string &input_path, const std::string &save_path, const bool video) override;
-    virtual void Inference(const std::string &input_path, const std::string &save_path) override;
+    void Inference(const std::string &input_path, const std::string &save_path) override;
     void Visualize(const std::vector<Segmentations> &segmentations, std::vector<cv::Mat> &imgBatch,
                      std::vector<std::string> image_names);
     void Visualize(const std::vector<Segmentations> &segmentations, std::vector<cv::Mat> &imgBatch,
                      cv::String save_name, int fps, cv::Size size); 
-    static float DIoU(const Instance &det_a, const Instance &det_b);
     cv::Mat scale_mask(cv::Mat mask, cv::Mat img);
 
 protected:
     virtual std::vector<Segmentations> PostProcess(const std::vector<cv::Mat> &vec_Mat, float* output1, float* output2)=0;
-    void NMS(std::vector<Instance> &segmentations);
     uint8_t* cpu_mask_buffer = nullptr;
     int num_classes;
     float conf_thr;
     float nms_thr;
     std::string type;
     std::vector<std::string> class_labels;
     std::vector<cv::Scalar> class_colors;
-    std::vector<int> strides;
-    int num_rows = 0;
+    int num_bboxes = 0;
 };
 
 #endif
diff --git a/include/rtdetr.h b/include/rtdetr.h
@@ -7,8 +7,9 @@
 class RTDETR : public Detection {
 public:
     explicit RTDETR(const YAML::Node &config);
+    std::vector<Detections> InferenceImages(std::vector<cv::Mat> &imgBatch) noexcept;
 protected:
-    std::vector<Detections> PostProcess(const std::vector<cv::Mat> &vec_Mat, float* output) override;
+    std::vector<Detections> PostProcess(const std::vector<cv::Mat> &vec_Mat, float* output1, float* output2);
 };
 
 #endif
diff --git a/include/yolo.h b/include/yolo.h
@@ -7,16 +7,21 @@
 class YOLO : public Detection {
 public:
     explicit YOLO(const YAML::Node &config);
+    std::vector<Detections> InferenceImages(std::vector<cv::Mat> &imgBatch) noexcept;
 protected:
     std::vector<Detections> PostProcess(const std::vector<cv::Mat> &vec_Mat, float* output) override;
+    float nms_thr;
+    std::vector<int> strides;
 };
 
 class YOLO_seg : public InstanceSegmentation {
 public:
     explicit YOLO_seg(const YAML::Node &config);
-    cv::Rect get_downscale_rect(float bbox[4], float scale);
+    std::vector<Segmentations> InferenceImages(std::vector<cv::Mat> &imgBatch) noexcept;
 protected:
     std::vector<Segmentations> PostProcess(const std::vector<cv::Mat> &vec_Mat, float* output1, float* output2) override;
+    float nms_thr;
+    std::vector<int> strides;
 };
 
 #endif
diff --git a/src/basemodel.cpp b/src/basemodel.cpp
@@ -5,7 +5,6 @@ Model::Model(const YAML::Node &config) {
     engine_file = config["engine_file"].as<std::string>();
     mode = config["mode"].as<std::string>();
     batchSize = config["batchSize"].as<int>();
-    inputChannel = config["inputChannel"].as<int>();
     imageWidth = config["imageWidth"].as<int>();
     imageHeight = config["imageHeight"].as<int>();
     imgMean = config["imgMean"].as<std::vector<float>>();
@@ -44,9 +43,7 @@ void Model::OnnxToTRTModel() {
     
     std::ofstream file;
     file.open(engine_file, std::ios::binary | std::ios::out);
-    // sample::gLogInfo << "writing engine file..." << std::endl;
     file.write((const char *) data->data(), data->size());
-    // sample::gLogInfo << "save engine file done" << std::endl;
     file.close();
 
     delete parser;
diff --git a/src/cuda_function.cu b/src/cuda_function.cu
diff --git a/src/detection.cpp b/src/detection.cpp
diff --git a/src/instance_segmentation.cpp b/src/instance_segmentation.cpp
diff --git a/src/rtdetr.cpp b/src/rtdetr.cpp
diff --git a/src/yolo.cpp b/src/yolo.cpp
diff --git a/src/yolov8.cpp b/src/yolov8.cpp
