tucan9389 / ObjectDetection-CoreML

An example running Object Detection using Core ML (YOLOv8, YOLOv5, YOLOv3, MobileNetV2+SSDLite)
https://github.com/motlabs/awesome-ml-demos-with-ios
MIT License
295 stars 51 forks source link

Support YOLOv5 #6

Closed tucan9389 closed 2 years ago

tucan9389 commented 2 years ago

Convert pre-trained YOLOv5 pt to Core ML

Clone and install pre-requirements

Note that you should clone mshamash/yolov5 repo and checkout fix/coreml_export_nms_layer branch, instead of the original repo (ultralytics/yolov5) because the PR wasn't merged.

git clone https://github.com/mshamash/yolov5
cd yolov5
git checkout fix/coreml_export_nms_layer
pip install -qr requirements.txt  # install

Export to Core ML model

python export.py --weights yolov5n.pt --include coreml
python export.py --weights yolov5s.pt --include coreml
python export.py --weights yolov5m.pt --include coreml
python export.py --weights yolov5l.pt --include coreml
python export.py --weights yolov5x.pt --include coreml

python export.py --weights yolov5n6.pt --include coreml
python export.py --weights yolov5s6.pt --include coreml
python export.py --weights yolov5m6.pt --include coreml
python export.py --weights yolov5l6.pt --include coreml
python export.py --weights yolov5x6.pt --include coreml

Check the model input/output with https://netron.app

yolov5s.mlmodel example

image

Check the model with Xcode

image

Related Issue and PR

tucan9389 commented 2 years ago

Class names of YOLOv5's pre-trained with COCO dataset