Merge pull request #25794 from Abdurrahheem:ash/yolov10-support

Add sample support of YOLOv9 and YOLOv10 in OpenCV #25794

This PR adds sample support of  [`YOLOv9`](https://github.com/WongKinYiu/yolov9) and [`YOLOv10`](https://github.com/THU-MIG/yolov10/tree/main)) in OpenCV. Models for this test are located in this [PR](https://github.com/opencv/opencv_extra/pull/1186). 

**Running YOLOv10 using OpenCV.** 
1. In oder to run `YOLOv10` one needs to cut off postporcessing with dynamic shapes from torch and then convert it to ONNX. If someone is looking for ready solution, there is [this forked branch](https://github.com/Abdurrahheem/yolov10/tree/ash/opencv-export) from official YOLOv10.  Particularty follow this proceduce. 

```bash
git clone git@github.com:Abdurrahheem/yolov10.git
conda create -n yolov10 python=3.9
conda activate yolov10
pip install -r requirements.txt
python export_opencv.py --model=<model-name> --imgsz=<input-img-size>
```
By default `model="yolov10s"` and `imgsz=(480,640)`. This will generate file `yolov10s.onnx`, which can be use for inference in OpenCV

2. For inference part on OpenCV.  one can use `yolo_detector.cpp` [sample](https://github.com/opencv/opencv/blob/4.x/samples/dnn/yolo_detector.cpp). If you have followed above exporting procedure, then you can use following command to run the model. 

``` bash
build opencv from source 
cd build 
./bin/example_dnn_yolo_detector --model=<path-to-yolov10s.onnx-file> --yolo=yolov10 --width=640 --height=480 --input=<path-to-image> --scale=0.003921568627 --padvalue=114
```
If you do not specify `--input` argument, OpenCV will grab first camera that is avaliable on your platform. 
For more deatils on how to run the `yolo_detector.cpp` file see this [guide](https://docs.opencv.org/4.x/da/d9d/tutorial_dnn_yolo.html#autotoc_md443) 


**Running YOLOv9 using OpenCV**

1. Export model following [official guide](https://github.com/WongKinYiu/yolov9)of the YOLOv9 repository. Particularly you can do following for converting.

```bash
git clone https://github.com/WongKinYiu/yolov9.git
cd yolov9
conda create -n yolov9 python=3.9
conda activate yolov9
pip install -r requirements.txt
wget https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-t-converted.pt
python export.py --weights=./yolov9-t-converted.pt --include=onnx --img-size=(480,640) 
```

This will generate <yolov9-t-converted.onnx> file.

2.  Inference on OpenCV.

```bash
build opencv from source 
cd build 
./bin/example_dnn_yolo_detector --model=<path-to-yolov9-t-converted.onnx> --yolo=yolov9 --width=640 --height=480 --scale=0.003921568627 --padvalue=114 --path=<path-to-image>
```

### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [x] The PR is proposed to the proper branch
- [x] There is a reference to the original bug report and related work
- [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [x] The feature is well documented and sample code can be built with the project CMake
pull/25857/head
Abduragim Shtanchaev 5 months ago committed by GitHub
parent 939cb58cd6
commit a8d1373919
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 51
      doc/tutorials/dnn/dnn_yolo/dnn_yolo.markdown
  2. 96
      modules/dnn/test/test_onnx_importer.cpp
  3. 32
      samples/dnn/yolo_detector.cpp

@ -24,7 +24,9 @@ model, but the methodology applies to other supported models.
@note Currently, OpenCV supports the following YOLO models: @note Currently, OpenCV supports the following YOLO models:
- [YOLOX](https://github.com/Megvii-BaseDetection/YOLOX/blob/main), - [YOLOX](https://github.com/Megvii-BaseDetection/YOLOX/blob/main),
- [YoloNas](https://github.com/Deci-AI/super-gradients/tree/master), - [YOLONas](https://github.com/Deci-AI/super-gradients/tree/master),
- [YOLOv10](https://github.com/THU-MIG/yolov10/tree/main),
- [YOLOv9](https://github.com/WongKinYiu/yolov9),
- [YOLOv8](https://github.com/ultralytics/ultralytics/tree/main), - [YOLOv8](https://github.com/ultralytics/ultralytics/tree/main),
- [YOLOv7](https://github.com/WongKinYiu/yolov7/tree/main), - [YOLOv7](https://github.com/WongKinYiu/yolov7/tree/main),
- [YOLOv6](https://github.com/meituan/YOLOv6/blob/main), - [YOLOv6](https://github.com/meituan/YOLOv6/blob/main),
@ -79,7 +81,7 @@ the ONNX graph, a process that we will detail further in the subsequent sections
Now that we know know the parameters of the pre-precessing we can go on and export the model from Now that we know know the parameters of the pre-precessing we can go on and export the model from
Pytorch to ONNX graph. Since in this tutorial we are using YOLOX as our sample model, lets use its Pytorch to ONNX graph. Since in this tutorial we are using YOLOX as our sample model, lets use its
export for demonstration purposes (the process is identical for the rest of the YOLO detectors). export for demonstration purposes (the process is identical for the rest of the YOLO detectors except `YOLOv10` model, see details on how to export it later in the post).
To exporting YOLOX we can just use [export script](https://github.com/Megvii-BaseDetection/YOLOX/blob/ac58e0a5e68e57454b7b9ac822aced493b553c53/tools/export_onnx.py). Particularly we need following commands: To exporting YOLOX we can just use [export script](https://github.com/Megvii-BaseDetection/YOLOX/blob/ac58e0a5e68e57454b7b9ac822aced493b553c53/tools/export_onnx.py). Particularly we need following commands:
@code{.bash} @code{.bash}
@ -125,6 +127,20 @@ than YOLOX) in case it is needed. However, usually each YOLO repository has pred
onnx.save(model_simp, args.output_name) onnx.save(model_simp, args.output_name)
@endcode @endcode
#### Exporting YOLOv10 model
In oder to run YOLOv10 one needs to cut off postporcessing with dynamic shapes from torch and then convert it to ONNX. If someone is looking for on how to cut off the postprocessing, there is this [forked branch](https://github.com/Abdurrahheem/yolov10/tree/ash/opencv-export) from official YOLOv10. The forked branch cuts of the postprocessing by [returning output](https://github.com/Abdurrahheem/yolov10/blob/4fdaafd912c8891642bfbe85751ea66ec20f05ad/ultralytics/nn/modules/head.py#L522) of the model before postprocessing procedure itself. To convert torch model to ONNX follow this proceduce.
@code{.bash}
git clone git@github.com:Abdurrahheem/yolov10.git
conda create -n yolov10 python=3.9
conda activate yolov10
pip install -r requirements.txt
python export_opencv.py --model=<model-name> --imgsz=<input-img-size>
@endcode
By default `--model="yolov10s"` and `--imgsz=(480,640)`. This will generate file `yolov10s.onnx`, which can be use for inference in OpenCV
### Running Yolo ONNX detector with OpenCV Sample ### Running Yolo ONNX detector with OpenCV Sample
Once we have our ONNX graph of the model, we just simply can run with OpenCV's sample. To that we need to make sure: Once we have our ONNX graph of the model, we just simply can run with OpenCV's sample. To that we need to make sure:
@ -144,24 +160,25 @@ Once we have our ONNX graph of the model, we just simply can run with OpenCV's s
--padvalue=<padding_value> \ --padvalue=<padding_value> \
--paddingmode=<padding_mode> \ --paddingmode=<padding_mode> \
--backend=<computation_backend> \ --backend=<computation_backend> \
--target=<target_computation_device> --target=<target_computation_device> \
--width=<model_input_width> \
--height=<model_input_height> \
@endcode @endcode
VIDEO DEMO:
@youtube{NHtRlndE2cg}
- --input: File path to your input image or video. If omitted, it will capture frames from a camera. - --input: File path to your input image or video. If omitted, it will capture frames from a camera.
- --classes: File path to a text file containing class names for object detection. - --classes: File path to a text file containing class names for object detection.
- --thr: Confidence threshold for detection (e.g., 0.5). - --thr: Confidence threshold for detection (e.g., 0.5).
- --nms: Non-maximum suppression threshold (e.g., 0.4). - --nms: Non-maximum suppression threshold (e.g., 0.4).
- --mean: Mean normalization value (e.g., 0.0 for no mean normalization). - --mean: Mean normalization value (e.g., 0.0 for no mean normalization).
- --scale: Scale factor for input normalization (e.g., 1.0). - --scale: Scale factor for input normalization (e.g., 1.0, 1/255.0, etc).
- --yolo: YOLO model version (e.g., YOLOv3, YOLOv4, etc.). - --yolo: YOLO model version (e.g., YOLOv3, YOLOv4, etc.).
- --padvalue: Padding value used in pre-processing (e.g., 114.0). - --padvalue: Padding value used in pre-processing (e.g., 114.0).
- --paddingmode: Method for handling image resizing and padding. Options: 0 (resize without extra processing), 1 (crop after resize), 2 (resize with aspect ratio preservation). - --paddingmode: Method for handling image resizing and padding. Options: 0 (resize without extra processing), 1 (crop after resize), 2 (resize with aspect ratio preservation).
- --backend: Selection of computation backend (0 for automatic, 1 for Halide, 2 for OpenVINO, etc.). - --backend: Selection of computation backend (0 for automatic, 1 for Halide, 2 for OpenVINO, etc.).
- --target: Selection of target computation device (0 for CPU, 1 for OpenCL, etc.). - --target: Selection of target computation device (0 for CPU, 1 for OpenCL, etc.).
- --device: Camera device number (0 for default camera). If `--input` is not provided camera with index 0 will used by default. - --device: Camera device number (0 for default camera). If `--input` is not provided camera with index 0 will used by default.
- --width: Model input width. Not to be confused with the image width. (e.g., 416, 480, 640, 1280, etc).
- --height: Model input height. Not to be confused with the image height. (e.g., 416, 480, 640, 1280, etc).
Here `mean`, `scale`, `padvalue`, `paddingmode` should exactly match those that we discussed Here `mean`, `scale`, `padvalue`, `paddingmode` should exactly match those that we discussed
in pre-processing section in order for the model to match result in PyTorch in pre-processing section in order for the model to match result in PyTorch
@ -183,7 +200,8 @@ cd <build directory of OpenCV>
./bin/example_dnn_yolo_detector ./bin/example_dnn_yolo_detector
@endcode @endcode
This will execute the YOLOX detector with your camera. For YOLOv8 (for instance), follow these additional steps: This will execute the YOLOX detector with your camera.
For YOLOv8 (for instance), follow these additional steps:
@code{.sh} @code{.sh}
cd opencv_extra/testdata/dnn cd opencv_extra/testdata/dnn
@ -195,6 +213,23 @@ cd <build directory of OpenCV>
./bin/example_dnn_yolo_detector --model=onnx/models/yolov8n.onnx --yolo=yolov8 --mean=0.0 --scale=0.003921568627 --paddingmode=2 --padvalue=144.0 --thr=0.5 --nms=0.4 --rgb=0 ./bin/example_dnn_yolo_detector --model=onnx/models/yolov8n.onnx --yolo=yolov8 --mean=0.0 --scale=0.003921568627 --paddingmode=2 --padvalue=144.0 --thr=0.5 --nms=0.4 --rgb=0
@endcode @endcode
For YOLOv10, follow these steps:
@code{.sh}
cd opencv_extra/testdata/dnn
python download_models.py yolov10
cd ..
export OPENCV_TEST_DATA_PATH=$(pwd)
cd <build directory of OpenCV>
./bin/example_dnn_yolo_detector --model=onnx/models/yolov8n.onnx --yolo=yolov10 --width=640 --height=480 --scale=0.003921568627 --padvalue=114
@endcode
This will run `YOLOv10` detector on first camera found on your system. If you want to run it on a image/video file, you can use `--input` option to specify the path to the file.
VIDEO DEMO:
@youtube{NHtRlndE2cg}
### Building a Custom Pipeline ### Building a Custom Pipeline

@ -19,7 +19,8 @@ void yoloPostProcessing(
std::vector<Rect2d>& keep_boxes, std::vector<Rect2d>& keep_boxes,
float conf_threshold, float conf_threshold,
float iou_threshold, float iou_threshold,
const std::string& test_name); const std::string& model_name,
const int nc=80);
template<typename TString> template<typename TString>
static std::string _tf(TString filename, bool required = true) static std::string _tf(TString filename, bool required = true)
@ -2670,7 +2671,8 @@ void yoloPostProcessing(
std::vector<Rect2d>& keep_boxes, std::vector<Rect2d>& keep_boxes,
float conf_threshold, float conf_threshold,
float iou_threshold, float iou_threshold,
const std::string& test_name const std::string& model_name,
const int nc
){ ){
// Retrieve // Retrieve
@ -2678,11 +2680,13 @@ void yoloPostProcessing(
std::vector<float> confidences; std::vector<float> confidences;
std::vector<Rect2d> boxes; std::vector<Rect2d> boxes;
if (test_name == "yolov8"){ if (model_name == "yolov8" || model_name == "yolov10" ||
model_name == "yolov9")
{
cv::transposeND(outs[0], {0, 2, 1}, outs[0]); cv::transposeND(outs[0], {0, 2, 1}, outs[0]);
} }
if (test_name == "yolonas"){ if (model_name == "yolonas"){
// outs contains 2 elemets of shape [1, 8400, 80] and [1, 8400, 4]. Concat them to get [1, 8400, 84] // outs contains 2 elemets of shape [1, 8400, 80] and [1, 8400, 4]. Concat them to get [1, 8400, 84]
Mat concat_out; Mat concat_out;
// squeeze the first dimension // squeeze the first dimension
@ -2696,22 +2700,27 @@ void yoloPostProcessing(
outs[0] = outs[0].reshape(0, std::vector<int>{1, 8400, 84}); outs[0] = outs[0].reshape(0, std::vector<int>{1, 8400, 84});
} }
// assert if last dim is 85 or 84
CV_CheckEQ(outs[0].dims, 3, "Invalid output shape. The shape should be [1, #anchors, 85 or 84]");
CV_CheckEQ((outs[0].size[2] == nc + 5 || outs[0].size[2] == 80 + 4), true, "Invalid output shape: ");
for (auto preds : outs){ for (auto preds : outs){
preds = preds.reshape(1, preds.size[1]); // [1, 8400, 85] -> [8400, 85] preds = preds.reshape(1, preds.size[1]); // [1, 8400, 85] -> [8400, 85]
for (int i = 0; i < preds.rows; ++i) for (int i = 0; i < preds.rows; ++i)
{ {
// filter out non object // filter out non object
float obj_conf = (test_name == "yolov8" || test_name == "yolonas") ? 1.0f : preds.at<float>(i, 4) ; float obj_conf = (model_name == "yolov8" || model_name == "yolonas" ||
model_name == "yolov9" || model_name == "yolov10") ? 1.0f : preds.at<float>(i, 4) ;
if (obj_conf < conf_threshold) if (obj_conf < conf_threshold)
continue; continue;
Mat scores = preds.row(i).colRange((test_name == "yolov8" || test_name == "yolonas") ? 4 : 5, preds.cols); Mat scores = preds.row(i).colRange((model_name == "yolov8" || model_name == "yolonas" || model_name == "yolov9" || model_name == "yolov10") ? 4 : 5, preds.cols);
double conf; double conf;
Point maxLoc; Point maxLoc;
minMaxLoc(scores, 0, &conf, 0, &maxLoc); minMaxLoc(scores, 0, &conf, 0, &maxLoc);
conf = (test_name == "yolov8" || test_name == "yolonas") ? conf : conf * obj_conf; conf = (model_name == "yolov8" || model_name == "yolonas" || model_name == "yolov9" || model_name == "yolov10") ? conf : conf * obj_conf;
if (conf < conf_threshold) if (conf < conf_threshold)
continue; continue;
@ -2722,15 +2731,14 @@ void yoloPostProcessing(
double w = det[2]; double w = det[2];
double h = det[3]; double h = det[3];
// std::cout << "cx: " << cx << " cy: " << cy << " w: " << w << " h: " << h << " conf: " << conf << " idx: " << maxLoc.x << std::endl;
// [x1, y1, x2, y2] // [x1, y1, x2, y2]
if (test_name == "yolonas"){ if (model_name == "yolonas" || model_name == "yolov10"){
boxes.push_back(Rect2d(cx, cy, w, h)); boxes.push_back(Rect2d(cx, cy, w, h));
} else { } else {
boxes.push_back(Rect2d(cx - 0.5 * w, cy - 0.5 * h, boxes.push_back(Rect2d(cx - 0.5 * w, cy - 0.5 * h,
cx + 0.5 * w, cy + 0.5 * h)); cx + 0.5 * w, cy + 0.5 * h));
} }
classIds.push_back(maxLoc.x); classIds.push_back(maxLoc.x);
confidences.push_back(conf); confidences.push_back(conf);
} }
} }
@ -2747,7 +2755,75 @@ void yoloPostProcessing(
} }
} }
TEST_P(Test_ONNX_nets, YOLOv10)
{
std::string weightPath = _tf("models/yolov10s.onnx", false);
Size targetSize{640, 480};
float conf_threshold = 0.50;
float iou_threshold = 0.50;
std::vector<int> refClassIds{1, 16, 7};
std::vector<float> refScores{0.9510f, 0.9454f, 0.8404f};
std::vector<Rect2d> refBoxes{
Rect2d(105.5014, 112.8838, 472.9274, 350.0603),
Rect2d(109.8231, 185.7994, 258.5916, 452.9302),
Rect2d(388.5018, 62.1034, 576.6399, 143.3986)
};
Image2BlobParams imgParams(
Scalar::all(1 / 255.0),
targetSize,
Scalar::all(0),
true,
CV_32F,
DNN_LAYOUT_NCHW,
DNN_PMODE_LETTERBOX,
Scalar::all(114)
);
testYOLO(
weightPath, refClassIds, refScores, refBoxes,
imgParams, conf_threshold, iou_threshold,
1.0e-4, 1.0e-4, "yolov10");
}
TEST_P(Test_ONNX_nets, YOLOv9)
{
std::string weightPath = _tf("models/yolov9t.onnx", false);
Size targetSize{640, 480};
float conf_threshold = 0.50;
float iou_threshold = 0.50;
std::vector<int> refClassIds{1, 16, 2}; // wrong class mapping for yolov9
std::vector<float> refScores{0.959274f, 0.901125, 0.559396f};
std::vector<Rect2d> refBoxes{
Rect2d(106.255, 107.927, 472.497, 350.309),
Rect2d(108.633, 185.256, 259.287, 450.672),
Rect2d(390.701, 62.1454, 576.928, 141.795)
};
Image2BlobParams imgParams(
Scalar::all(1 / 255.0),
targetSize,
Scalar::all(0),
true,
CV_32F,
DNN_LAYOUT_NCHW,
DNN_PMODE_LETTERBOX,
Scalar::all(114)
);
testYOLO(
weightPath, refClassIds, refScores, refBoxes,
imgParams, conf_threshold, iou_threshold,
1.0e-4, 1.0e-4, "yolov9");
}
TEST_P(Test_ONNX_nets, YOLOX) TEST_P(Test_ONNX_nets, YOLOX)
{ {
applyTestTag(CV_TEST_TAG_DEBUG_VERYLONG); applyTestTag(CV_TEST_TAG_DEBUG_VERYLONG);

@ -27,7 +27,8 @@ void yoloPostProcessing(
std::vector<Rect2d>& keep_boxes, std::vector<Rect2d>& keep_boxes,
float conf_threshold, float conf_threshold,
float iou_threshold, float iou_threshold,
const std::string& test_name const std::string& model_name,
const int nc
); );
std::vector<std::string> classes; std::vector<std::string> classes;
@ -40,6 +41,7 @@ std::string keys =
"{ yolo | yolox | yolo model version. }" "{ yolo | yolox | yolo model version. }"
"{ input i | | Path to input image or video file. Skip this argument to capture frames from a camera. }" "{ input i | | Path to input image or video file. Skip this argument to capture frames from a camera. }"
"{ classes | | Optional path to a text file with names of classes to label detected objects. }" "{ classes | | Optional path to a text file with names of classes to label detected objects. }"
"{ nc | 80 | Number of classes. Default is 80 (coming from COCO dataset). }"
"{ thr | .5 | Confidence threshold. }" "{ thr | .5 | Confidence threshold. }"
"{ nms | .4 | Non-maximum suppression threshold. }" "{ nms | .4 | Non-maximum suppression threshold. }"
"{ mean | 0.0 | Normalization constant. }" "{ mean | 0.0 | Normalization constant. }"
@ -107,19 +109,21 @@ void yoloPostProcessing(
std::vector<Rect2d>& keep_boxes, std::vector<Rect2d>& keep_boxes,
float conf_threshold, float conf_threshold,
float iou_threshold, float iou_threshold,
const std::string& test_name) const std::string& model_name,
const int nc=80)
{ {
// Retrieve // Retrieve
std::vector<int> classIds; std::vector<int> classIds;
std::vector<float> confidences; std::vector<float> confidences;
std::vector<Rect2d> boxes; std::vector<Rect2d> boxes;
if (test_name == "yolov8") if (model_name == "yolov8" || model_name == "yolov10" ||
model_name == "yolov9")
{ {
cv::transposeND(outs[0], {0, 2, 1}, outs[0]); cv::transposeND(outs[0], {0, 2, 1}, outs[0]);
} }
if (test_name == "yolonas") if (model_name == "yolonas")
{ {
// outs contains 2 elemets of shape [1, 8400, 80] and [1, 8400, 4]. Concat them to get [1, 8400, 84] // outs contains 2 elemets of shape [1, 8400, 80] and [1, 8400, 4]. Concat them to get [1, 8400, 84]
Mat concat_out; Mat concat_out;
@ -131,25 +135,30 @@ void yoloPostProcessing(
// remove the second element // remove the second element
outs.pop_back(); outs.pop_back();
// unsqueeze the first dimension // unsqueeze the first dimension
outs[0] = outs[0].reshape(0, std::vector<int>{1, 8400, 84}); outs[0] = outs[0].reshape(0, std::vector<int>{1, 8400, nc + 4});
} }
// assert if last dim is 85 or 84
CV_CheckEQ(outs[0].dims, 3, "Invalid output shape. The shape should be [1, #anchors, 85 or 84]");
CV_CheckEQ((outs[0].size[2] == nc + 5 || outs[0].size[2] == 80 + 4), true, "Invalid output shape: ");
for (auto preds : outs) for (auto preds : outs)
{ {
preds = preds.reshape(1, preds.size[1]); // [1, 8400, 85] -> [8400, 85] preds = preds.reshape(1, preds.size[1]); // [1, 8400, 85] -> [8400, 85]
for (int i = 0; i < preds.rows; ++i) for (int i = 0; i < preds.rows; ++i)
{ {
// filter out non object // filter out non object
float obj_conf = (test_name == "yolov8" || test_name == "yolonas") ? 1.0f : preds.at<float>(i, 4) ; float obj_conf = (model_name == "yolov8" || model_name == "yolonas" ||
model_name == "yolov9" || model_name == "yolov10") ? 1.0f : preds.at<float>(i, 4) ;
if (obj_conf < conf_threshold) if (obj_conf < conf_threshold)
continue; continue;
Mat scores = preds.row(i).colRange((test_name == "yolov8" || test_name == "yolonas") ? 4 : 5, preds.cols); Mat scores = preds.row(i).colRange((model_name == "yolov8" || model_name == "yolonas" || model_name == "yolov9" || model_name == "yolov10") ? 4 : 5, preds.cols);
double conf; double conf;
Point maxLoc; Point maxLoc;
minMaxLoc(scores, 0, &conf, 0, &maxLoc); minMaxLoc(scores, 0, &conf, 0, &maxLoc);
conf = (test_name == "yolov8" || test_name == "yolonas") ? conf : conf * obj_conf; conf = (model_name == "yolov8" || model_name == "yolonas" || model_name == "yolov9" || model_name == "yolov10") ? conf : conf * obj_conf;
if (conf < conf_threshold) if (conf < conf_threshold)
continue; continue;
@ -161,7 +170,7 @@ void yoloPostProcessing(
double h = det[3]; double h = det[3];
// [x1, y1, x2, y2] // [x1, y1, x2, y2]
if (test_name == "yolonas"){ if (model_name == "yolonas" || model_name == "yolov10"){
boxes.push_back(Rect2d(cx, cy, w, h)); boxes.push_back(Rect2d(cx, cy, w, h));
} else { } else {
boxes.push_back(Rect2d(cx - 0.5 * w, cy - 0.5 * h, boxes.push_back(Rect2d(cx - 0.5 * w, cy - 0.5 * h,
@ -203,6 +212,7 @@ int main(int argc, char** argv)
// if model is default, use findFile to get the full path otherwise use the given path // if model is default, use findFile to get the full path otherwise use the given path
std::string weightPath = findFile(parser.get<String>("model")); std::string weightPath = findFile(parser.get<String>("model"));
std::string yolo_model = parser.get<String>("yolo"); std::string yolo_model = parser.get<String>("yolo");
int nc = parser.get<int>("nc");
float confThreshold = parser.get<float>("thr"); float confThreshold = parser.get<float>("thr");
float nmsThreshold = parser.get<float>("nms"); float nmsThreshold = parser.get<float>("nms");
@ -219,6 +229,7 @@ int main(int argc, char** argv)
// check if yolo model is valid // check if yolo model is valid
if (yolo_model != "yolov5" && yolo_model != "yolov6" if (yolo_model != "yolov5" && yolo_model != "yolov6"
&& yolo_model != "yolov7" && yolo_model != "yolov8" && yolo_model != "yolov7" && yolo_model != "yolov8"
&& yolo_model != "yolov10" && yolo_model !="yolov9"
&& yolo_model != "yolox" && yolo_model != "yolonas") && yolo_model != "yolox" && yolo_model != "yolonas")
CV_Error(Error::StsError, "Invalid yolo model: " + yolo_model); CV_Error(Error::StsError, "Invalid yolo model: " + yolo_model);
@ -331,7 +342,8 @@ int main(int argc, char** argv)
yoloPostProcessing( yoloPostProcessing(
outs, keep_classIds, keep_confidences, keep_boxes, outs, keep_classIds, keep_confidences, keep_boxes,
confThreshold, nmsThreshold, confThreshold, nmsThreshold,
yolo_model); yolo_model,
nc);
//![postprocess] //![postprocess]
// covert Rect2d to Rect // covert Rect2d to Rect

Loading…
Cancel
Save