brain-facens/cellphone-detection-yolov6

Cellphone detection using YOLOv6

Python scripts performing object detection using the YOLOv6 model in ONNX.

Important

• About the YOLOv6 name: https://github.com/meituan/YOLOv6/blob/main/docs/About_naming_yolov6.md
• The input images are directly resized to match the input size of the model. I skipped adding the pad to the input image, it might affect the accuracy of the model if the input image has a different aspect ratio compared to the input size of the model. Always try to get an input size with a ratio close to the input images you will use.
• The original YOLOv6 repository is still in development, so things might not work in the future.

Requirements

• Check the requirements.txt file.
• For ONNX, if you have a NVIDIA GPU, then install the onnxruntime-gpu, otherwise use the onnxruntime library.
• Additionally, pafy and youtube-dl are required for youtube video inference.

Installation

git clone https://github.com/brain-facens/YOLOv6
cd ONNX-YOLOv6-Object-Detection
pip install -r requirements.txt

ONNX Runtime

For Nvidia GPU computers: pip install onnxruntime-gpu

Otherwise: pip install onnxruntime

For youtube video inference

pip install youtube_dl
pip install git+https://github.com/zizo-pro/pafy@b8976f22c19e4ab5515cacbfae0a3970370c102b

ONNX model

The original model was converted to ONNX using the following Colab notebook from the original repository, run the notebook and save the download model into the models folder:

• Convert YOLOv6 ONNX for Inference
• The License of the models is GPL-3.0 license: License

Pytorch model

The original Pytorch model can be found in this repository: YOLOv6 Repository

• Runing:

./run.sh

• Comparison with YOLOv5:

python video_object_detection.py

Convert YOLOv5 model to ONNX:

References:

• YOLOv6 model: https://github.com/meituan/YOLOv6
• YOLOv5 model: https://github.com/ultralytics/yolov5