Initial commit: Video detection platform with YOLO models
Features: - Fire detection (YOLOv10) - Helmet detection (YOLOv8) - Crowd detection (YOLOv8) - Smoking detection (YOLOv8) - Loitering detection (YOLOv8) Tech Stack: - Frontend: Vue 3 + Vite + Element Plus - Backend: FastAPI + WebSocket - Monorepo: pnpm workspace + Turbo - Docker support included
This commit is contained in:
wwh
147
apps/server/services/model_service_updated.py
Normal file
147
apps/server/services/model_service_updated.py
Normal file
@@ -0,0 +1,147 @@
|
||||
import os
|
||||
import logging
|
||||
from ultralytics import YOLO
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ModelService:
|
||||
def __init__(self):
|
||||
self.models: Dict[str, YOLO] = {}
|
||||
self.model_configs = {
|
||||
'fire_detection': {
|
||||
'path': os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
|
||||
'fire_detection', 'models', 'best.pt'),
|
||||
'type': 'yolov10',
|
||||
'classes': ['Fire', 'Smoke'],
|
||||
'labels': {'Fire': '火焰', 'Smoke': '烟雾'},
|
||||
'size': '61MB',
|
||||
'description': '基于YOLOv10的火灾烟雾检测模型',
|
||||
'name': '火灾检测'
|
||||
},
|
||||
'helmet_detection': {
|
||||
'path': os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
|
||||
'yolov', 'yolov8n.pt'),
|
||||
'type': 'yolov8',
|
||||
'classes': ['person', 'helmet'],
|
||||
'labels': {'person': '人员', 'helmet': '安全帽'},
|
||||
'size': '6MB',
|
||||
'description': '基于YOLOv8的安全帽检测模型',
|
||||
'name': '安全帽检测'
|
||||
},
|
||||
'crowd_detection': {
|
||||
'path': os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
|
||||
'behavior_detection', 'Crowd-Gathering', 'models', 'yolov8l.pt'),
|
||||
'type': 'yolov8',
|
||||
'classes': ['person'],
|
||||
'labels': {'person': '人员'},
|
||||
'size': '100MB',
|
||||
'description': '基于YOLOv8的人群聚集检测模型',
|
||||
'name': '人群检测'
|
||||
},
|
||||
'smoking_detection': {
|
||||
'path': 'PADDLE_DETECTION', # 特殊标记，表示使用 PaddleDetection
|
||||
'type': 'paddle',
|
||||
'classes': ['cigarette'],
|
||||
'labels': {'cigarette': '香烟'},
|
||||
'size': '27MB',
|
||||
'description': '基于PP-YOLOE的抽烟检测模型',
|
||||
'name': '抽烟检测',
|
||||
'docker_image': 'smoking-detection:test',
|
||||
'model_dir': 'output_inference/ppyoloe_crn_s_80e_smoking_visdrone'
|
||||
}
|
||||
}
|
||||
|
||||
def get_available_models(self) -> List[Dict]:
|
||||
available_models = []
|
||||
for model_id, config in self.model_configs.items():
|
||||
# 对于 PaddleDetection 模型，不需要检查本地文件
|
||||
if config.get('type') == 'paddle':
|
||||
# 检查 Docker 是否可用
|
||||
import subprocess
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "image", "inspect", config['docker_image']],
|
||||
capture_output=True,
|
||||
timeout=5
|
||||
)
|
||||
if result.returncode == 0:
|
||||
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"PaddleDetection Docker 镜像不可用: {config['docker_image']}")
|
||||
except Exception as e:
|
||||
logger.warning(f"检查 PaddleDetection Docker 镜像失败: {e}")
|
||||
elif os.path.exists(config['path']):
|
||||
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"模型文件不存在: {config['path']}")
|
||||
return available_models
|
||||
|
||||
async def load_model(self, model_id: str):
|
||||
if model_id not in self.model_configs:
|
||||
logger.error(f"未知模型ID: {model_id}")
|
||||
return None
|
||||
|
||||
# 如果已经加载，直接返回
|
||||
if model_id in self.models:
|
||||
logger.info(f"模型已加载: {model_id}")
|
||||
return self.models[model_id]
|
||||
|
||||
config = self.model_configs[model_id]
|
||||
|
||||
# 处理 PaddleDetection 模型
|
||||
if config.get('type') == 'paddle':
|
||||
try:
|
||||
from .paddle_detection_service import SmokingDetectionModel
|
||||
|
||||
logger.info(f"正在加载 PaddleDetection 抽烟检测模型...")
|
||||
model = SmokingDetectionModel()
|
||||
self.models[model_id] = model
|
||||
logger.info(f"PaddleDetection 模型加载成功: {model_id}")
|
||||
return model
|
||||
except Exception as e:
|
||||
logger.error(f"PaddleDetection 模型加载失败: {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"正在加载模型: {model_id} from {model_path}")
|
||||
model = YOLO(model_path)
|
||||
self.models[model_id] = model
|
||||
logger.info(f"模型加载成功: {model_id}")
|
||||
return model
|
||||
except Exception as e:
|
||||
logger.error(f"模型加载失败: {model_id}, 错误: {e}")
|
||||
return None
|
||||
|
||||
def get_model(self, model_id: str):
|
||||
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
|
||||
Reference in New Issue
Block a user