linghu8812 / tensorrt_inference

703 stars 206 forks source link

Yolov5 classification & Segmentation #152

Open sctrueew opened 2 years ago

sctrueew commented 2 years ago

Hi, Thanks for your work, Do you have any plans for implementing Yolov5 classification and segmentation?

Thanks

linghu8812 commented 2 years ago

Hi, Thanks for your work, Do you have any plans for implementing Yolov5 classification and segmentation?

Thanks

https://github.com/linghu8812/tensorrt_inference/commit/82838ba34bb3b2ddab338e39aafafe34dd894e66 support yolov5 classification models, please test it!

sctrueew commented 2 years ago

@linghu8812 Hi Thanks for the implementation. I'll give it a try

sctrueew commented 2 years ago

@linghu8812 Hi, I tested the classification task but the result it's not right:

OS: win10 CV: 4.5.4 TRT: 8.2.0.6 ONNX: onnxruntime.gpu.1.11.0 Config:

yolov5_cls:
    onnx_file:     "./yolov5l-cls.onnx"
    engine_file:   "./yolov5l-cls.trt"
    labels_file:   "./configs/yolov5/imagenet-classes.names"
    BATCH_SIZE:    1
    INPUT_CHANNEL: 3
    IMAGE_WIDTH:   224
    IMAGE_HEIGHT:  224
    image_order:   "BCHW"
    channel_order: "BGR"
    img_mean:      [ 0, 0, 0 ]
    img_std:       [ 1, 1, 1 ]
    alpha:         255.0
    resize:        "directly"
    YAML::Node root = YAML::LoadFile(config_file);
    YOLOv5_cls YOLOv5(root["yolov5_cls"]);
    YOLOv5.LoadEngine();
    YOLOv5.InferenceFolder(folder_name);

result: Processing: ./samples/classification/train.jpg classification prepare image take: 2.4224 ms. classification inference take: 6.2786 ms. classification postprocess take: 0.0042 ms. classic: 623 :: lens cap, lens cover :: 4.26172 <=== I modified the code to print this line ` void Classification::DrawResults(const std::vector& results, std::vector& vec_img,

std::vector<std::string> image_names = std::vector<std::string>()) {

for (int i = 0; i < (int)vec_img.size(); i++) {
    auto org_img = vec_img[i];
    if (!org_img.data)
        continue;
    auto result = results[i];
    auto classId = result.classes;
    if (!image_names.empty()) {
        std::string rst_name = class_labels[classId];// +".jpg";;
        std::cout << result.classes << " :: " << rst_name << " :: " << result.prob << std::endl;
    }
}}
linghu8812 commented 2 years ago

@linghu8812 Hi, I tested the classification task but the result it's not right:

OS: win10 CV: 4.5.4 TRT: 8.2.0.6 ONNX: onnxruntime.gpu.1.11.0 Config:

yolov5_cls:
    onnx_file:     "./yolov5l-cls.onnx"
    engine_file:   "./yolov5l-cls.trt"
    labels_file:   "./configs/yolov5/imagenet-classes.names"
    BATCH_SIZE:    1
    INPUT_CHANNEL: 3
    IMAGE_WIDTH:   224
    IMAGE_HEIGHT:  224
    image_order:   "BCHW"
    channel_order: "BGR"
    img_mean:      [ 0, 0, 0 ]
    img_std:       [ 1, 1, 1 ]
    alpha:         255.0
    resize:        "directly"
    YAML::Node root = YAML::LoadFile(config_file);
    YOLOv5_cls YOLOv5(root["yolov5_cls"]);
    YOLOv5.LoadEngine();
    YOLOv5.InferenceFolder(folder_name);

result: Processing: ./samples/classification/train.jpg classification prepare image take: 2.4224 ms. classification inference take: 6.2786 ms. classification postprocess take: 0.0042 ms. classic: 623 :: lens cap, lens cover :: 4.26172 <=== I modified the code to print this line ` void Classification::DrawResults(const std::vector& results, std::vectorcv::Mat& vec_img,

std::vector<std::string> image_names = std::vector<std::string>()) {

for (int i = 0; i < (int)vec_img.size(); i++) {
  auto org_img = vec_img[i];
  if (!org_img.data)
      continue;
  auto result = results[i];
  auto classId = result.classes;
  if (!image_names.empty()) {
      std::string rst_name = class_labels[classId];// +".jpg";;
      std::cout << result.classes << " :: " << rst_name << " :: " << result.prob << std::endl;
  }
}}

my result for train.jpg is:

Processing: ../samples/classification/train.jpg
classification prepare image take: 0.861332 ms.
classification inference take: 0.634246 ms.
classification postprocess take: 0.004999 ms.
bullet train, bullet.jpg
sctrueew commented 2 years ago

Could you please tell me what steps I need to check again? I'm not sure what the problem is.