✨ feat: python模型管理

python-inference-service/app/main.py

@@ -37,7 +37,7 @@ async def startup_event():
     model_manager = ModelManager()
     
     # Look for models.json configuration file
-    models_json_path = os.getenv("MODELS_JSON", os.path.join(os.path.dirname(__file__), "..", "models", "models.json"))
+    models_json_path = os.getenv("MODELS_JSON", os.path.join(os.path.dirname(__file__), "..", "models.json"))
     
     if os.path.exists(models_json_path):
         try:

python-inference-service/models.json

@@ -0,0 +1,14 @@
+[
+  {
+    "name": "smoke",
+    "path": "models/smoke_model.py",
+    "size": [640, 640],
+    "comment": "烟雾检测模型"
+  },
+  {
+    "name": "garbage",
+    "path": "models/garbage_model.py",
+    "size": [640, 640],
+    "comment": "垃圾检测模型"
+  }
+] 

python-inference-service/models/models.json

@@ -1,8 +0,0 @@
-[
-  {
-    "name": "yolov8_detector",
-    "path": "models/yolov8_model.py",
-    "size": [640, 640],
-    "comment": "YOLOv8检测模型，确保将训练好的best.pt文件放在models目录下"
-  }
-] 

python-inference-service/models/smoke_model.py

@@ -0,0 +1,207 @@
+import os
+import numpy as np
+import cv2
+from typing import List, Dict, Any
+import torch
+
+class Model:
+    """
+    垃圾识别模型 - 直接加载 PyTorch 模型文件
+    """
+    
+    def __init__(self):
+        """初始化模型"""
+        # 获取当前文件所在目录路径
+        model_dir = os.path.dirname(os.path.abspath(__file__))
+        # 模型文件路径
+        model_path = os.path.join(model_dir, "smoke.pt")
+        
+        print(f"正在加载垃圾识别模型: {model_path}")
+        
+        # 加载 PyTorch 模型
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"使用设备: {self.device}")
+        
+        # 使用 YOLOv5 或通用方式加载模型
+        try:
+            # 尝试使用 YOLOv5 加载
+            import sys
+            sys.path.append(os.path.dirname(model_dir))  # 添加父目录到路径
+            
+            try:
+                # 方法1: 如果安装了 YOLOv5
+                import yolov5
+                self.model = yolov5.load(model_path, device=self.device)
+                self.yolov5_api = True
+                print("使用 YOLOv5 包加载模型")
+            except (ImportError, ModuleNotFoundError):
+                # 方法2: 直接加载 YOLO 代码
+                from models.yolov5_utils import attempt_load
+                self.model = attempt_load(model_path, device=self.device)
+                self.yolov5_api = False
+                print("使用内置 YOLOv5 工具加载模型")
+                
+        except Exception as e:
+            # 方法3: 通用 PyTorch 加载
+            print(f"YOLOv5 加载失败: {e}")
+            print("使用通用 PyTorch 加载")
+            self.model = torch.load(model_path, map_location=self.device)
+            if isinstance(self.model, dict) and 'model' in self.model:
+                self.model = self.model['model']
+            self.yolov5_api = False
+            
+            # 如果是 ScriptModule，设置为评估模式
+            if isinstance(self.model, torch.jit.ScriptModule):
+                self.model.eval()
+            elif hasattr(self.model, 'eval'):
+                self.model.eval()
+                
+        # 加载类别名称
+        self.classes = []
+        classes_path = os.path.join(model_dir, "classes.txt")
+        if os.path.exists(classes_path):
+            with open(classes_path, 'r', encoding='utf-8') as f:
+                self.classes = [line.strip() for line in f.readlines() if line.strip()]
+            print(f"已加载 {len(self.classes)} 个类别")
+        else:
+            # 如果模型自带类别信息
+            if hasattr(self.model, 'names') and self.model.names:
+                self.classes = self.model.names
+                print(f"使用模型自带类别，共 {len(self.classes)} 个类别")
+            else:
+                print("未找到类别文件，将使用数字索引作为类别名")
+        
+        # 设置识别参数
+        self.conf_threshold = 0.25  # 置信度阈值
+        self.img_size = 640        # 默认输入图像大小
+        
+        print("垃圾识别模型加载完成")
+    
+    def preprocess(self, image: np.ndarray) -> np.ndarray:
+        """预处理图像"""
+        # 如果是使用 YOLOv5 API，不需要预处理
+        if hasattr(self, 'yolov5_api') and self.yolov5_api:
+            return image
+            
+        # 默认预处理：调整大小并归一化
+        img = cv2.resize(image, (self.img_size, self.img_size))
+        
+        # BGR 转 RGB
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        
+        # 归一化 [0, 255] -> [0, 1]
+        img = img / 255.0
+        
+        # HWC -> CHW (高度，宽度，通道 -> 通道，高度，宽度)
+        img = img.transpose(2, 0, 1)
+        
+        # 转为 torch tensor
+        img = torch.from_numpy(img).float()
+        
+        # 添加批次维度
+        img = img.unsqueeze(0)
+        
+        # 移至设备
+        img = img.to(self.device)
+        
+        return img
+    
+    def predict(self, image: np.ndarray) -> List[Dict[str, Any]]:
+        """模型推理"""
+        original_height, original_width = image.shape[:2]
+        
+        try:
+            # 如果使用 YOLOv5 API
+            if hasattr(self, 'yolov5_api') and self.yolov5_api:
+                # YOLOv5 API 直接处理图像
+                results = self.model(image)
+                
+                # 提取检测结果
+                predictions = results.pred[0]  # 第一批次的预测
+                
+                detections = []
+                for *xyxy, conf, cls_id in predictions.cpu().numpy():
+                    x1, y1, x2, y2 = xyxy
+                    
+                    # 转换为归一化坐标 (x, y, w, h)
+                    x = x1 / original_width
+                    y = y1 / original_height
+                    w = (x2 - x1) / original_width
+                    h = (y2 - y1) / original_height
+                    
+                    # 整数类别 ID
+                    cls_id = int(cls_id)
+                    
+                    # 获取类别名称
+                    class_name = f"cls{cls_id}"
+                    if 0 <= cls_id < len(self.classes):
+                        class_name = self.classes[cls_id]
+                    
+                    # 添加检测结果
+                    if conf >= self.conf_threshold:
+                        detections.append({
+                            'bbox': (x, y, w, h),
+                            'class_id': cls_id,
+                            'confidence': float(conf)
+                        })
+                        
+                return detections
+                
+            else:
+                # 通用 PyTorch 模型处理
+                # 预处理图像
+                img = self.preprocess(image)
+                
+                # 推理
+                with torch.no_grad():
+                    outputs = self.model(img)
+                
+                # 后处理结果（这里需要根据模型输出格式调整）
+                detections = []
+                
+                # 假设输出格式是 YOLO 风格：[batch_idx, x1, y1, x2, y2, conf, cls_id]
+                if isinstance(outputs, torch.Tensor) and outputs.dim() == 2 and outputs.size(1) >= 6:
+                    for *xyxy, conf, cls_id in outputs.cpu().numpy():
+                        if conf >= self.conf_threshold:
+                            x1, y1, x2, y2 = xyxy
+                            
+                            # 转换为归一化坐标 (x, y, w, h)
+                            x = x1 / original_width
+                            y = y1 / original_height
+                            w = (x2 - x1) / original_width
+                            h = (y2 - y1) / original_height
+                            
+                            # 整数类别 ID
+                            cls_id = int(cls_id)
+                            
+                            detections.append({
+                                'bbox': (x, y, w, h),
+                                'class_id': cls_id,
+                                'confidence': float(conf)
+                            })
+                # 处理其他可能的输出格式
+                else:
+                    # 这里需要根据模型的实际输出格式进行适配
+                    print("警告：无法识别的模型输出格式，请检查模型类型")
+                    
+                return detections
+                
+        except Exception as e:
+            print(f"推理过程中出错: {str(e)}")
+            # 出错时返回空结果
+            return []
+    
+    @property
+    def applies_nms(self) -> bool:
+        """模型是否内部应用了 NMS"""
+        # YOLOv5 会自动应用 NMS
+        return True
+    
+    def close(self):
+        """释放资源"""
+        if hasattr(self, 'model'):
+            # 删除模型以释放 GPU 内存
+            del self.model
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        print("垃圾识别模型已关闭")