diff --git a/python-inference-service/app/main.py b/python-inference-service/app/main.py index 3af65bc..f9e1948 100644 --- a/python-inference-service/app/main.py +++ b/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: diff --git a/python-inference-service/models.json b/python-inference-service/models.json new file mode 100644 index 0000000..b81c848 --- /dev/null +++ b/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": "垃圾检测模型" + } +] \ No newline at end of file diff --git a/python-inference-service/models/models.json b/python-inference-service/models/models.json deleted file mode 100644 index a3eaf3b..0000000 --- a/python-inference-service/models/models.json +++ /dev/null @@ -1,8 +0,0 @@ -[ - { - "name": "yolov8_detector", - "path": "models/yolov8_model.py", - "size": [640, 640], - "comment": "YOLOv8检测模型，确保将训练好的best.pt文件放在models目录下" - } -] \ No newline at end of file diff --git a/python-inference-service/models/smoke.pt b/python-inference-service/models/smoke.pt new file mode 100644 index 0000000..ce69aea Binary files /dev/null and b/python-inference-service/models/smoke.pt differ diff --git a/python-inference-service/models/smoke_model.py b/python-inference-service/models/smoke_model.py new file mode 100644 index 0000000..9ab6639 --- /dev/null +++ b/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("垃圾识别模型已关闭") \ No newline at end of file diff --git a/rtsp-vue/src/views/video/model/dict.vue b/rtsp-vue/src/views/video/model/dict.vue new file mode 100644 index 0000000..4084aeb --- /dev/null +++ b/rtsp-vue/src/views/video/model/dict.vue @@ -0,0 +1,341 @@ + + +