import os import sys import json import asyncio from datetime import datetime from pathlib import Path from typing import Dict, Any, Optional import logging # 添加algorithm目录到Python路径 sys.path.append(str(Path(__file__).parent.parent.parent / "algorithm")) from detection import DetectionProcessor, SegmentationProcessor from core.database import get_db from models.device import Device from models.algorithm import Algorithm logger = logging.getLogger(__name__) class VideoProcessor: """视频处理服务""" def __init__(self): self.detection_processor = None self.segmentation_processor = None def _get_processor(self, algorithm_id: int) -> Optional[DetectionProcessor]: """根据算法ID获取对应的处理器""" db = next(get_db()) try: algorithm = db.query(Algorithm).filter(Algorithm.id == algorithm_id).first() if not algorithm or not algorithm.model_path: logger.error(f"算法 {algorithm_id} 不存在或模型路径为空") return None # 检查模型文件是否存在 if not os.path.exists(algorithm.model_path): logger.error(f"模型文件不存在: {algorithm.model_path}") return None # 根据算法名称判断使用检测还是分割 if "分割" in algorithm.name or "segmentation" in algorithm.name.lower(): if not self.segmentation_processor: self.segmentation_processor = SegmentationProcessor(model_path=algorithm.model_path) return self.segmentation_processor else: if not self.detection_processor: self.detection_processor = DetectionProcessor(model_path=algorithm.model_path) return self.detection_processor except Exception as e: logger.error(f"获取处理器失败: {e}") return None finally: db.close() async def process_device_video(self, device_id: int, demo_video_path: str) -> Dict[str, Any]: """ 处理设备的演示视频 Args: device_id: 设备ID demo_video_path: 演示视频路径 Returns: 处理结果字典 """ db = next(get_db()) try: # 获取设备信息 device = db.query(Device).filter(Device.id == device_id).first() if not device: return {"success": False, "error": "设备不存在"} if not device.algorithm_id: return {"success": False, "error": "设备未关联算法"} # 检查视频文件是否存在 if not os.path.exists(demo_video_path): return {"success": False, "error": "演示视频文件不存在"} # 更新设备状态为处理中 device.processing_status = "processing" device.last_processed_at = datetime.now() db.commit() # 获取处理器 processor = self._get_processor(device.algorithm_id) if not processor: device.processing_status = "failed" device.processing_result = json.dumps({"error": "无法获取处理器"}) db.commit() return {"success": False, "error": "无法获取处理器"} # 生成输出视频路径 uploads_dir = Path(__file__).parent.parent / "uploads" / "results" uploads_dir.mkdir(parents=True, exist_ok=True) output_filename = f"processed_{device_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4" output_video_path = str(uploads_dir / output_filename) # 获取算法配置 algorithm = db.query(Algorithm).filter(Algorithm.id == device.algorithm_id).first() detection_classes = json.loads(algorithm.detection_classes) if algorithm.detection_classes else None # 处理视频 logger.info(f"开始处理设备 {device_id} 的视频: {demo_video_path}") if isinstance(processor, DetectionProcessor): # 目标检测处理 result = processor.process_video( input_video_path=demo_video_path, output_video_path=output_video_path, confidence_threshold=0.5, classes=[0] if detection_classes and "person" in detection_classes else None, show_live=False, save_annotated=True ) else: # 图像分割处理 result = processor.process_video( input_video_path=demo_video_path, output_video_path=output_video_path, confidence_threshold=0.3, classes=[0] if detection_classes and "person" in detection_classes else None, show_live=False, save_annotated=True ) # 更新设备状态和结果 if result.get("success"): device.processing_status = "completed" device.processed_video_path = output_video_path device.processing_result = json.dumps(result) device.last_processed_at = datetime.now() logger.info(f"设备 {device_id} 视频处理完成: {output_video_path}") else: device.processing_status = "failed" device.processing_result = json.dumps(result) logger.error(f"设备 {device_id} 视频处理失败: {result.get('error', '未知错误')}") db.commit() return result except Exception as e: logger.error(f"处理设备 {device_id} 视频时发生错误: {e}") try: device.processing_status = "failed" device.processing_result = json.dumps({"error": str(e)}) db.commit() except: pass return {"success": False, "error": str(e)} finally: db.close() def get_processing_status(self, device_id: int) -> Dict[str, Any]: """获取设备处理状态""" db = next(get_db()) try: device = db.query(Device).filter(Device.id == device_id).first() if not device: return {"success": False, "error": "设备不存在"} return { "success": True, "device_id": device_id, "processing_status": device.processing_status, "demo_video_path": device.demo_video_path, "processed_video_path": device.processed_video_path, "last_processed_at": device.last_processed_at.isoformat() if device.last_processed_at else None, "processing_result": json.loads(device.processing_result) if device.processing_result else None } except Exception as e: logger.error(f"获取设备 {device_id} 处理状态时发生错误: {e}") return {"success": False, "error": str(e)} finally: db.close() # 全局视频处理器实例 video_processor = VideoProcessor()