""" PaddleDetection 抽烟检测服务适配器 通过 Docker 调用 Paddle 模型 """ import os import cv2 import numpy as np import subprocess import tempfile import logging from typing import Dict, List, Optional from pathlib import Path logger = logging.getLogger(__name__) class PaddleDetectionService: """PaddleDetection 服务适配器""" def __init__(self): self.model_name = "smoking_detection" self.docker_image = "smoking-detection:test" self.model_dir = "output_inference/ppyoloe_crn_s_80e_smoking_visdrone" self.threshold = 0.1 # 抽烟检测需要较低的阈值 # 检查 Docker 和镜像 self._check_docker() def _check_docker(self): """检查 Docker 环境""" try: result = subprocess.run( ["docker", "info"], capture_output=True, text=True, timeout=5 ) if result.returncode != 0: logger.error("Docker 未运行") self.available = False return # 检查镜像 result = subprocess.run( ["docker", "image", "inspect", self.docker_image], capture_output=True, text=True, timeout=5 ) self.available = result.returncode == 0 if self.available: logger.info(f"PaddleDetection 服务已就绪: {self.docker_image}") else: logger.error(f"Docker 镜像不存在: {self.docker_image}") except Exception as e: logger.error(f"Docker 检查失败: {e}") self.available = False def detect_image(self, image: np.ndarray) -> Dict: """ 检测图片中的抽烟行为 Args: image: OpenCV 图片 (BGR格式) Returns: 检测结果字典 """ if not self.available: return { 'success': False, 'message': 'PaddleDetection 服务不可用', 'detections': [], 'stats': None } try: # 创建临时文件 with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as f: temp_input = f.name with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as f: temp_output = f.name # 保存输入图片 cv2.imwrite(temp_input, image) # 构建 Docker 命令 cmd = [ "docker", "run", "--rm", "-v", f"{temp_input}:/workspace/input.jpg", "-v", f"{os.path.dirname(temp_output)}:/workspace/output", self.docker_image, "python", "deploy/python/infer.py", f"--model_dir={self.model_dir}", "--image_file=/workspace/input.jpg", "--device=CPU", "--output_dir=/workspace/output", f"--threshold={self.threshold}" ] # 执行检测 logger.info(f"执行抽烟检测: {temp_input}") result = subprocess.run( cmd, capture_output=True, text=True, timeout=60 ) # 解析结果 detections = self._parse_detection_output(result.stdout) # 读取输出图片 output_image = None output_path = temp_output.replace('.jpg', '') + '_result.jpg' if os.path.exists(output_path): output_image = cv2.imread(output_path) # 清理临时文件 self._cleanup_temp_files([temp_input, temp_output, output_path]) return { 'success': True, 'message': '检测完成', 'detections': detections, 'output_image': output_image, 'stats': { 'total_detections': len(detections), 'model_used': 'ppyoloe_crn_s_80e_smoking_visdrone', 'threshold': self.threshold } } except subprocess.TimeoutExpired: logger.error("检测超时") return { 'success': False, 'message': '检测超时', 'detections': [], 'stats': None } except Exception as e: logger.error(f"检测失败: {e}") return { 'success': False, 'message': f'检测失败: {str(e)}', 'detections': [], 'stats': None } def _parse_detection_output(self, output: str) -> List[Dict]: """解析检测输出""" detections = [] # 查找检测结果行 for line in output.split('\n'): if 'class_id:' in line and 'confidence:' in line: try: # 解析: class_id:0, confidence:0.8921, left_top:[268.66,231.64],right_bottom:[351.87,258.66] parts = line.split(',') # 提取置信度 conf_part = [p for p in parts if 'confidence:' in p][0] confidence = float(conf_part.split(':')[1]) # 提取坐标 left_top_part = [p for p in parts if 'left_top:' in p][0] right_bottom_part = [p for p in parts if 'right_bottom:' in p][0] # 解析坐标 left_top = eval(left_top_part.split(':')[1]) right_bottom = eval(right_bottom_part.split(':')[1]) x1, y1 = left_top x2, y2 = right_bottom detections.append({ 'class': 'cigarette', 'label': '香烟', 'confidence': round(confidence, 3), 'bbox': [int(x1), int(y1), int(x2), int(y2)] }) except Exception as e: logger.warning(f"解析检测结果失败: {e}") continue return detections def _cleanup_temp_files(self, files: List[str]): """清理临时文件""" for f in files: try: if os.path.exists(f): os.remove(f) except Exception as e: logger.warning(f"清理临时文件失败: {f}, {e}") # 兼容性包装，保持与 YOLO 模型相同的接口 class SmokingDetectionModel: """抽烟检测模型包装器，兼容 YOLO 接口""" def __init__(self): self.service = PaddleDetectionService() self.names = {0: 'cigarette'} def __call__(self, image, conf=0.1, iou=0.45, verbose=False): """ 模拟 YOLO 模型的调用接口 Args: image: OpenCV 图片 conf: 置信度阈值 iou: IoU 阈值 verbose: 是否输出详细信息 Returns: 模拟 YOLO 结果的对象 """ result = self.service.detect_image(image) # 创建模拟的 YOLO 结果对象 return [PaddleDetectionResult(result, self.names)] class PaddleDetectionResult: """模拟 YOLO 检测结果对象""" def __init__(self, detection_result: Dict, names: Dict): self.detection_result = detection_result self.names = names # 创建模拟的 boxes 对象 self.boxes = self._create_boxes() def _create_boxes(self): """创建模拟的 boxes 对象""" detections = self.detection_result.get('detections', []) if not detections: return MockBoxes([]) # 转换为 YOLO 格式 xyxy = [] conf = [] cls = [] for det in detections: xyxy.append(det['bbox']) conf.append(det['confidence']) cls.append(0) # cigarette 类别 return MockBoxes(xyxy, conf, cls) class MockBoxes: """模拟 YOLO boxes 对象""" def __init__(self, xyxy_list, conf_list=None, cls_list=None): import torch if xyxy_list: self.xyxy = torch.tensor(xyxy_list, dtype=torch.float32) self.conf = torch.tensor(conf_list, dtype=torch.float32).reshape(-1, 1) self.cls = torch.tensor(cls_list, dtype=torch.int64).reshape(-1, 1) else: self.xyxy = torch.empty((0, 4)) self.conf = torch.empty((0, 1)) self.cls = torch.empty((0, 1), dtype=torch.int64)