jc-video-recognize/apps/server/services/model_service.py

import os
import logging
from ultralytics import YOLO
from typing import Dict, List, Optional, Union

logger = logging.getLogger(__name__)

class ModelService:
    def __init__(self):
        self.models: Dict[str, Union[YOLO, object]] = {}
        # 基础路径：从 apps/server/services/model_service.py 到 jc-video-web 根目录
        base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

        self.model_configs = {
            'fire_detection': {
                'path': os.path.join(base_dir, 'models', 'fire_detection', 'best.pt'),
                'type': 'yolov10',
                'classes': ['Fire', 'Smoke'],
                'labels': {'Fire': '火焰', 'Smoke': '烟雾'},
                'size': '61MB',
                'description': '基于YOLOv10的火灾烟雾检测模型',
                'name': '火灾检测'
            },
            'helmet_detection': {
                'path': os.path.join(base_dir, 'models', 'helmet_detection', 'yolov8n.pt'),
                'type': 'yolov8',
                'classes': ['person', 'helmet'],
                'labels': {'person': '人员', 'helmet': '安全帽'},
                'size': '6MB',
                'description': '基于YOLOv8的安全帽检测模型',
                'name': '安全帽检测'
            },
            'crowd_detection': {
                'path': os.path.join(base_dir, 'models', 'crowd_detection', 'yolov8l.pt'),
                'type': 'yolov8',
                'classes': ['person'],
                'labels': {'person': '人员'},
                'size': '100MB',
                'description': '基于YOLOv8的人群聚集检测模型',
                'name': '人群检测'
            },
            'smoking_detection': {
                'path': os.path.join(base_dir, 'models', 'smoking_detection', 'smoking_yolov8n.pt'),
                'type': 'yolov8',
                'classes': ['cigarette', 'smoke'],
                'labels': {'cigarette': '香烟', 'smoke': '烟雾'},
                'size': '6MB',
                'description': '基于YOLOv8的抽烟检测模型',
                'name': '抽烟检测 (YOLOv8)'
            },
            'smoking_detection_paddle': {
                'path': os.path.join(base_dir, 'models', 'smoking_detection_paddle', 'model.pdmodel'),
                'type': 'paddle',
                'classes': ['cigarette'],
                'labels': {'cigarette': '香烟'},
                'size': '27MB',
                'description': '基于PaddlePaddle PP-YOLOE-s的抽烟检测模型（更高准确率）',
                'name': '抽烟检测 (Paddle)'
            },
            'loitering_detection': {
                'path': os.path.join(base_dir, 'models', 'loitering_detection', 'yolov8n.pt'),
                'type': 'yolov8',
                'classes': ['person'],
                'labels': {'person': '人员'},
                'size': '6MB',
                'description': '基于YOLOv8的徘徊检测模型',
                'name': '徘徊检测'
            }
        }
    
    def get_available_models(self) -> List[Dict]:
        available_models = []
        for model_id, config in self.model_configs.items():
            model_path = config['path']
            
            # 检查模型是否存在（Paddle模型检查目录，YOLO模型检查文件）
            model_exists = False
            if config['type'] == 'paddle':
                model_dir = os.path.dirname(model_path)
                required_files = ['model.pdmodel', 'model.pdiparams', 'infer_cfg.yml']
                model_exists = all(
                    os.path.exists(os.path.join(model_dir, f))
                    for f in required_files
                )
            else:
                model_exists = os.path.exists(model_path)
            
            if model_exists:
                available_models.append({
                    'id': model_id,
                    'name': config['name'],
                    'description': config['description'],
                    'classes': config['classes'],
                    'labels': config['labels'],
                    'size': config['size'],
                    'type': config['type']
                })
            else:
                logger.warning(f"模型文件不存在: {model_path}")
        return available_models
    
    async def load_model(self, model_id: str) -> Optional[Union[YOLO, object]]:
        if model_id not in self.model_configs:
            logger.error(f"未知模型ID: {model_id}")
            return None
        
        if model_id in self.models:
            return self.models[model_id]
        
        config = self.model_configs[model_id]

        # 处理 PaddleDetection 模型
        if config['type'] == 'paddle':
            try:
                from .paddle_detection_service import SmokingDetectionModel
                logger.info(f"正在加载 PaddlePaddle Docker 服务: {model_id}")
                model = SmokingDetectionModel()
                self.models[model_id] = model
                logger.info(f"PaddlePaddle Docker 服务加载成功: {model_id}")
                return model
            except Exception as e:
                logger.error(f"PaddlePaddle Docker 服务加载失败: {model_id}, 错误: {e}")
                return None

        # 处理 YOLO 模型
        model_path = config['path']
        
        if not os.path.exists(model_path):
            logger.error(f"模型文件不存在: {model_path}")
            return None
        
        try:
            logger.info(f"正在加载 YOLO 模型: {model_id} from {model_path}")
            model = YOLO(model_path)
            self.models[model_id] = model
            logger.info(f"YOLO 模型加载成功: {model_id}")
            return model
        except Exception as e:
            logger.error(f"YOLO 模型加载失败: {model_id}, 错误: {e}")
            return None
    
    def get_model(self, model_id: str) -> Optional[Union[YOLO, object]]:
        return self.models.get(model_id)
    
    async def unload_model(self, model_id: str) -> bool:
        if model_id in self.models:
            del self.models[model_id]
            logger.info(f"模型已卸载: {model_id}")
            return True
        return False