环境:
Mac M1 (MacOS Sonoma 14.3.1)
Python 3.11+PyTorch 2.1.2
一、准备工作
使用YOLO一般都会接触ultralytics这个框架,今天来试试用该框架进行YOLO V9模型的推理。
YOLOv9目前提供了四种模型下载:yolov9-c.pt、yolov9-e.pt、gelan-c.pt、gelan-e.pt
wget -P /Users/zhujiahui/Local/model/yolov9 -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-c.pt
wget -P /Users/zhujiahui/Local/model/yolov9 -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e.pt
wget -P /Users/zhujiahui/Local/model/yolov9 -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c.pt
wget -P /Users/zhujiahui/Local/model/yolov9 -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-e.pt
将下载好的模型放到指定的位置(实际发现这一步不用做,白费了)。
下载示例图片(也可手动下载放置),放到指定位置下:
wget -P /Users/zhujiahui/Local/dataset -q https://media.roboflow.com/notebooks/examples/dog.jpeg
安装Python依赖
pip install opencv-python
pip install ultralytics
二、推理
编写以下代码:
from ultralytics import YOLO
def yolov9_inference():
yolo_model = YOLO("/Users/zhujiahui/Local/model/yolov9/yolov9-e.pt")
yolo_model.info()
if __name__ == '__main__':
yolov9_inference()
直接运行,发现报错:
TypeError: ERROR ❌️ /Users/zhujiahui/Local/model/yolov9/yolov9-e.pt appears to be an Ultralytics YOLOv5 model originally trained with https://github.com/ultralytics/yolov5.
This model is NOT forwards compatible with YOLOv8 at https://github.com/ultralytics/ultralytics.
Recommend fixes are to train a new model using the latest 'ultralytics' package or to run a command with an official YOLOv8 model, i.e. 'yolo predict model=yolov8n.pt'
意思是说从https://github.com/WongKinYiu/yolov9下载的模型yolov9-e.pt与本训练推理框架(ultralytics)不match,必须使用经过ultralytics训练的模型。
于是改动如下,选用ultralytics提供的YOLO V9模型:
from ultralytics import YOLO
def yolov9_inference():
yolo_model = YOLO("yolov9e.pt")
yolo_model.info()
if __name__ == '__main__':
yolov9_inference()
首次运行会下载模型到当前代码所在的文件夹下
对示例图片进行检测:
from ultralytics import YOLO
def yolov9_inference():
yolo_model = YOLO("yolov9e.pt")
yolo_model.info()
results = yolo_model.predict(source='/Users/zhujiahui/Local/dataset/dog.jpeg', save=True, save_txt=True)
print(results)
if __name__ == '__main__':
yolov9_inference()
结果如下:
具体检测后的结果图片在runs/detect/predict/dog.jpeg下,效果:
原图 | 检测结果 |
输入的图片也可以通过OpenCV读入:
def yolov9_inference2():
yolo_model = YOLO("yolov9e.pt")
yolo_model.info()
input_image = cv2.imread("/Users/zhujiahui/Local/dataset/dog.jpeg")
results = yolo_model.predict(source=input_image, save=True, save_txt=True)
print(results)