YOLO模型物體檢測
下面是一個簡單的python程序,他的功能是使用yolo11n模型對coco_bike.jpg照片進行檢測,并顯示檢測結果
代碼步驟如下:
- coco_bike.jpg照片
- 加載yolo模型
- 使用模型進行detect推理
- 顯示推理結果
下面是完整的代碼
from pathlib import Pathimport urllib.request
import tkinter as tk
from PIL import Image, ImageTkfrom ultralytics import YOLO# ----------------------------- #
# 全局配置和路徑定義
# ----------------------------- #MODEL_VARIANTS = ["yolo11n", "yolo11s", "yolo11m", "yolo11l", "yolo11x"]
MODEL_NAME = MODEL_VARIANTS[0] # 默認使用最輕量的 yolo11n 模型
PT_MODEL_PATH = f"{MODEL_NAME}.pt"
IR_MODEL_DIR = Path(f"{MODEL_NAME}_openvino_model")
IR_MODEL_PATH = IR_MODEL_DIR / f"{MODEL_NAME}.xml"
INT8_MODEL_PATH = Path(f"{MODEL_NAME}_openvino_int8_model/{MODEL_NAME}.xml")IMAGE_PATH = Path("./coco_bike.jpg")
OUT_DIR = Path("./")# ----------------------------- #
# 工具函數模塊
# ----------------------------- #def download_file_if_needed(url: str, filename: str, dest_dir: Path) -> Path:"""下載文件(若文件已存在則跳過)"""dest_dir.mkdir(parents=True, exist_ok=True)file_path = dest_dir / filenameif not file_path.exists():print(f"Downloading: {filename}")urllib.request.urlretrieve(url, file_path)else:print(f"File already exists: {file_path}")return file_pathdef prepare_test_image():"""確保測試圖片存在,如無則從官方地址下載"""if not IMAGE_PATH.exists():download_file_if_needed("https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg",IMAGE_PATH.name, IMAGE_PATH.parent)def predict_and_show_image(det_model: YOLO):"""使用模型對圖像進行目標檢測,并通過 Tkinter GUI 顯示檢測結果"""results = det_model(IMAGE_PATH)result_img = Image.fromarray(results[0].plot()[:, :, ::-1])root = tk.Tk()root.title("YOLOv11 Detection Result")tk_img = ImageTk.PhotoImage(result_img)label = tk.Label(root, image=tk_img)label.pack()root.mainloop()# ----------------------------- #
# 主執行流程
# ----------------------------- #def main():# 1. 準備測試圖像(如無則下載)prepare_test_image()# 2. 構造 Ultralytics YOLO 接口,用于推理/驗證det_model = YOLO(PT_MODEL_PATH).to("cpu")# 3. 使用openvino模型進行推理并顯示結果predict_and_show_image(det_model)if __name__ == "__main__":main()
下面是推理結果:
OpenVINO物體檢測
基于上面程序我們修改,將yolo模型轉換為openvino IR模型,然后使用IR模型進行物體檢測.完整代碼如下:
代碼步驟如下:
- coco_bike.jpg照片
- 導出和加載OpenVINO IR模型
- 用 Ultralytics 的 YOLO 接口重新封裝 之前導出的 OpenVINO 模型
- 使用模型進行detect推理
- 顯示推理結果
下面是完整的代碼
from pathlib import Pathimport urllib.request
import tkinter as tk
from PIL import Image, ImageTkfrom ultralytics import YOLO
import openvino as ov# ----------------------------- #
# 全局配置和路徑定義
# ----------------------------- #MODEL_VARIANTS = ["yolo11n", "yolo11s", "yolo11m", "yolo11l", "yolo11x"]
MODEL_NAME = MODEL_VARIANTS[0] # 默認使用最輕量的 yolo11n 模型
PT_MODEL_PATH = f"{MODEL_NAME}.pt"
IR_MODEL_DIR = Path(f"{MODEL_NAME}_openvino_model")
IR_MODEL_PATH = IR_MODEL_DIR / f"{MODEL_NAME}.xml"
INT8_MODEL_PATH = Path(f"{MODEL_NAME}_openvino_int8_model/{MODEL_NAME}.xml")IMAGE_PATH = Path("./coco_bike.jpg")
OUT_DIR = Path("./")# ----------------------------- #
# 工具函數模塊
# ----------------------------- #def download_file_if_needed(url: str, filename: str, dest_dir: Path) -> Path:"""下載文件(若文件已存在則跳過)"""dest_dir.mkdir(parents=True, exist_ok=True)file_path = dest_dir / filenameif not file_path.exists():print(f"Downloading: {filename}")urllib.request.urlretrieve(url, file_path)else:print(f"File already exists: {file_path}")return file_pathdef prepare_test_image():"""確保測試圖片存在,如無則從官方地址下載"""if not IMAGE_PATH.exists():download_file_if_needed("https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg",IMAGE_PATH.name, IMAGE_PATH.parent)def load_or_export_openvino_model() -> ov.CompiledModel:"""加載或導出 YOLOv11 OpenVINO IR 模型,并編譯為 CPU 運行時模型"""model = YOLO(PT_MODEL_PATH).to("cpu")if not IR_MODEL_PATH.exists():model.export(format="openvino", dynamic=True, half=True)core = ov.Core()ir_model = core.read_model(IR_MODEL_PATH)return core.compile_model(ir_model, "CPU")def build_ultralytics_model() -> YOLO:"""創建 Ultralytics 的 YOLO 模型接口,用于調用預測器"""return YOLO(IR_MODEL_DIR, task="detect")def predict_and_show_image(det_model: YOLO, compiled_model: ov.CompiledModel):"""使用模型對圖像進行目標檢測,并通過 Tkinter GUI 顯示檢測結果"""if det_model.predictor is None:config = {"conf": 0.25, "batch": 1, "save": False, "mode": "predict"}args = {**det_model.overrides, **config}det_model.predictor = det_model._smart_load("predictor")(overrides=args, _callbacks=det_model.callbacks)det_model.predictor.setup_model(model=det_model.model)det_model.predictor.model.ov_compiled_model = compiled_modelresults = det_model(IMAGE_PATH)result_img = Image.fromarray(results[0].plot()[:, :, ::-1])root = tk.Tk()root.title("YOLOv11 (OpenVINO) Detection Result")tk_img = ImageTk.PhotoImage(result_img)label = tk.Label(root, image=tk_img)label.pack()root.mainloop()# ----------------------------- #
# 主執行流程
# ----------------------------- #def main():# 1. 準備測試圖像(如無則下載)prepare_test_image()# 2. 加載或導出 OpenVINO IR 模型,并編譯運行(用于量化或預測)compiled_fp_model = load_or_export_openvino_model()# 3. 構造 Ultralytics YOLO 接口,用于推理/驗證det_model = build_ultralytics_model()# 4. 使用openvino模型進行推理并顯示結果predict_and_show_image(det_model, compiled_fp_model)if __name__ == "__main__":main()
下面是推理結果: