border_Inspection/server/services/video_processor.py

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()