環境:
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 YOLOdef 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 YOLOdef yolov9_inference():yolo_model = YOLO("yolov9e.pt")yolo_model.info()if __name__ == '__main__':yolov9_inference()
首次運行會下載模型到當前代碼所在的文件夾下
對示例圖片進行檢測:
from ultralytics import YOLOdef 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)