|
|
|
|
@ -20,6 +20,12 @@ except ImportError as e:
|
|
|
|
|
print(f"YOLOv5导入失败: {e}")
|
|
|
|
|
YOLOV5_AVAILABLE = False
|
|
|
|
|
|
|
|
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
|
|
|
PROJECT_ROOT = SCRIPT_DIR.parents[2] # /Src/vision_models/services -> root
|
|
|
|
|
MODELS_DIR = SCRIPT_DIR.parent / "models"
|
|
|
|
|
YOLOROOT_DIR = SCRIPT_DIR.parent / "yolov5"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class YOLOv5Detector:
|
|
|
|
|
"""YOLOv5目标检测器 - 简化版本"""
|
|
|
|
|
|
|
|
|
|
@ -33,17 +39,32 @@ class YOLOv5Detector:
|
|
|
|
|
|
|
|
|
|
# 加载模型
|
|
|
|
|
try:
|
|
|
|
|
if weights_path and os.path.exists(weights_path):
|
|
|
|
|
self.model = YOLO(weights_path)
|
|
|
|
|
else:
|
|
|
|
|
# 优先使用现有的yolov5su.pt模型
|
|
|
|
|
if os.path.exists('yolov5su.pt'):
|
|
|
|
|
self.model = YOLO('yolov5su.pt')
|
|
|
|
|
elif os.path.exists('../models/yolov5su.pt'):
|
|
|
|
|
self.model = YOLO('../models/yolov5su.pt')
|
|
|
|
|
else:
|
|
|
|
|
# 使用默认的YOLOv5s模型
|
|
|
|
|
self.model = YOLO('yolov5s.pt')
|
|
|
|
|
candidate_paths = []
|
|
|
|
|
if weights_path:
|
|
|
|
|
candidate_paths.append(Path(weights_path))
|
|
|
|
|
|
|
|
|
|
# 优先选择项目内模型
|
|
|
|
|
candidate_paths.extend([
|
|
|
|
|
MODELS_DIR / "yolov5su.pt",
|
|
|
|
|
SCRIPT_DIR / "yolov5su.pt",
|
|
|
|
|
Path.cwd() / "yolov5su.pt",
|
|
|
|
|
MODELS_DIR / "yolov5s.pt",
|
|
|
|
|
SCRIPT_DIR / "yolov5s.pt",
|
|
|
|
|
Path.cwd() / "yolov5s.pt",
|
|
|
|
|
])
|
|
|
|
|
|
|
|
|
|
model_path = None
|
|
|
|
|
for path in candidate_paths:
|
|
|
|
|
if path.exists():
|
|
|
|
|
model_path = path
|
|
|
|
|
break
|
|
|
|
|
|
|
|
|
|
if model_path is None:
|
|
|
|
|
raise FileNotFoundError(
|
|
|
|
|
"未找到YOLO模型权重,请将 yolov5su.pt 或 yolov5s.pt 放在 Src/vision_models/models 目录。"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
self.model = YOLO(str(model_path))
|
|
|
|
|
print(f"✅ 模型加载成功: {self.model.model_name if hasattr(self.model, 'model_name') else 'YOLOv5'}")
|
|
|
|
|
except Exception as e:
|
|
|
|
|
raise RuntimeError(f"模型加载失败: {e}")
|