Skip to content

Commit

Permalink
Merge branch 'wang-xinyu:master' into master
Browse files Browse the repository at this point in the history
  • Loading branch information
@mantuoluozk
mantuoluozk authored Jan 11, 2023
2 parents 7edc937 + 772d353 commit 7e4da87
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 21 deletions.
12 changes: 8 additions & 4 deletions yolov5/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
47 changes: 44 additions & 3 deletions yolov5/src/utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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_

52 changes: 38 additions & 14 deletions yolov5/yolov5_seg.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]);
Expand All @@ -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;
}
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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,
Expand All @@ -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);

}
}

Expand All @@ -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;
}

Expand Down Expand Up @@ -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];
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit 7e4da87

Please sign in to comment.