62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
|
import os
|
|||
|
|
import numpy as np
|
|||
|
|
import traceback
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
from typing import Optional
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_yolo_model(model_path: str) -> Optional[YOLO]:
|
|||
|
|
"""
|
|||
|
|
加载YOLO模型（支持v5/v8），并校验模型有效性
|
|||
|
|
:param model_path: 模型文件的绝对路径
|
|||
|
|
:return: 加载成功返回YOLO模型实例，失败返回None
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
# 加载前的基础信息检查
|
|||
|
|
print(f"\n[模型工具] 开始加载模型：{model_path}")
|
|||
|
|
print(f"[模型工具] 文件是否存在：{os.path.exists(model_path)}")
|
|||
|
|
if os.path.exists(model_path):
|
|||
|
|
print(f"[模型工具] 文件大小：{os.path.getsize(model_path) / 1024 / 1024:.2f} MB")
|
|||
|
|
|
|||
|
|
# 强制重新加载模型，避免缓存问题
|
|||
|
|
model = YOLO(model_path)
|
|||
|
|
|
|||
|
|
# 兼容性校验：使用numpy空数组测试模型
|
|||
|
|
dummy_image = np.zeros((640, 640, 3), dtype=np.uint8)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 优先使用新版本参数
|
|||
|
|
model.predict(
|
|||
|
|
source=dummy_image,
|
|||
|
|
imgsz=640,
|
|||
|
|
conf=0.25,
|
|||
|
|
verbose=False,
|
|||
|
|
stream=False
|
|||
|
|
)
|
|||
|
|
except Exception as pred_e:
|
|||
|
|
print(f"[模型工具] 预测校验兼容处理：{str(pred_e)}")
|
|||
|
|
# 兼容旧版本YOLO参数
|
|||
|
|
model.predict(
|
|||
|
|
img=dummy_image,
|
|||
|
|
imgsz=640,
|
|||
|
|
conf=0.25,
|
|||
|
|
verbose=False
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 验证模型基本属性
|
|||
|
|
if not hasattr(model, 'names'):
|
|||
|
|
print("[模型工具] 警告：模型缺少类别名称属性")
|
|||
|
|
else:
|
|||
|
|
print(f"[模型工具] 模型包含类别：{list(model.names.values())[:5]}...") # 显示前5个类别
|
|||
|
|
|
|||
|
|
print(f"[模型工具] 模型加载成功！")
|
|||
|
|
return model
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
# 详细错误信息输出
|
|||
|
|
print(f"\n[模型工具] 加载模型失败！路径：{model_path}")
|
|||
|
|
print(f"[模型工具] 异常类型：{type(e).__name__}")
|
|||
|
|
print(f"[模型工具] 异常详情：{str(e)}")
|
|||
|
|
print(f"[模型工具] 堆栈跟踪：\n{traceback.format_exc()}")
|
|||
|
|
return None
|