diff --git a/python-inference-service/app/detector.py b/python-inference-service/app/detector.py index d155ea1..473f162 100644 --- a/python-inference-service/app/detector.py +++ b/python-inference-service/app/detector.py @@ -12,7 +12,7 @@ from app.models import Detection class PythonModelDetector: """Object detector using native Python models""" - def __init__(self, model_name: str, model_path: str, input_width: int, input_height: int, color: int = 0x00FF00): + def __init__(self, model_name: str, model_path: str, input_width: int, input_height: int, color: int = 0x00FF00, model_config: dict = None): """ Initialize detector with Python model @@ -22,11 +22,13 @@ class PythonModelDetector: input_width: Input width for the model input_height: Input height for the model color: RGB color for detection boxes (default: green) + model_config: Additional configuration to pass to the model """ self.model_name = model_name self.input_width = input_width self.input_height = input_height self.color = color + self.model_config = model_config or {} # Convert color from RGB to BGR (OpenCV uses BGR) self.color_bgr = ((color & 0xFF) << 16) | (color & 0xFF00) | ((color >> 16) & 0xFF) @@ -72,8 +74,18 @@ class PythonModelDetector: if not hasattr(model_module, "Model"): raise AttributeError(f"Model module must define a 'Model' class: {model_path}") - # Create model instance - self.model = model_module.Model() + # Create model instance with config + # Try to pass config to model constructor if it accepts parameters + import inspect + model_class = model_module.Model + sig = inspect.signature(model_class.__init__) + + if len(sig.parameters) > 1: # Has parameters beyond 'self' + # Pass all config as keyword arguments + self.model = model_class(**self.model_config) + else: + # No parameters, create without arguments + self.model = model_class() # Check if model has the required methods if not hasattr(self.model, "predict"): @@ -113,18 +125,16 @@ class PythonModelDetector: # Original image dimensions img_height, img_width = img.shape[:2] - # Preprocess image - processed_img = self.preprocess(img) - # Measure inference time start_time = time.time() try: # Run inference using model's predict method + # Note: Pass original image to model, let it handle preprocessing # Expected return format from model's predict: # List of dicts with keys: 'bbox', 'class_id', 'confidence' # bbox: (x, y, w, h) normalized [0-1] - model_results = self.model.predict(processed_img) + model_results = self.model.predict(img) # Calculate inference time in milliseconds inference_time = (time.time() - start_time) * 1000 @@ -279,13 +289,22 @@ class ModelManager: # Use color from palette color = palette[i % len(palette)] + # Extract model-specific config (model_file, model_name, etc.) + # These will be passed to the Model class __init__ + model_init_config = {} + if "model_file" in model_config: + model_init_config["model_file"] = model_config["model_file"] + if "display_name" in model_config: + model_init_config["model_name"] = model_config["display_name"] + # Create detector for Python model detector = PythonModelDetector( model_name=name, model_path=path, input_width=size[0], input_height=size[1], - color=color + color=color, + model_config=model_init_config ) self.models[name] = detector diff --git a/python-inference-service/app/main.py b/python-inference-service/app/main.py index 0bebd9c..b0044ea 100644 --- a/python-inference-service/app/main.py +++ b/python-inference-service/app/main.py @@ -129,17 +129,13 @@ async def detect_file( model_name: str, file: UploadFile = File(...) ): - # 1. 打印 model_name（直接打印字符串即可） - print(f"接收到的 model_name: {model_name}") - - # 2. 打印 file 的基本信息（文件名、内容类型等） - print(f"文件名: {file.filename}") - print(f"文件内容类型: {file.content_type}") # 例如 image/jpeg、text/plain 等 """Detect objects in an uploaded image file""" + print(f"接收到的 model_name: {model_name}") + print(f"文件名: {file.filename}") + print(f"文件内容类型: {file.content_type}") + global model_manager - - if not model_manager: raise HTTPException(status_code=500, detail="Model manager not initialized") @@ -151,16 +147,39 @@ async def detect_file( # Read uploaded file try: contents = await file.read() + print(f"文件大小: {len(contents)} 字节") + + if len(contents) == 0: + raise HTTPException(status_code=400, detail="Empty file") + nparr = np.frombuffer(contents, np.uint8) + print(f"numpy数组形状: {nparr.shape}, dtype: {nparr.dtype}") + image = cv2.imdecode(nparr, cv2.IMREAD_COLOR) if image is None: - raise HTTPException(status_code=400, detail="Invalid image data") + print("错误: cv2.imdecode 返回 None") + raise HTTPException(status_code=400, detail="Invalid image data - failed to decode") + + print(f"解码后图像形状: {image.shape}, dtype: {image.dtype}") + + except HTTPException: + raise except Exception as e: + print(f"处理图像时出错: {str(e)}") + import traceback + traceback.print_exc() raise HTTPException(status_code=400, detail=f"Failed to process image: {str(e)}") # Run detection - detections, inference_time = detector.detect(image) + try: + detections, inference_time = detector.detect(image) + print(f"检测完成: 找到 {len(detections)} 个目标, 耗时 {inference_time:.2f}ms") + except Exception as e: + print(f"推理过程中出错: {str(e)}") + import traceback + traceback.print_exc() + raise HTTPException(status_code=500, detail=f"Detection failed: {str(e)}") return DetectionResponse( model_name=model_name, diff --git a/python-inference-service/models.json b/python-inference-service/models.json index b81c848..7c30b86 100644 --- a/python-inference-service/models.json +++ b/python-inference-service/models.json @@ -1,14 +1,18 @@ [ { "name": "smoke", - "path": "models/smoke_model.py", + "path": "models/universal_yolo_model.py", + "model_file": "smoke.pt", + "display_name": "吸烟检测", "size": [640, 640], - "comment": "烟雾检测模型" + "comment": "吸烟检测模型 - YOLOv11" }, { "name": "garbage", - "path": "models/garbage_model.py", + "path": "models/universal_yolo_model.py", + "model_file": "garbage.pt", + "display_name": "垃圾识别", "size": [640, 640], - "comment": "垃圾检测模型" + "comment": "垃圾检测模型 - YOLOv8" } ] \ No newline at end of file diff --git a/python-inference-service/models/best.pt b/python-inference-service/models/garbage.pt similarity index 100% rename from python-inference-service/models/best.pt rename to python-inference-service/models/garbage.pt diff --git a/python-inference-service/models/classes.txt b/python-inference-service/models/garbage.txt similarity index 100% rename from python-inference-service/models/classes.txt rename to python-inference-service/models/garbage.txt diff --git a/python-inference-service/models/garbage_model.py b/python-inference-service/models/garbage_model.py deleted file mode 100644 index b53dc81..0000000 --- a/python-inference-service/models/garbage_model.py +++ /dev/null @@ -1,211 +0,0 @@ -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, "best.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, - weights_only=False # 允许加载模型类结构（解决 PyTorch 2.6+ 兼容性问题） - ) - 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/python-inference-service/models/smoke.txt b/python-inference-service/models/smoke.txt new file mode 100644 index 0000000..ec7cbc7 --- /dev/null +++ b/python-inference-service/models/smoke.txt @@ -0,0 +1 @@ +垃圾 \ No newline at end of file diff --git a/python-inference-service/models/smoke_detector.py b/python-inference-service/models/smoke_detector.py deleted file mode 100644 index caea6e0..0000000 --- a/python-inference-service/models/smoke_detector.py +++ /dev/null @@ -1,126 +0,0 @@ -import numpy as np -import cv2 -from typing import List, Dict, Any, Tuple - -class Model: - """ - Smoke detection model implementation - - This is a simple example that could be replaced with an actual - TensorFlow, PyTorch, or other ML framework implementation. - """ - - def __init__(self): - """Initialize smoke detection model""" - # In a real implementation, you would load your model here - print("Smoke detection model initialized") - - # Define smoke class IDs - self.smoke_classes = { - 0: "smoke", - 1: "fire" - } - - def preprocess(self, image: np.ndarray) -> np.ndarray: - """Preprocess image for model input""" - # Convert BGR to grayscale for smoke detection - gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - # Convert back to 3 channels to match model expected input shape - gray_3ch = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR) - - # In a real implementation, you would do normalization, etc. - return gray_3ch - - def predict(self, image: np.ndarray) -> List[Dict[str, Any]]: - """ - Run smoke detection on the image - - This is a simplified example that uses basic image processing - In a real implementation, you would use your ML model - """ - # Convert to grayscale for processing - gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - - # Apply Gaussian blur to reduce noise - blurred = cv2.GaussianBlur(gray, (15, 15), 0) - - # Simple thresholding to find potential smoke regions - # In a real implementation, you'd use a trained model - _, thresh = cv2.threshold(blurred, 100, 255, cv2.THRESH_BINARY) - - # Find contours in the thresholded image - contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - - # Process contours to find potential smoke regions - detections = [] - height, width = image.shape[:2] - - for contour in contours: - # Get bounding box - x, y, w, h = cv2.boundingRect(contour) - - # Filter small regions - if w > width * 0.05 and h > height * 0.05: - # Calculate area ratio - area = cv2.contourArea(contour) - rect_area = w * h - fill_ratio = area / rect_area if rect_area > 0 else 0 - - # Smoke tends to have irregular shapes - # This is just for demonstration purposes - if fill_ratio > 0.2 and fill_ratio < 0.8: - # Normalize coordinates - x_norm = x / width - y_norm = y / height - w_norm = w / width - h_norm = h / height - - # Determine if it's smoke or fire (just a simple heuristic for demo) - # In a real model, this would be determined by the model prediction - class_id = 0 # Default to smoke - - # Check if the region has high red values (fire) - roi = image[y:y+h, x:x+w] - if roi.size > 0: # Make sure ROI is not empty - avg_color = np.mean(roi, axis=(0, 1)) - if avg_color[2] > 150 and avg_color[2] > avg_color[0] * 1.5: # High red, indicating fire - class_id = 1 # Fire - - # Calculate confidence based on fill ratio - # This is just for demonstration - confidence = 0.5 + fill_ratio * 0.3 - - # Add to detections - detections.append({ - 'bbox': (x_norm, y_norm, w_norm, h_norm), - 'class_id': class_id, - 'confidence': confidence - }) - - # For demo purposes, if no smoke detected by algorithm, - # add a small chance of random detection - if not detections and np.random.random() < 0.1: # 10% chance - # Random smoke detection - x = np.random.random() * 0.7 - y = np.random.random() * 0.7 - w = 0.1 + np.random.random() * 0.2 - h = 0.1 + np.random.random() * 0.2 - confidence = 0.5 + np.random.random() * 0.3 - - detections.append({ - 'bbox': (x, y, w, h), - 'class_id': 0, # Smoke - 'confidence': confidence - }) - - return detections - - @property - def applies_nms(self) -> bool: - """Model does not apply NMS internally""" - return False - - def close(self): - """Release resources""" - # In a real implementation, you would release model resources here - pass \ No newline at end of file diff --git a/python-inference-service/models/smoke_model.py b/python-inference-service/models/smoke_model.py deleted file mode 100644 index a61db2a..0000000 --- a/python-inference-service/models/smoke_model.py +++ /dev/null @@ -1,211 +0,0 @@ -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, - weights_only=False # 允许加载模型类结构，解决 PyTorch 2.6+ 兼容性问题 - ) - 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/python-inference-service/models/yolov8_model.py b/python-inference-service/models/universal_yolo_model.py similarity index 62% rename from python-inference-service/models/yolov8_model.py rename to python-inference-service/models/universal_yolo_model.py index e0167a3..dc12bc6 100644 --- a/python-inference-service/models/yolov8_model.py +++ b/python-inference-service/models/universal_yolo_model.py @@ -6,17 +6,39 @@ import torch class Model: """ - YOLOv8 模型包装类 - 使用 Ultralytics YOLO + 通用 YOLO 模型 - 支持 YOLOv8/YOLOv11 等基于 Ultralytics 的模型 """ - def __init__(self): - """初始化YOLOv8模型""" + def __init__(self, model_file: str = None, model_name: str = "YOLO"): + """ + 初始化模型 + + Args: + model_file: 模型文件名（如 smoke.pt, best.pt） + model_name: 模型显示名称（用于日志） + """ # 获取当前文件所在目录路径 model_dir = os.path.dirname(os.path.abspath(__file__)) - # 模型文件路径 - model_path = os.path.join(model_dir, "best.pt") - print(f"正在加载YOLOv8模型: {model_path}") + # 如果没有指定模型文件，尝试常见的文件名 + if model_file is None: + for possible_file in ['garbage.pt', 'smoke.pt', 'best.pt', 'yolov8.pt', 'model.pt']: + test_path = os.path.join(model_dir, possible_file) + if os.path.exists(test_path): + model_file = possible_file + break + + if model_file is None: + raise FileNotFoundError(f"未找到模型文件，请在初始化时指定 model_file 参数") + + # 模型文件路径 + model_path = os.path.join(model_dir, model_file) + + if not os.path.exists(model_path): + raise FileNotFoundError(f"模型文件不存在: {model_path}") + + self.model_name = model_name + print(f"正在加载{model_name}模型: {model_path}") # 检查设备 self.device = "cuda" if torch.cuda.is_available() else "cpu" @@ -26,19 +48,30 @@ class Model: try: from ultralytics import YOLO self.model = YOLO(model_path) - print("使用 Ultralytics YOLO 加载模型成功") + print(f"使用 Ultralytics YOLO 加载模型成功") except ImportError: raise ImportError("请安装 ultralytics: pip install ultralytics>=8.0.0") except Exception as e: - raise Exception(f"加载YOLOv8模型失败: {str(e)}") + raise Exception(f"加载{model_name}模型失败: {str(e)}") # 加载类别名称 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: + + # 1. 首先尝试加载与模型文件同名的类别文件（如 smoke.txt） + model_base_name = os.path.splitext(model_file)[0] + classes_path_specific = os.path.join(model_dir, f"{model_base_name}.txt") + + # 2. 然后尝试加载通用的 classes.txt + classes_path_generic = os.path.join(model_dir, "classes.txt") + + if os.path.exists(classes_path_specific): + with open(classes_path_specific, 'r', encoding='utf-8') as f: self.classes = [line.strip() for line in f.readlines() if line.strip()] - print(f"已加载 {len(self.classes)} 个类别") + print(f"已加载类别文件: {model_base_name}.txt ({len(self.classes)} 个类别)") + elif os.path.exists(classes_path_generic): + with open(classes_path_generic, 'r', encoding='utf-8') as f: + self.classes = [line.strip() for line in f.readlines() if line.strip()] + print(f"已加载类别文件: classes.txt ({len(self.classes)} 个类别)") else: # 使用模型自带的类别信息 if hasattr(self.model, 'names') and self.model.names: @@ -51,10 +84,10 @@ class Model: self.conf_threshold = 0.25 # 置信度阈值 self.img_size = 640 # 默认输入图像大小 - print("YOLOv8模型加载完成") + print(f"{model_name}模型加载完成") def preprocess(self, image: np.ndarray) -> np.ndarray: - """预处理图像 - YOLOv8会自动处理，这里直接返回""" + """预处理图像 - Ultralytics YOLO 会自动处理，这里直接返回""" return image def predict(self, image: np.ndarray) -> List[Dict[str, Any]]: @@ -62,7 +95,7 @@ class Model: original_height, original_width = image.shape[:2] try: - # YOLOv8推理 + # YOLO 推理 results = self.model( image, conf=self.conf_threshold, @@ -122,7 +155,7 @@ class Model: @property def applies_nms(self) -> bool: """模型是否内部应用了 NMS""" - # YOLOv8会自动应用 NMS + # Ultralytics YOLO 会自动应用 NMS return True def close(self): @@ -132,4 +165,5 @@ class Model: del self.model if torch.cuda.is_available(): torch.cuda.empty_cache() - print("YOLOv8模型已关闭") \ No newline at end of file + print(f"{self.model_name}模型已关闭") + diff --git a/python-inference-service/models/yolov5_utils.py b/python-inference-service/models/yolov5_utils.py deleted file mode 100644 index 9ea29aa..0000000 --- a/python-inference-service/models/yolov5_utils.py +++ /dev/null @@ -1,56 +0,0 @@ -import torch -import torch.nn as nn -import sys -import os - -def attempt_load(weights, device=''): - """尝试加载YOLOv5模型""" - # 加载模型 - model = torch.load(weights, map_location=device) - - # 确定模型格式 - if isinstance(model, dict): - if 'model' in model: # state_dict格式 - model = model['model'] - elif 'state_dict' in model: # state_dict格式 - model = model['state_dict'] - - # 如果是state_dict，则需要创建模型架构 - if isinstance(model, dict): - print("警告：加载的是权重字典，尝试创建默认模型结构") - from models.yolov5_model import YOLOv5 - model_arch = YOLOv5() - model_arch.load_state_dict(model) - model = model_arch - - # 设置为评估模式 - if isinstance(model, nn.Module): - model.eval() - - # 检查是否有类别信息 - if not hasattr(model, 'names') or not model.names: - print("模型没有类别信息，尝试加载默认类别") - # 设置通用类别 - model.names = ['object'] - - return model - -class YOLOv5: - """简化版YOLOv5模型结构，用于加载权重""" - def __init__(self): - super(YOLOv5, self).__init__() - self.names = [] # 类别名称 - # 这里应该添加真实的网络结构 - # 但为了简单起见，我们只提供一个占位符 - # 在实际使用中，您应该实现完整的网络架构 - - def forward(self, x): - # 这里应该是实际的前向传播逻辑 - # 这只是一个占位符 - raise NotImplementedError("这是一个占位符模型，请使用完整的YOLOv5模型实现") - - def load_state_dict(self, state_dict): - print("尝试加载模型权重") - # 实际的权重加载逻辑 - # 这只是一个占位符 - return self \ No newline at end of file diff --git a/ruoyi-video/src/main/java/com/ruoyi/video/service/VideoAnalysisService.java b/ruoyi-video/src/main/java/com/ruoyi/video/service/VideoAnalysisService.java index 8b73bd3..40ce9e5 100644 --- a/ruoyi-video/src/main/java/com/ruoyi/video/service/VideoAnalysisService.java +++ b/ruoyi-video/src/main/java/com/ruoyi/video/service/VideoAnalysisService.java @@ -65,9 +65,9 @@ public class VideoAnalysisService { @Autowired private com.ruoyi.video.mapper.InspectionTaskRecordMapper inspectionTaskRecordMapper; - // 检测器配置 - 使用容器名而不是localhost - private static final String PYTHON_API_URL = "http://localhost:8000/api/detect/file"; - private static final String MODEL_NAME = "yolov8_detector"; + // 检测器配置 - 支持环境变量配置 + private static final String PYTHON_API_URL = System.getenv().getOrDefault("PYTHON_API_URL", "http://localhost:8000") + "/api/detect/file"; + private static final String MODEL_NAME = "smoke"; // 默认使用吸烟检测模型 /** * 分析视频并更新记录（同步调用） diff --git a/ruoyi-video/src/main/resources/libs/models/models.json b/ruoyi-video/src/main/resources/libs/models/models.json index 38e5aa4..cfe4917 100644 --- a/ruoyi-video/src/main/resources/libs/models/models.json +++ b/ruoyi-video/src/main/resources/libs/models/models.json @@ -1,4 +1,16 @@ [ - {"name":"smoke","path":"libs/models/smoke","size":[640,640],"backend":"opencv"}, - {"name":"garbage","path":"libs/models/garbage","size":[640,640],"backend":"opencv"} + { + "name": "smoke", + "pythonModelName": "smoke", + "pythonApiUrl": "http://localhost:8000/api/detect/file", + "size": [640, 640], + "backend": "python" + }, + { + "name": "garbage", + "pythonModelName": "garbage", + "pythonApiUrl": "http://localhost:8000/api/detect/file", + "size": [640, 640], + "backend": "python" + } ]