abhishek-gola's picture
Add_Yolov4_onnx_model (#5)
c902f68
|
Raw
History Blame
4.15 kB
# YOLOv4, YOLOv4-tiny and YOLOv4x-mish ONNX Conversion
## Prerequisites
### 1. Model files (cfg + weights)
The Darknet `.cfg` files are already provided in [opencv_extra/testdata/dnn](https://github.com/opencv/opencv_extra/tree/master/testdata/dnn) (`yolov4.cfg`, `yolov4-tiny-2020-12.cfg`, `yolov4x-mish.cfg`).
Download the matching `.weights` files using the OpenCV test data download script:
```bash
git clone https://github.com/opencv/opencv_extra.git
cd opencv_extra/testdata/dnn
python download_models.py YOLOv4 YOLOv4-tiny-2020-12 YOLOv4x-mish
```
### 2. Python environment for `pytorch-YOLOv4`
The conversion uses [`pytorch-YOLOv4`](https://github.com/Tianxiaomo/pytorch-YOLOv4). Create a Python environment with the required dependencies:
Supported Python versions: **3.7 – 3.10**.
```bash
conda create -n <env_name> python=<3.7-3.10> -y
conda activate <env_name>
pip install "torch<2.4" "torchvision<0.19" "numpy<2" onnx onnxruntime onnxscript
```
---
## Conversion of YOLOv4 to ONNX
```bash
git clone https://github.com/Tianxiaomo/pytorch-YOLOv4.git
cd pytorch-YOLOv4
# Convert (dynamic batch, batch_size=0)
python -c "from tool.darknet2onnx import transform_to_onnx; transform_to_onnx('yolov4.cfg', 'yolov4.weights', 0)"
```
---
## Conversion of YOLOv4-tiny to ONNX
```bash
git clone https://github.com/Tianxiaomo/pytorch-YOLOv4.git
cd pytorch-YOLOv4
# Convert YOLOv4-tiny (dynamic batch, batch_size=0)
python -c "from tool.darknet2onnx import transform_to_onnx; transform_to_onnx('yolov4-tiny-2020-12.cfg', 'yolov4-tiny.weights', 0)"
```
---
## Conversion of YOLOv4x-mish to ONNX
### Why it differs from YOLOv4
The `yolov4x-mish.cfg` uses `new_coords=1` — a Scaled-YOLOv4 optimization. The network is trained to output values directly in the `[0, 1]` range (no sigmoid needed) and uses a squared formula `(t_w * 2)² * anchor` for width/height instead of `exp(t_w) * anchor`. This is more numerically stable for large models.
### The Issue
The `pytorch-YOLOv4` converter was written for the original YOLOv4 (`new_coords=0`) and always applies `sigmoid` + `exp`. With `new_coords=1` weights, `sigmoid` gets applied on top of already-activated values, squishing confidences from ~0.93 down to ~0.36. As a result, the model produces garbage detections.
### The Fix (Modified scripts provided)
To resolve this, modified versions of **`darknet2pytorch.py`** and **`yolo_layer.py`** are provided in this repository. These scripts include the following patches:
* **`darknet2pytorch.py`**: Properly reads the `new_coords` flag from the `.cfg`.
* **`yolo_layer.py`**: Skips the redundant sigmoid activation for `xy/obj/cls` and implements the squared `wh` formula when `new_coords=1` is detected.
### Conversion Steps
```bash
git clone https://github.com/Tianxiaomo/pytorch-YOLOv4.git
cd pytorch-YOLOv4
# [!] Replace tool/darknet2pytorch.py and tool/yolo_layer.py with the patched
# versions from this repository before running the conversion.
# Convert YOLOv4x-mish (dynamic batch, batch_size=0)
python -c "from tool.darknet2onnx import transform_to_onnx; transform_to_onnx('yolov4x-mish.cfg', 'yolov4x-mish.weights', 0)"
```
---
## Usage
A demo script is provided to run inference using OpenCV DNN:
```bash
# YOLOv4 (input size: 608x608)
python demo.py --model yolov4.onnx --image example_outputs/input.jpg --output example_outputs/yolov4_output.jpg
# YOLOv4-tiny (input size: 416x416)
python demo.py --model yolov4-tiny.onnx --image example_outputs/input.jpg --output example_outputs/yolov4-tiny_output.jpg
# YOLOv4x-mish (input size: 640x640)
python demo.py --model yolov4x-mish.onnx --image example_outputs/input.jpg --output example_outputs/yolov4x-mish_output.jpg
```
The demo prints the detected COCO classes, confidence scores, and bounding boxes, and saves an annotated output image.
---
## License
See [License.txt](./License.txt) — This conversion tool is based on [pytorch-YOLOv4](https://github.com/Tianxiaomo/pytorch-YOLOv4) (Apache-2.0). Original YOLOv4 model weights and configuration are released by Alexey Bochkovskiy ([AlexeyAB/darknet](https://github.com/AlexeyAB/darknet)).