Draw label text for yolov5_seg (wang-xinyu#1197)

PrinceP · wang-xinyu · web-flow · commit 772d3533997b · 2023-01-11T10:54:16.000+08:00
* convert class id to class name + added confidence + different colors for class ids already present + background box to match with yolov5 repo + added coco.txt

* Required changes: removed coco.txt and *.jpg files, updated readme

* Resolving comments

* Update utils.h

* Update yolov5_seg.cpp

* Update README.md

Co-authored-by: Wang Xinyu &lt;shaywxy@gmail.com&gt;
diff --git a/yolov5/README.md b/yolov5/README.md
@@ -30,7 +30,8 @@ TensorRTx inference code base for [ultralytics/yolov5](https://github.com/ultral
 <a href="https://github.com/triple-Mu"><img src="https://avatars.githubusercontent.com/u/92794867?s=48&v=4" width="40px;" alt=""/></a>
 <a href="https://github.com/xiang-wuu"><img src="https://avatars.githubusercontent.com/u/107029401?s=48&v=4" width="40px;" alt=""/></a>
 <a href="https://github.com/uyolo1314"><img src="https://avatars.githubusercontent.com/u/101853326?s=48&v=4" width="40px;" alt=""/></a>
-<a href="https://github.com/Rex-LK"><img src="https://avatars.githubusercontent.com/u/74702576?s=96&v=4" width="40px;" alt=""/></a>
+<a href="https://github.com/Rex-LK"><img src="https://avatars.githubusercontent.com/u/74702576?s=48&v=4" width="40px;" alt=""/></a>
+<a href="https://github.com/PrinceP"><img src="https://avatars.githubusercontent.com/u/10251537?s=48&v=4" width="40px;" alt=""/></a>
 
 ## Different versions of yolov5
 
@@ -128,12 +129,15 @@ wget https://github.com/joannzhang00/ImageNet-dataset-classes-labels/blob/main/i
 # Build and serialize TensorRT engine
 ./yolov5_seg -s yolov5s-seg.wts yolov5s-seg.engine s
 
-# Run inference
-./yolov5_seg -d yolov5s-seg.engine ../samples
+# Download the labels file
+wget -O coco.txt https://raw.githubusercontent.com/amikelive/coco-labels/master/coco-labels-2014_2017.txt
+
+# Run inference with labels file
+./yolov5_seg -d yolov5s-seg.engine ../samples coco.txt
 ```
 
 <p align="center">
-<img src="https://user-images.githubusercontent.com/15235574/208305921-0a2ee358-6550-4d36-bb86-867685bfe069.jpg" height="360px;">
+<img src="https://user-images.githubusercontent.com/10251537/211291625-1b912483-b6a6-4e92-80c1-434d165b6776.jpg" height="360px;">
 </p>
 
 # INT8 Quantization
diff --git a/yolov5/src/utils.h b/yolov5/src/utils.h
@@ -3,11 +3,15 @@
 
 #include <dirent.h>
 #include <opencv2/opencv.hpp>
+#include <fstream>
+#include <unordered_map>
+#include <string>
+#include <sstream>
 
 static inline cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) {
     int w, h, x, y;
-    float r_w = input_w / (img.cols*1.0);
-    float r_h = input_h / (img.rows*1.0);
+    float r_w = input_w / (img.cols * 1.0);
+    float r_h = input_h / (img.rows * 1.0);
     if (r_h > r_w) {
         w = input_w;
         h = r_w * img.rows;
@@ -26,7 +30,7 @@ static inline cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) {
     return out;
 }
 
-static inline int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
+static inline int read_files_in_dir(const char* p_dir_name, std::vector<std::string>& file_names) {
     DIR *p_dir = opendir(p_dir_name);
     if (p_dir == nullptr) {
         return -1;
@@ -48,5 +52,42 @@ static inline int read_files_in_dir(const char *p_dir_name, std::vector<std::str
     return 0;
 }
 
+// Function to trim leading and trailing whitespace from a string
+static inline std::string trim_leading_whitespace(const std::string& str) {
+    size_t first = str.find_first_not_of(' ');
+    if (std::string::npos == first) {
+        return str;
+    }
+    size_t last = str.find_last_not_of(' ');
+    return str.substr(first, (last - first + 1));
+}
+// Src: https://stackoverflow.com/questions/16605967
+static inline std::string to_string_with_precision(const float a_value, const int n = 2) {
+    std::ostringstream out;
+    out.precision(n);
+    out << std::fixed << a_value;
+    return out.str();
+}
+
+static inline int read_labels(const std::string labels_filename, std::unordered_map<int, std::string>& labels_map) {
+
+    std::ifstream file(labels_filename);
+    // Read each line of the file
+    std::string line;
+    int index = 0;
+    while (std::getline(file, line)) {
+        // Strip the line of any leading or trailing whitespace
+        line = trim_leading_whitespace(line);
+
+        // Add the stripped line to the labels_map, using the loop index as the key
+        labels_map[index] = line;
+        index++;
+    }
+    // Close the file
+    file.close();
+
+    return 0;
+}
+
 #endif  // TRTX_YOLOV5_UTILS_H_
 
diff --git a/yolov5/yolov5_seg.cpp b/yolov5/yolov5_seg.cpp
@@ -155,7 +155,7 @@ void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffer
     cudaStreamSynchronize(stream);
 }
 
-bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, float& gd, float& gw, std::string& img_dir) {
+bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, float& gd, float& gw, std::string& img_dir, std::string& labels_filename) {
     if (argc < 4) return false;
     if (std::string(argv[1]) == "-s" && (argc == 5 || argc == 7)) {
         wts = std::string(argv[2]);
@@ -182,9 +182,10 @@ bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, fl
         } else {
             return false;
         }
-    } else if (std::string(argv[1]) == "-d" && argc == 4) {
+    } else if (std::string(argv[1]) == "-d" && argc == 5) {
         engine = std::string(argv[2]);
         img_dir = std::string(argv[3]);
+        labels_filename = std::string(argv[4]);
     } else {
         return false;
     }
@@ -216,12 +217,6 @@ std::vector<cv::Mat> process_mask(const float* proto, std::vector<Yolo::Detectio
         }
         e = 1.0f / (1.0f + expf(-e));
         mask_mat.at<float>(y, x) = e;
-        // if (e > 0.5) {
-        //   // TODO(Call for PR): Use different colors for different class ids
-        //   mask_mat.at<cv::Vec3b>(y, x)[2] = 0xFF;
-        //   mask_mat.at<cv::Vec3b>(y, x)[1] = 0x38;
-        //   mask_mat.at<cv::Vec3b>(y, x)[0] = 0x38;
-        // }
       }
     }
     cv::resize(mask_mat, mask_mat, cv::Size(INPUT_W, INPUT_H));
@@ -251,7 +246,7 @@ cv::Mat scale_mask(cv::Mat mask, cv::Mat img) {
   return res;
 }
 
-void draw_mask_bbox(cv::Mat& img, std::vector<Yolo::Detection>& dets, std::vector<cv::Mat>& masks) {
+void draw_mask_bbox(cv::Mat& img, std::vector<Yolo::Detection>& dets, std::vector<cv::Mat>& masks, std::unordered_map<int, std::string>& labels_map) {
   static std::vector<uint32_t> colors = {0xFF3838, 0xFF9D97, 0xFF701F, 0xFFB21D, 0xCFD231, 0x48F90A,
                                          0x92CC17, 0x3DDB86, 0x1A9334, 0x00D4BB, 0x2C99A8, 0x00C2FF,
                                          0x344593, 0x6473FF, 0x0018EC, 0x8438FF, 0x520085, 0xCB38FF,
@@ -273,8 +268,23 @@ void draw_mask_bbox(cv::Mat& img, std::vector<Yolo::Detection>& dets, std::vecto
     }
 
     cv::rectangle(img, r, bgr, 2);
-    // TODO(Call for PR): convert class id to class name
-    cv::putText(img, std::to_string((int)dets[i].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar::all(0xFF), 2);
+    
+    // Get the size of the text
+    cv::Size textSize = cv::getTextSize(labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::FONT_HERSHEY_PLAIN, 1.2, 2, NULL);
+    // Set the top left corner of the rectangle
+    cv::Point topLeft(r.x, r.y - textSize.height);
+
+    // Set the bottom right corner of the rectangle
+    cv::Point bottomRight(r.x + textSize.width, r.y + textSize.height);
+
+    // Set the thickness of the rectangle lines
+    int lineThickness = 2;
+
+    // Draw the rectangle on the image
+    cv::rectangle(img, topLeft, bottomRight, bgr, -1);
+
+    cv::putText(img, labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::Point(r.x, r.y + 4), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar::all(0xFF), 2);
+
   }
 }
 
@@ -283,12 +293,14 @@ int main(int argc, char** argv) {
 
     std::string wts_name = "";
     std::string engine_name = "";
+    std::string labels_filename = "";
+    
     float gd = 0.0f, gw = 0.0f;
     std::string img_dir;
-    if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) {
+    if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir, labels_filename)) {
         std::cerr << "arguments not right!" << std::endl;
         std::cerr << "./yolov5_seg -s [.wts] [.engine] [n/s/m/l/x or c gd gw]  // serialize model to plan file" << std::endl;
-        std::cerr << "./yolov5_seg -d [.engine] ../samples  // deserialize plan file and run inference" << std::endl;
+        std::cerr << "./yolov5_seg -d [.engine] ../samples coco.txt  // deserialize plan file, read the labels file and run inference" << std::endl;
         return -1;
     }
 
@@ -328,6 +340,18 @@ int main(int argc, char** argv) {
         std::cerr << "read_files_in_dir failed." << std::endl;
         return -1;
     }
+    
+    // read the txt file for classnames
+    std::ifstream labels_file(labels_filename, std::ios::binary);
+    if (!labels_file.good()) {
+        std::cerr << "read " << labels_filename << " error!" << std::endl;
+        return -1;
+    }
+    std::unordered_map<int, std::string> labels_map;
+    read_labels(labels_filename, labels_map);
+    
+    assert(CLASS_NUM == labels_map.size());
+
 
     static float prob[BATCH_SIZE * OUTPUT_SIZE1];
     static float proto[BATCH_SIZE * OUTPUT_SIZE2];
@@ -398,7 +422,7 @@ int main(int argc, char** argv) {
             cv::Mat img = imgs_buffer[b];
 
             auto masks = process_mask(&proto[b * OUTPUT_SIZE2], res);
-            draw_mask_bbox(img, res, masks);
+            draw_mask_bbox(img, res, masks, labels_map);
             cv::imwrite("_" + file_names[f - fcount + 1 + b], img);
         }
         fcount = 0;
