Add yolov5 version6 model (#1204)

* add yolov5s version6 model

* add flag to control cut focus from model
Co-authored-by: Nchao.tang <chao.tang@lynxi.com>

examples/tm_yolov5s.cpp

@@ -38,6 +38,8 @@
 #include "tengine/c_api.h"
 #include "tengine_operations.h"
 
+#define YOLOV5_VERSION6 1
+
 struct Object
 {
     cv::Rect_<float> rect;
@@ -334,6 +336,60 @@ void get_input_data_focus(const char* image_file, float* input_data, int letterb
     }
 }
 
+void get_input_data_nofocus(const char* image_file, float* input_data, int letterbox_rows, int letterbox_cols, const float* mean, const float* scale)
+{
+    cv::Mat sample = cv::imread(image_file, 1);
+    cv::Mat img;
+
+    if (sample.channels() == 1)
+        cv::cvtColor(sample, img, cv::COLOR_GRAY2RGB);
+    else
+        cv::cvtColor(sample, img, cv::COLOR_BGR2RGB);
+
+    /* letterbox process to support different letterbox size */
+    float scale_letterbox;
+    int resize_rows;
+    int resize_cols;
+    if ((letterbox_rows * 1.0 / img.rows) < (letterbox_cols * 1.0 / img.cols))
+    {
+        scale_letterbox = letterbox_rows * 1.0 / img.rows;
+    }
+    else
+    {
+        scale_letterbox = letterbox_cols * 1.0 / img.cols;
+    }
+    resize_cols = int(scale_letterbox * img.cols);
+    resize_rows = int(scale_letterbox * img.rows);
+
+    cv::resize(img, img, cv::Size(resize_cols, resize_rows));
+    img.convertTo(img, CV_32FC3);
+    // Generate a gray image for letterbox using opencv
+    cv::Mat img_new(letterbox_cols, letterbox_rows, CV_32FC3, cv::Scalar(0.5 / scale[0] + mean[0], 0.5 / scale[1] + mean[1], 0.5 / scale[2] + mean[2]));
+    int top = (letterbox_rows - resize_rows) / 2;
+    int bot = (letterbox_rows - resize_rows + 1) / 2;
+    int left = (letterbox_cols - resize_cols) / 2;
+    int right = (letterbox_cols - resize_cols + 1) / 2;
+    // Letterbox filling
+    cv::copyMakeBorder(img, img_new, top, bot, left, right, cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0));
+
+    img_new.convertTo(img_new, CV_32FC3);
+    float* img_data = (float*)img_new.data;
+
+    /* nhwc to nchw */
+    for (int h = 0; h < letterbox_rows; h++)
+    {
+        for (int w = 0; w < letterbox_cols; w++)
+        {
+            for (int c = 0; c < 3; c++)
+            {
+                int in_index = h * letterbox_cols * 3 + w * 3 + c;
+                int out_index = c * letterbox_rows * letterbox_cols + h * letterbox_cols + w;
+                input_data[out_index] = (img_data[in_index] - mean[c]) * scale[c];
+            }
+        }
+    }
+}
+
 int main(int argc, char* argv[])
 {
     const char* model_file = nullptr;
@@ -424,7 +480,11 @@ int main(int argc, char* argv[])
     }
 
     int img_size = letterbox_rows * letterbox_cols * img_c;
+#if YOLOV5_VERSION6
+    int dims[] = {1, 3, letterbox_rows, letterbox_cols};
+#else
     int dims[] = {1, 12, int(letterbox_rows / 2), int(letterbox_cols / 2)};
+#endif
     std::vector<float> input_data(img_size);
 
     tensor_t input_tensor = get_graph_input_tensor(graph, 0, 0);
@@ -454,7 +514,11 @@ int main(int argc, char* argv[])
     }
 
     /* prepare process input data, set the data mem to input tensor */
+#if YOLOV5_VERSION6
+    get_input_data_nofocus(image_file, input_data.data(), letterbox_rows, letterbox_cols, mean, scale);
+#else
     get_input_data_focus(image_file, input_data.data(), letterbox_rows, letterbox_cols, mean, scale);
+#endif
 
     /* run graph */
     double min_time = DBL_MAX;

tools/optimize/yolov5s-opt.py

@@ -50,6 +50,7 @@ def parse_args():
     parser.add_argument('--output', help='output model path', default='./yolov5s-opt.onnx', type=str)
     parser.add_argument('--in_tensor', help='input tensor name', default='167', type=str)
     parser.add_argument('--out_tensor', help='output tensor names', default='381,420,459', type=str)
+    parser.add_argument('--cut_focus', action='store_false', help='cut focus from model if true')
     parser.add_argument('--verbose', action='store_true', help='show verbose info')
     
     args = parser.parse_args()
@@ -180,8 +181,9 @@ def main():
     new_nodes = old_node[:]
 
     # cut the focus and postprocess nodes
-    print("[Quant Tools Info]: Step 1, Remove the focus and postprocess nodes.")
-    new_nodes = cut_focus_output(old_node, in_tensor, out_tensor)
+    if args.cut_focus:
+        print("[Quant Tools Info]: Step 1, Remove the focus and postprocess nodes.")
+        new_nodes = cut_focus_output(old_node, in_tensor, out_tensor)
 
     # op fusion, using HardSwish replace the Sigmoid and Mul
     print("[Quant Tools Info]: Step 2, Using hardswish replace the sigmoid and mul.")