Fix normalization

Li-Hongda · Li-Hongda · commit 84bbd232c71f · 2023-06-04T16:53:10.000+08:00
diff --git a/README.md b/README.md
@@ -45,12 +45,12 @@ This repo use TensorRT-8.x to deploy well-trained models, both image preprocessi
 
 - [x] [YOLOv5](https://github.com/ultralytics/yolov5)<br>
 - [x] [YOLOv5-seg](https://github.com/ultralytics/yolov5)<br>
+- [x] [YOLOv6](https://github.com/meituan/YOLOv6)<br>
 - [x] [YOLOv7](https://github.com/WongKinYiu/yolov7)<br>
 - [x] [YOLOv8](https://github.com/ultralytics/ultralytics)<br>
 - [x] [YOLOv8-seg](https://github.com/ultralytics/ultralytics)<br>
+- [x] [YOLOX](https://github.com/Megvii-BaseDetection/YOLOX)<br>
 - [x] [RT-DETR](https://github.com/PaddlePaddle/PaddleDetection/tree/develop/configs/rtdetr)<br>
-- [x] [YOLOv6](https://github.com/meituan/YOLOv6)<br>
-- [ ] [YOLO-NAS](https://github.com/Deci-AI/super-gradients) (to be continued)<br>
 </details>
 
 All speed tests were performed on RTX 3090 with COCO Val set.The time calculated here is the sum of the time of image loading, preprocess, inference and postprocess, so it's going to be slower than what's reported in the paper.
diff --git a/configs/rtdetr.yaml b/configs/rtdetr.yaml
@@ -9,5 +9,6 @@ rtdetr:
     imageHeight:   640
     conf_thr:      0.5
     num_queries:   300
+    imgScale:      255
     imgMean:       [ 0, 0, 0 ]
     imgStd:        [ 1, 1, 1 ]
diff --git a/configs/yolov5-seg.yaml b/configs/yolov5-seg.yaml
@@ -10,5 +10,6 @@ yolov5-seg:
     conf_thr:      0.25
     nms_thr:       0.45
     strides:       [8, 16, 32]
+    imgScale:      255
     imgMean:       [ 0, 0, 0 ]
     imgStd:        [ 1, 1, 1 ] 
diff --git a/configs/yolov5.yaml b/configs/yolov5.yaml
@@ -10,5 +10,6 @@ yolov5:
     conf_thr:      0.25
     nms_thr:       0.45
     strides:       [8, 16, 32]
+    imgScale:      255
     imgMean:       [ 0, 0, 0 ]
     imgStd:        [ 1, 1, 1 ]
diff --git a/configs/yolov6.yaml b/configs/yolov6.yaml
@@ -10,5 +10,6 @@ yolov6:
     conf_thr:      0.25
     nms_thr:       0.45
     strides:       [8, 16, 32]
+    imgScale:      255
     imgMean:       [ 0, 0, 0 ]
     imgStd:        [ 1, 1, 1 ]
diff --git a/configs/yolov7-p6.yaml b/configs/yolov7-p6.yaml
@@ -10,5 +10,6 @@ yolov7:
     conf_thr:      0.4
     nms_thr:       0.45
     strides:       [8, 16, 32, 64]
+    imgScale:      255
     imgMean:       [ 0, 0, 0 ]
     imgStd:        [ 1, 1, 1 ]
diff --git a/configs/yolov7.yaml b/configs/yolov7.yaml
@@ -10,5 +10,6 @@ yolov7:
     conf_thr:      0.25
     nms_thr:       0.45
     strides:       [8, 16, 32]
+    imgScale:      255
     imgMean:       [ 0, 0, 0 ]
     imgStd:        [ 1, 1, 1 ]
diff --git a/configs/yolov8-seg.yaml b/configs/yolov8-seg.yaml
@@ -1,27 +1,15 @@
 yolov8-seg:
-    onnx_file:     "../weights/yolov8/yolov8s1280-seg.onnx"
-    engine_file:   "../weights/yolov8/yolov8s1280-seg.trt"
+    onnx_file:     "../weights/yolov8/yolov8s-seg.onnx"
+    engine_file:   "../weights/yolov8/yolov8s-seg.trt"
     type:          "coco"
     mode:          "fp32"
     dynamic:       1
     batchSize:     8
-    imageWidth:    1280
-    imageHeight:   1280
-    conf_thr:      0.45
+    imageWidth:    640
+    imageHeight:   640
+    conf_thr:      0.25
     nms_thr:       0.45
     strides:       [8, 16, 32]
+    imgScale:      255
     imgMean:       [ 0, 0, 0 ]
     imgStd:        [ 1, 1, 1 ]
-    # onnx_file:     "../weights/yolov8/yolov8s-seg.onnx"
-    # engine_file:   "../weights/yolov8/yolov8s-seg.trt"
-    # type:          "coco"
-    # mode:          "fp32"
-    # dynamic:       1
-    # batchSize:     8
-    # imageWidth:    640
-    # imageHeight:   640
-    # conf_thr:      0.25
-    # nms_thr:       0.45
-    # strides:       [8, 16, 32]
-    # imgMean:       [ 0, 0, 0 ]
-    # imgStd:        [ 1, 1, 1 ]
diff --git a/configs/yolov8.yaml b/configs/yolov8.yaml
@@ -10,5 +10,6 @@ yolov8:
     conf_thr:      0.25
     nms_thr:       0.45
     strides:       [8, 16, 32]
+    imgScale:      255
     imgMean:       [ 0, 0, 0 ]
     imgStd:        [ 1, 1, 1 ]
diff --git a/configs/yolox.yaml b/configs/yolox.yaml
@@ -0,0 +1,15 @@
+yolox:
+    onnx_file:     "../weights/yolox/yoloxs.onnx"
+    engine_file:   "../weights/yolox/yoloxs.trt"
+    type:          "coco"
+    mode:          "fp32"
+    dynamic:       1
+    batchSize:     1
+    imageWidth:    640
+    imageHeight:   640
+    conf_thr:      0.25
+    nms_thr:       0.45
+    strides:       [8, 16, 32]
+    imgScale:      1
+    imgMean:       [ 0, 0, 0 ]
+    imgStd:        [ 1, 1, 1 ]    
diff --git a/include/basemodel.h b/include/basemodel.h
@@ -31,7 +31,9 @@ class Model {
     std::unique_ptr<nvinfer1::IExecutionContext> context;
 
     cudaStream_t stream;
-    std::vector<float> imgMean;
-    std::vector<float> imgStd;   
+    Norm norm;
+    // float imgScale;
+    // std::vector<float> imgMean;
+    // std::vector<float> imgStd;   
 };
 #endif
diff --git a/include/build.h b/include/build.h
@@ -6,6 +6,7 @@
 #include "yolov6.h"
 #include "yolov7.h"
 #include "yolov8.h"
+#include "yolox.h"
 #include "rtdetr.h"
 
 std::shared_ptr<Model> build_model(std::string model_arch, std::string cfg);
diff --git a/include/cuda_function.h b/include/cuda_function.h
@@ -3,6 +3,11 @@
 #include "common.h"
 #include <cstdint>
 
+struct Norm {
+    float mean[3];
+    float std[3];
+	float scale;
+};
 
 struct AffineMatrix {
     float v0, v1, v2;
@@ -18,7 +23,7 @@ void cuda_postprocess_destroy();
 void process_mask_init(int num_out, int width, int height);
 
 void preprocess(uint8_t* src, AffineMatrix d2s, int src_width, int src_height,
-                     float* dst, int dst_width, int dst_height,
+                     float* dst, int dst_width, int dst_height, Norm norm,
                      cudaStream_t stream);					 
 
 void postprocess_box(float* predict, int num_bboxes, int num_classes, int num_out, float conf_thr,
diff --git a/include/detection.h b/include/detection.h
@@ -61,8 +61,7 @@ struct Detections {
     std::vector<Box> dets;
 };
 
-class Detection : public Model
-{
+class Detection : public Model {
 public:
 
     explicit Detection(const YAML::Node &config);
diff --git a/include/instance_segmentation.h b/include/instance_segmentation.h
@@ -17,8 +17,7 @@ struct Segmentations {
     std::vector<Instance> segs;
 };
 
-class InstanceSegmentation : public Model
-{
+class InstanceSegmentation : public Model {
 public:
     explicit InstanceSegmentation(const YAML::Node &config);
     virtual std::vector<Segmentations> InferenceImages(std::vector<cv::Mat> &imgBatch) = 0;
diff --git a/include/yolox.h b/include/yolox.h
@@ -0,0 +1,11 @@
+#ifndef YOLOX_H
+#define YOLOX_H
+
+#include "yolov8.h"
+
+class YOLOX : public YOLO {
+public:
+    explicit YOLOX(const YAML::Node &config);
+};
+
+#endif
diff --git a/object_detection/CMakeLists.txt b/object_detection/CMakeLists.txt
@@ -65,11 +65,12 @@ add_subdirectory(yolov5)
 add_subdirectory(yolov6)
 add_subdirectory(yolov7)
 add_subdirectory(yolov8)
+add_subdirectory(yolox)
 add_subdirectory(rtdetr)
 set(LIBRARY_OUTPUT_PATH ../libs)
 add_library(build SHARED ../src/build.cpp)
 
 set(EXECUTABLE_OUTPUT_PATH ../../bin)
 add_executable(object_detection main.cpp)
-target_link_libraries(object_detection yaml-cpp build yolov5 yolov6 yolov7 yolov8 rtdetr cudart)
+target_link_libraries(object_detection yaml-cpp build yolov5 yolov6 yolov7 yolov8 yolox rtdetr cudart)
 
diff --git a/object_detection/rtdetr/main.cpp b/object_detection/rtdetr/main.cpp
@@ -1,7 +1,6 @@
 #include "../../include/rtdetr.h"
 
-int main(int argc, char** argv)
-{
+int main(int argc, char** argv) {
     std::string inputpath = argv[1];
     std::string cfg_dir = "../configs";
     std::string cfg_suffix = ".yaml";
diff --git a/object_detection/yolov5/main.cpp b/object_detection/yolov5/main.cpp
@@ -1,8 +1,7 @@
 #include "../../include/yolov5.h"
 
 
-int main(int argc, char** argv)
-{
+int main(int argc, char** argv) {
     std::string inputpath = argv[1];
     std::string cfg_dir = "../configs";
     std::string cfg_suffix = ".yaml";
diff --git a/object_detection/yolov6/main.cpp b/object_detection/yolov6/main.cpp
@@ -1,8 +1,7 @@
 #include "../../include/yolov6.h"
 
 
-int main(int argc, char** argv)
-{
+int main(int argc, char** argv) {
     std::string inputpath = argv[1];
     std::string cfg_dir = "../configs";
     std::string cfg_suffix = ".yaml";
diff --git a/object_detection/yolov7/main.cpp b/object_detection/yolov7/main.cpp
@@ -1,8 +1,7 @@
 #include "../../include/yolov7.h"
 
 
-int main(int argc, char** argv)
-{
+int main(int argc, char** argv) {
     std::string inputpath = argv[1];
     std::string cfg_dir = "../configs";
     std::string cfg_suffix = ".yaml";
diff --git a/object_detection/yolov8/main.cpp b/object_detection/yolov8/main.cpp
@@ -1,8 +1,7 @@
 #include "../../include/yolov8.h"
 
 
-int main(int argc, char** argv)
-{
+int main(int argc, char** argv) {
     std::string inputpath = argv[1];
     std::string cfg_dir = "../configs";
     std::string cfg_suffix = ".yaml";
diff --git a/object_detection/yolox/CMakeLists.txt b/object_detection/yolox/CMakeLists.txt
@@ -0,0 +1,32 @@
+cmake_minimum_required(VERSION 3.10)
+
+set(CMAKE_BUILD_TYPE "Debug")
+# set(CMAKE_BUILD_TYPE "Release")
+
+project(yolox)
+
+set(CMAKE_CXX_STANDARD 14)
+
+# YAML
+set(YAML_LIBRARY ../../yaml-cpp/build)
+set(YAML_INCLUDE ../../yaml-cpp/include)
+link_directories(${YAML_LIBRARY})
+list(APPEND ALL_INCLUDE ${YAML_INCLUDE})
+
+
+# Project
+set(PROJECT_INCLUDE  ../../include)
+message(STATUS "Find project include at ${PROJECT_INCLUDE}")
+list(APPEND ALL_INCLUDE ${PROJECT_INCLUDE})
+set(PROJECT_LIBRARY ../../src/common.cpp  
+                    ../../src/basemodel.cpp 
+                    ../../src/detection.cpp 
+                    ../../src/yolo.cpp 
+                    ../../src/yolox.cpp
+                    ../../src/cuda_function.cu)
+
+include_directories(${ALL_INCLUDE})
+
+set(LIBRARY_OUTPUT_PATH ../../libs)
+add_library(yolox SHARED ${PROJECT_LIBRARY} ${SAMPLES_LIBRARY})
+target_link_libraries(yolox ${ALL_LIBS} yaml-cpp)
diff --git a/object_detection/yolox/main.cpp b/object_detection/yolox/main.cpp
@@ -0,0 +1,16 @@
+#include "../../include/yolox.h"
+
+
+int main(int argc, char** argv) {
+    std::string inputpath = argv[1];
+    std::string cfg_dir = "../configs";
+    std::string cfg_suffix = ".yaml";
+    std::string savedir = "../results";
+    auto savepath = savedir + "/" + "yolox" + "/";
+    auto cfg = cfg_dir + "/" + "yolox" + cfg_suffix;    
+    YAML::Node root = YAML::LoadFile(cfg);
+    YOLOX YOLOX(root["yolox"]);
+    YOLOX.LoadEngine();
+    YOLOX.Inference(inputpath, savepath);
+    return 0;
+}
diff --git a/src/basemodel.cpp b/src/basemodel.cpp
@@ -8,8 +8,12 @@ Model::Model(const YAML::Node &config) {
     batchSize = config["batchSize"].as<int>();
     imageWidth = config["imageWidth"].as<int>();
     imageHeight = config["imageHeight"].as<int>();
-    imgMean = config["imgMean"].as<std::vector<float>>();
-    imgStd = config["imgStd"].as<std::vector<float>>();
+    auto imgMean = config["imgMean"].as<std::vector<float>>();
+    auto imgStd = config["imgStd"].as<std::vector<float>>();
+    auto imgScale = config["imgScale"].as<float>();
+    memcpy(norm.mean, imgMean.data(), sizeof(float) * 3);
+    memcpy(norm.std, imgStd.data(), sizeof(float) * 3);
+    norm.scale = imgScale;
 }
 
 Model::~Model() {
@@ -107,7 +111,7 @@ void Model::LoadEngine(){
     for (int i = 0; i < nbBindings; ++i) {
         nvinfer1::Dims dims = engine->getBindingDimensions(i);
         if (dims.d[0] == -1)
-                dims.d[0] = batchSize;
+            dims.d[0] = batchSize;
         nvinfer1::DataType dtype = engine->getBindingDataType(i);
         int64_t totalSize = sample::volume(dims) * sample::dataTypeSize(dtype);
         bufferSize[i] = totalSize;
@@ -136,7 +140,7 @@ void Model::PreProcess(std::vector<cv::Mat>& img_batch) {
         AffineMatrix mat;
         memcpy(&mat, d2s.ptr(), sizeof(mat));
         dst2src.emplace_back(mat);
-        preprocess(img_batch[i].ptr(), mat, width, height, &gpu_buffers[0][size * i], imageWidth, imageHeight, stream); 
+        preprocess(img_batch[i].ptr(), mat, width, height, &gpu_buffers[0][size * i], imageWidth, imageHeight, norm, stream); 
         CUDA_CHECK(cudaStreamSynchronize(stream));
     }
 }
diff --git a/src/build.cpp b/src/build.cpp
@@ -15,8 +15,10 @@ std::shared_ptr<Model> build_model(std::string model_arch, std::string cfg) {
         model = std::make_shared<YOLOv8>(root[model_arch]);
     else if (model_arch == "yolov8-seg")
         model = std::make_shared<YOLOv8_seg>(root[model_arch]); 
+    else if (model_arch == "yolox")
+        model = std::make_shared<YOLOX>(root[model_arch]);
     else if (model_arch == "rtdetr")
-        model = std::make_shared<RTDETR>(root[model_arch]);                
+        model = std::make_shared<RTDETR>(root[model_arch]);
     else
         std::cout << "No model arch matched!" << std::endl;
     return model;
diff --git a/src/common.cpp b/src/common.cpp
@@ -14,20 +14,16 @@ bool __check_cuda_runtime(cudaError_t code, const char* op, const char* file, in
 	return true;
 }
 
-std::vector<std::string> get_names(const std::string &image_path)
-{
+std::vector<std::string> get_names(const std::string &image_path) {
     std::vector<std::string> image_names;
     auto dir = opendir(image_path.c_str());
 
-    if ((dir) != nullptr)
-    {
+    if ((dir) != nullptr) {
         struct dirent *entry;
         entry = readdir(dir);
-        while (entry)
-        {
+        while (entry) {
             auto temp = image_path + "/" + entry->d_name;
-            if (strcmp(entry->d_name, "") == 0 || strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0)
-            {
+            if (strcmp(entry->d_name, "") == 0 || strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) {
                 entry = readdir(dir);
                 continue;
             }
@@ -80,17 +76,16 @@ static void createDirectory(const std::string &path) {
 int check_dir(const std::string & path, const bool is_mkdir) noexcept {
     DIR *p_dir;
     struct dirent *entry;
-    if(isPathExists(path)){
-        if(isDirectory(path)){
+    if(isPathExists(path)) {
+        if(isDirectory(path)) {
             if((p_dir = opendir(path.c_str())) == NULL ) {  
                 std::cout << "Opendir error: " << strerror(errno) << std::endl;
                 return -1;  
             }  
     
-            while((entry = readdir(p_dir)) != NULL){
+            while((entry = readdir(p_dir)) != NULL) {
                 std::string file_name = path + "/" + entry->d_name;
-                if((0 != strcmp(entry->d_name, ".")) && (0 != strcmp(entry->d_name, "..")))
-                {
+                if((0 != strcmp(entry->d_name, ".")) && (0 != strcmp(entry->d_name, ".."))) {
                     remove(file_name.c_str());
                 }
             }
diff --git a/src/cuda_function.cu b/src/cuda_function.cu
diff --git a/src/instance_segmentation.cpp b/src/instance_segmentation.cpp
diff --git a/src/yolo.cpp b/src/yolo.cpp
diff --git a/src/yolov6.cpp b/src/yolov6.cpp
diff --git a/src/yolov8.cpp b/src/yolov8.cpp
diff --git a/src/yolox.cpp b/src/yolox.cpp
