目前可以成功动态更换模型运行的
This commit is contained in:
		
							
								
								
									
										47
									
								
								core/yolo.py
									
									
									
									
									
								
							
							
						
						
									
										47
									
								
								core/yolo.py
									
									
									
									
									
								
							| @ -1,37 +1,43 @@ | ||||
| import os | ||||
|  | ||||
| import numpy as np | ||||
| from ultralytics import YOLO | ||||
| from service.model_service import get_current_yolo_model  # 从模型管理模块获取模型 | ||||
|  | ||||
| # 全局变量 | ||||
| # 全局模型变量 | ||||
| _yolo_model = None | ||||
|  | ||||
|  | ||||
| model_path = os.path.join(os.path.dirname(__file__), "models", "best.pt") | ||||
|  | ||||
|  | ||||
| def load_model(): | ||||
|     """加载YOLO目标检测模型""" | ||||
| def load_model(model_path=None): | ||||
|     """加载YOLO模型(优先使用模型管理模块的默认模型)""" | ||||
|     global _yolo_model | ||||
|  | ||||
|     if model_path is None: | ||||
|         _yolo_model = get_current_yolo_model() | ||||
|         return _yolo_model is not None | ||||
|  | ||||
|     try: | ||||
|         _yolo_model = YOLO(model_path) | ||||
|         return True | ||||
|     except Exception as e: | ||||
|         print(f"YOLO model load failed: {e}") | ||||
|         print(f"YOLO模型加载失败(指定路径):{str(e)}") | ||||
|         return False | ||||
|  | ||||
|     return True if _yolo_model else False | ||||
|  | ||||
|  | ||||
| def detect(frame, conf_threshold=0.2): | ||||
|     """YOLO目标检测、返回(是否识别到, 结果字符串)""" | ||||
|     """执行目标检测,返回(是否成功, 结果字符串)""" | ||||
|     global _yolo_model | ||||
|  | ||||
|     if not _yolo_model or frame is None: | ||||
|         return (False, "未初始化或无效帧") | ||||
|     # 确保模型已加载 | ||||
|     if not _yolo_model: | ||||
|         if not load_model(): | ||||
|             return (False, "模型未初始化") | ||||
|  | ||||
|     if frame is None: | ||||
|         return (False, "无效输入帧") | ||||
|  | ||||
|     try: | ||||
|         results = _yolo_model(frame, conf=conf_threshold) | ||||
|         # 检查是否有检测结果 | ||||
|         # 执行检测(frame应为numpy数组) | ||||
|         results = _yolo_model(frame, conf=conf_threshold, verbose=False) | ||||
|         has_results = len(results[0].boxes) > 0 if results else False | ||||
|  | ||||
|         if not has_results: | ||||
| @ -42,13 +48,12 @@ def detect(frame, conf_threshold=0.2): | ||||
|         for box in results[0].boxes: | ||||
|             cls = int(box.cls[0]) | ||||
|             conf = float(box.conf[0]) | ||||
|             bbox = [float(x) for x in box.xyxy[0]] | ||||
|             bbox = [round(x, 2) for x in box.xyxy[0].tolist()]  # 保留两位小数 | ||||
|             class_name = _yolo_model.names[cls] if hasattr(_yolo_model, 'names') else f"类别{cls}" | ||||
|             result_parts.append(f"{class_name} (置信度: {conf:.2f}, 边界框: {bbox})") | ||||
|             result_parts.append(f"{class_name}(置信度:{conf:.2f},位置:{bbox})") | ||||
|  | ||||
|         result_str = "; ".join(result_parts) | ||||
|         return (has_results, result_str) | ||||
|         return (True, "; ".join(result_parts)) | ||||
|  | ||||
|     except Exception as e: | ||||
|         print(f"YOLO detect error: {e}") | ||||
|         return (False, f"检测错误: {str(e)}") | ||||
|         print(f"检测过程出错:{str(e)}") | ||||
|         return (False, f"检测错误:{str(e)}") | ||||
|  | ||||
		Reference in New Issue
	
	Block a user