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
274
apps/server/services/paddle_detection_service.py
Normal file
274
apps/server/services/paddle_detection_service.py
Normal file
@@ -0,0 +1,274 @@
|
||||
"""
|
||||
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)
|
||||
Reference in New Issue
Block a user