Update README.md

Author: Li-Hongda

README.md

@@ -1,7 +1,16 @@
 # TensorRT_Inference_Demo
 <div align="center">
-<img src="assets/000000005001.jpg" height="200px" >
-<img src="assets/000000007816.jpg" height="200px" >
+<img src="assets/000000005001.jpg" height="200px" width="240px" >
+<img src="assets/000000142324.jpg" height="200px" width="240px" >
+<img src="assets/000000007816.jpg" height="200px" width="240px" >
+<img src="assets/000000017899.jpg" height="200px" width="150px" >
+
+<img src="assets/000000157807.jpg" height="200px" width="240px" >
+<img src="assets/000000294695.jpg" height="200px" 
+width="240px" >
+<img src="assets/000000579158.jpg" height="200px" 
+width="240px" >
+<img src="assets/000000007977.jpg" height="200px" width="150px" >
 
 </div>
 
@@ -38,7 +47,7 @@ This repo use TensorRT-8.x to deploy well-trained models, both image preprocessi
 - [x] [YOLOv8](https://github.com/ultralytics/ultralytics)<br>
 - [x] [YOLOv8-seg](https://github.com/ultralytics/ultralytics)<br>
 - [x] [RT-DETR](https://github.com/PaddlePaddle/PaddleDetection/tree/develop/configs/rtdetr)<br>
-- [ ] [YOLOv6](https://github.com/meituan/YOLOv6)<br>
+- [x] [YOLOv6](https://github.com/meituan/YOLOv6)<br>
 - [ ] [YOLO-NAS](https://github.com/Deci-AI/super-gradients) (to be continued)<br>
 </details>
 

assets/000000017899.jpg

@@ -19,7 +19,7 @@ class Model {
     std::string onnx_file;
     std::string engine_file;
     std::string mode;
-    AffineMatrix dst2src;
+    std::vector<AffineMatrix> dst2src;
     int batchSize;
     int imageWidth;
     int imageHeight;

assets/000000142324.jpg

@@ -107,6 +107,7 @@ void Model::LoadEngine(){
 }
 
 void Model::PreProcess(std::vector<cv::Mat>& img_batch) {
+    dst2src.reserve(batchSize);
     for (size_t i = 0; i < img_batch.size(); i++) {
         int height = img_batch[i].rows; 
         int width = img_batch[i].cols;
@@ -115,9 +116,10 @@ void Model::PreProcess(std::vector<cv::Mat>& img_batch) {
         0.f, scale, (-scale * height + imageHeight + scale - 1) * 0.5);
         cv::Mat d2s = cv::Mat::zeros(2, 3, CV_32FC1);
         cv::invertAffineTransform(s2d, d2s);
-
-        memcpy(&dst2src, d2s.ptr(), sizeof(dst2src));
-        preprocess(img_batch[i].ptr(), dst2src, width, height, &gpu_buffers[0][bufferSize[0] * i], imageWidth, imageHeight, stream); 
+        AffineMatrix mat;
+        memcpy(&mat, d2s.ptr(), sizeof(mat));
+        dst2src.emplace_back(mat);
+        preprocess(img_batch[i].ptr(), mat, width, height, &gpu_buffers[0][bufferSize[0] * i], imageWidth, imageHeight, stream); 
         CUDA_CHECK(cudaStreamSynchronize(stream));
     }
 }

assets/000000157807.jpg

@@ -39,7 +39,7 @@ std::vector<Detections> RTDETR::PostProcess(const std::vector<cv::Mat> &imgBatch
         float* score_per_img = output2 + index * predscoreSize;
         cuda_postprocess_init(6, imageWidth, imageHeight);
         rtdetr_postprocess_box(box_per_img, score_per_img, num_bboxes, num_classes, 6, 
-                               conf_thr, imageWidth, imageHeight, dst2src, stream, cpu_buffers[2]);
+                               conf_thr, imageWidth, imageHeight, dst2src[index], stream, cpu_buffers[2]);
         int num_boxes = std::min((int)cpu_buffers[2][0], 300);
         for (int i = 0; i < num_boxes; i++) {
             Box box;
@@ -56,5 +56,6 @@ std::vector<Detections> RTDETR::PostProcess(const std::vector<cv::Mat> &imgBatch
         index++;
         // cuda_postprocess_destroy();
     }
+    dst2src.clear();    
     return vec_result;
 }

assets/000000294695.jpg

@@ -43,7 +43,7 @@ std::vector<Detections> YOLO::PostProcess(const std::vector<cv::Mat> &imgBatch,
         Detections result;
         float* pred_per_img = output + index * predSize;
         cuda_postprocess_init(7, imageWidth, imageHeight);
-        postprocess_box(pred_per_img, num_bboxes, num_classes, 7, conf_thr, nms_thr, dst2src, stream, cpu_buffers[1]);
+        postprocess_box(pred_per_img, num_bboxes, num_classes, 7, conf_thr, nms_thr, dst2src[index], stream, cpu_buffers[1]);
         int num_boxes = std::min((int)cpu_buffers[1][0], 1000);
         for (int i = 0; i < num_boxes; i++) {
             Box box;
@@ -59,7 +59,8 @@ std::vector<Detections> YOLO::PostProcess(const std::vector<cv::Mat> &imgBatch,
         }
         vec_result.emplace_back(result);
         index++;
-    }    
+    }
+    dst2src.clear();
     return vec_result;
 }
 
@@ -107,7 +108,7 @@ std::vector<Segmentations> YOLO_seg::PostProcess(const std::vector<cv::Mat> &img
         Segmentations result;
         float* proto = output1 + index * protoSize;
         float* pred_per_img = output2 + index * predSize;
-        postprocess_box_mask(pred_per_img, num_bboxes, num_classes, 39, conf_thr, nms_thr, dst2src, stream, cpu_buffers[1]);
+        postprocess_box_mask(pred_per_img, num_bboxes, num_classes, 39, conf_thr, nms_thr, dst2src[index], stream, cpu_buffers[1]);
         int num_boxes = std::min((int)cpu_buffers[1][0], 1000);
         for (int i = 0; i < num_boxes; i++) {
             Instance ins;
@@ -130,5 +131,6 @@ std::vector<Segmentations> YOLO_seg::PostProcess(const std::vector<cv::Mat> &img
         index++;
     }
     // cuda_postprocess_destroy();
+    dst2src.clear();    
     return vec_result;
 }