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
0
apps/server/services/__init__.py
Normal file
0
apps/server/services/__init__.py
Normal file
351
apps/server/services/camera_service.py
Normal file
351
apps/server/services/camera_service.py
Normal file
@@ -0,0 +1,351 @@
|
||||
import cv2
|
||||
import json
|
||||
import logging
|
||||
import asyncio
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import platform
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
from typing import Dict, Optional
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class CameraService:
|
||||
def __init__(self, model_service):
|
||||
self.model_service = model_service
|
||||
self.active_connections: Dict[str, WebSocket] = {}
|
||||
self.camera_captures: Dict[str, cv2.VideoCapture] = {}
|
||||
self.running = False
|
||||
self.camera_configs: Dict[str, Dict] = {}
|
||||
self._locks: Dict[str, asyncio.Lock] = {}
|
||||
self._stop_events: Dict[str, asyncio.Event] = {}
|
||||
|
||||
@staticmethod
|
||||
def force_release_cameras():
|
||||
"""强制释放被占用的摄像头资源
|
||||
|
||||
在以下场景调用：
|
||||
1. 服务启动前 - 清理之前异常退出残留的占用
|
||||
2. 服务关闭时 - 确保资源被释放
|
||||
3. 信号处理时 - 异常退出前的清理
|
||||
"""
|
||||
logger.info("强制释放摄像头资源...")
|
||||
|
||||
# 1. 尝试释放当前进程中的摄像头
|
||||
try:
|
||||
# 在macOS上，摄像头设备通常是 /dev/video* 或 AVFoundation 设备
|
||||
# 尝试打开并立即释放来清理状态
|
||||
for i in range(10): # 检查前10个可能的摄像头索引
|
||||
try:
|
||||
cap = cv2.VideoCapture(i)
|
||||
if cap.isOpened():
|
||||
cap.release()
|
||||
logger.info(f"已释放摄像头索引 {i}")
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.error(f"释放当前进程摄像头失败: {e}")
|
||||
|
||||
# 2. 终止占用摄像头的其他Python进程
|
||||
current_pid = os.getpid()
|
||||
system = platform.system()
|
||||
|
||||
try:
|
||||
if system == "Darwin": # macOS
|
||||
# 使用 lsof 查找占用摄像头的进程
|
||||
result = subprocess.run(
|
||||
['lsof', '-c', 'python'],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
if result.returncode == 0:
|
||||
pids_to_kill = set()
|
||||
for line in result.stdout.split('\n'):
|
||||
# 查找包含摄像头设备 (/dev/video* 或 V4L 相关) 的行
|
||||
if any(x in line for x in ['/dev/video', 'Camera', 'V4L']):
|
||||
parts = line.split()
|
||||
if len(parts) >= 2:
|
||||
try:
|
||||
pid = int(parts[1])
|
||||
if pid != current_pid:
|
||||
pids_to_kill.add(pid)
|
||||
except (ValueError, IndexError):
|
||||
pass
|
||||
|
||||
# 终止找到的进程
|
||||
for pid in pids_to_kill:
|
||||
try:
|
||||
logger.info(f"终止占用摄像头的进程: {pid}")
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
# 给进程一点时间优雅退出
|
||||
import time
|
||||
time.sleep(0.5)
|
||||
# 检查是否还在运行，如果是则强制终止
|
||||
try:
|
||||
os.kill(pid, 0) # 检查进程是否存在
|
||||
logger.warning(f"进程 {pid} 未响应 SIGTERM，使用 SIGKILL")
|
||||
os.kill(pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass # 进程已退出
|
||||
except ProcessLookupError:
|
||||
pass # 进程已不存在
|
||||
except Exception as e:
|
||||
logger.error(f"终止进程 {pid} 失败: {e}")
|
||||
|
||||
elif system == "Linux":
|
||||
# Linux 系统使用 fuser 或 lsof
|
||||
for device in ['/dev/video0', '/dev/video1', '/dev/video2']:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
['fuser', '-k', device],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
if result.returncode == 0:
|
||||
logger.info(f"已释放设备 {device}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"终止占用进程失败: {e}")
|
||||
|
||||
# 3. 给系统一点时间完成资源释放
|
||||
import time
|
||||
time.sleep(0.5)
|
||||
logger.info("摄像头资源释放完成")
|
||||
|
||||
async def handle_connection(self, websocket: WebSocket):
|
||||
await websocket.accept()
|
||||
connection_id = id(websocket)
|
||||
self.active_connections[connection_id] = websocket
|
||||
|
||||
logger.info(f"新连接: {connection_id}")
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
data = await websocket.receive_text()
|
||||
message = json.loads(data)
|
||||
|
||||
if message.get('action') == 'start':
|
||||
await self.start_camera(connection_id, websocket, message.get('config'))
|
||||
elif message.get('action') == 'stop':
|
||||
await self.stop_camera(connection_id)
|
||||
elif message.get('action') == 'update_config':
|
||||
await self.update_config(connection_id, message.get('config', {}))
|
||||
except WebSocketDisconnect:
|
||||
logger.info(f"WebSocket连接断开: {connection_id}")
|
||||
break
|
||||
|
||||
except WebSocketDisconnect:
|
||||
logger.info(f"连接断开: {connection_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"连接错误: {connection_id}, {e}")
|
||||
finally:
|
||||
await self.stop_camera(connection_id)
|
||||
if connection_id in self.active_connections:
|
||||
del self.active_connections[connection_id]
|
||||
logger.info(f"连接清理完成: {connection_id}")
|
||||
|
||||
async def start_camera(self, connection_id: str, websocket: WebSocket, config: Dict = None):
|
||||
try:
|
||||
camera_source = 0 # 默认本地摄像头
|
||||
|
||||
if not config:
|
||||
config = {
|
||||
'model_id': 'fire_detection',
|
||||
'confidence': 0.5,
|
||||
'iou': 0.45,
|
||||
'camera_source': 0 # 默认本地摄像头
|
||||
}
|
||||
|
||||
# 支持多种摄像头源
|
||||
if 'camera_source' in config:
|
||||
camera_source = config['camera_source']
|
||||
|
||||
# 如果该连接已有摄像头在运行，先停止它
|
||||
if connection_id in self.camera_captures:
|
||||
await self.stop_camera(connection_id)
|
||||
|
||||
# 初始化锁和停止事件
|
||||
self._locks[connection_id] = asyncio.Lock()
|
||||
self._stop_events[connection_id] = asyncio.Event()
|
||||
|
||||
# 尝试打开摄像头
|
||||
cap = cv2.VideoCapture(camera_source)
|
||||
if not cap.isOpened():
|
||||
await websocket.send_json({
|
||||
'type': 'error',
|
||||
'message': f'无法打开摄像头源: {camera_source}'
|
||||
})
|
||||
return
|
||||
|
||||
# 只对本地摄像头设置分辨率
|
||||
if isinstance(camera_source, int):
|
||||
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
|
||||
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
|
||||
|
||||
self.camera_captures[connection_id] = cap
|
||||
self.camera_configs[connection_id] = config
|
||||
|
||||
await websocket.send_json({
|
||||
'type': 'camera_started',
|
||||
'message': f'摄像头已启动: {camera_source}'
|
||||
})
|
||||
|
||||
await self.process_frames(connection_id, websocket, cap)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"启动摄像头失败: {e}")
|
||||
await websocket.send_json({
|
||||
'type': 'error',
|
||||
'message': f'启动摄像头失败: {str(e)}'
|
||||
})
|
||||
|
||||
async def process_frames(self, connection_id: str, websocket: WebSocket, cap: cv2.VideoCapture):
|
||||
from .detection_service import DetectionService
|
||||
detection_service = DetectionService(self.model_service)
|
||||
|
||||
try:
|
||||
frame_count = 0
|
||||
stop_event = self._stop_events.get(connection_id)
|
||||
|
||||
while connection_id in self.active_connections:
|
||||
# 检查是否收到停止信号
|
||||
if stop_event and stop_event.is_set():
|
||||
logger.info(f"收到停止信号，结束帧处理: {connection_id}")
|
||||
break
|
||||
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
await asyncio.sleep(0.1)
|
||||
continue
|
||||
|
||||
config = self.camera_configs.get(connection_id, {
|
||||
'model_id': 'fire_detection',
|
||||
'confidence': 0.5,
|
||||
'iou': 0.45
|
||||
})
|
||||
|
||||
model_id = config.get('model_id', 'fire_detection')
|
||||
confidence = config.get('confidence', 0.5)
|
||||
iou = config.get('iou', 0.45)
|
||||
draw = True
|
||||
|
||||
processed_frame, result = await detection_service.detect_frame(
|
||||
frame,
|
||||
model_id=model_id,
|
||||
confidence=confidence,
|
||||
iou=iou,
|
||||
draw=draw
|
||||
)
|
||||
|
||||
if result['success']:
|
||||
frame_count += 1
|
||||
|
||||
logger.info(f"发送检测结果: {len(result['detections'])} 个目标, {result['stats']}")
|
||||
|
||||
await websocket.send_json({
|
||||
'type': 'detection',
|
||||
'detections': result['detections'],
|
||||
'stats': result['stats']
|
||||
})
|
||||
|
||||
_, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
|
||||
import base64
|
||||
original_frame_data = base64.b64encode(buffer).decode('utf-8')
|
||||
|
||||
_, buffer = cv2.imencode('.jpg', processed_frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
|
||||
frame_data = base64.b64encode(buffer).decode('utf-8')
|
||||
|
||||
try:
|
||||
await websocket.send_json({
|
||||
'type': 'original_frame',
|
||||
'frame': original_frame_data
|
||||
})
|
||||
|
||||
await websocket.send_json({
|
||||
'type': 'annotated_frame',
|
||||
'frame': frame_data
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"发送帧数据失败: {e}")
|
||||
break
|
||||
|
||||
await asyncio.sleep(0.03)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"处理帧错误: {e}")
|
||||
finally:
|
||||
# 只在摄像头仍存在于字典中时才释放（避免重复释放）
|
||||
if connection_id in self.camera_captures:
|
||||
cap.release()
|
||||
del self.camera_captures[connection_id]
|
||||
logger.info(f"帧处理结束，摄像头已释放: {connection_id}")
|
||||
|
||||
async def stop_camera(self, connection_id: str):
|
||||
# 使用锁确保同一时间只有一个协程在操作该连接的摄像头资源
|
||||
lock = self._locks.get(connection_id)
|
||||
if lock:
|
||||
async with lock:
|
||||
await self._do_stop_camera(connection_id)
|
||||
else:
|
||||
await self._do_stop_camera(connection_id)
|
||||
|
||||
async def _do_stop_camera(self, connection_id: str):
|
||||
"""实际执行停止摄像头的操作（内部方法，应在获取锁后调用）"""
|
||||
# 设置停止事件，通知帧处理循环退出
|
||||
if connection_id in self._stop_events:
|
||||
self._stop_events[connection_id].set()
|
||||
|
||||
if connection_id in self.camera_captures:
|
||||
cap = self.camera_captures[connection_id]
|
||||
cap.release()
|
||||
del self.camera_captures[connection_id]
|
||||
|
||||
if connection_id in self.active_connections:
|
||||
try:
|
||||
await self.active_connections[connection_id].send_json({
|
||||
'type': 'camera_stopped',
|
||||
'message': '摄像头已停止'
|
||||
})
|
||||
except:
|
||||
pass
|
||||
|
||||
logger.info(f"摄像头已停止: {connection_id}")
|
||||
|
||||
# 清理锁和事件
|
||||
if connection_id in self._locks:
|
||||
del self._locks[connection_id]
|
||||
if connection_id in self._stop_events:
|
||||
del self._stop_events[connection_id]
|
||||
|
||||
async def update_config(self, connection_id: str, config: Dict):
|
||||
if connection_id in self.camera_configs:
|
||||
self.camera_configs[connection_id].update(config)
|
||||
|
||||
model_id = self.camera_configs[connection_id].get('model_id', 'fire_detection')
|
||||
confidence = self.camera_configs[connection_id].get('confidence', 0.5)
|
||||
iou = self.camera_configs[connection_id].get('iou', 0.45)
|
||||
|
||||
logger.info(f"配置更新: model_id={model_id}, confidence={confidence}, iou={iou}")
|
||||
|
||||
if connection_id in self.active_connections:
|
||||
try:
|
||||
await self.active_connections[connection_id].send_json({
|
||||
'type': 'config_updated',
|
||||
'message': '配置已更新'
|
||||
})
|
||||
except:
|
||||
pass
|
||||
|
||||
async def stop(self):
|
||||
for connection_id in list(self.camera_captures.keys()):
|
||||
await self.stop_camera(connection_id)
|
||||
|
||||
self.running = False
|
||||
logger.info("摄像头服务已停止")
|
||||
199
apps/server/services/detection_service.py
Normal file
199
apps/server/services/detection_service.py
Normal file
@@ -0,0 +1,199 @@
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import time
|
||||
import uuid
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class DetectionService:
|
||||
def __init__(self, model_service):
|
||||
self.model_service = model_service
|
||||
self.base_dir = os.path.dirname(os.path.dirname(__file__))
|
||||
self.results_dir = os.path.join(self.base_dir, "static", "results")
|
||||
self.temp_dir = os.path.join(self.base_dir, "static", "temp")
|
||||
|
||||
os.makedirs(self.results_dir, exist_ok=True)
|
||||
os.makedirs(self.temp_dir, exist_ok=True)
|
||||
|
||||
def draw_detections(self, frame: np.ndarray, detections: List[Dict], fps: float = 0) -> np.ndarray:
|
||||
try:
|
||||
img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
pil_img = Image.fromarray(img_rgb)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
|
||||
try:
|
||||
font = ImageFont.truetype("/System/Library/Fonts/PingFang.ttc", 20)
|
||||
font_large = ImageFont.truetype("/System/Library/Fonts/PingFang.ttc", 24)
|
||||
except:
|
||||
font = ImageFont.load_default()
|
||||
font_large = font
|
||||
|
||||
class_colors = {
|
||||
'Fire': (255, 0, 0),
|
||||
'Smoke': (128, 128, 128),
|
||||
'person': (0, 255, 0),
|
||||
'helmet': (255, 255, 0),
|
||||
'no_helmet': (255, 0, 255),
|
||||
'cigarette': (0, 165, 255) # 橙色，用于抽烟检测
|
||||
}
|
||||
|
||||
for det in detections:
|
||||
x1, y1, x2, y2 = det['bbox']
|
||||
class_name = det['class']
|
||||
conf = det['confidence']
|
||||
label = det['label']
|
||||
|
||||
color = class_colors.get(class_name, (0, 255, 0))
|
||||
|
||||
draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
|
||||
|
||||
label_text = f"{label} {conf:.2f}"
|
||||
bbox = draw.textbbox((0, 0), label_text, font=font)
|
||||
text_w = bbox[2] - bbox[0]
|
||||
text_h = bbox[3] - bbox[1]
|
||||
draw.rectangle([x1, y1 - text_h - 4, x1 + text_w + 4, y1], fill=color)
|
||||
draw.text((x1 + 2, y1 - text_h - 2), label_text, fill=(255, 255, 255), font=font)
|
||||
|
||||
if fps > 0:
|
||||
fps_text = f"FPS: {fps:.1f} | Detections: {len(detections)}"
|
||||
draw.text((10, 10), fps_text, fill=(0, 255, 0), font=font)
|
||||
|
||||
return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
|
||||
except Exception as e:
|
||||
logger.error(f"绘制检测结果失败: {e}")
|
||||
return frame
|
||||
|
||||
async def detect_image(
|
||||
self,
|
||||
image: np.ndarray,
|
||||
model_id: str,
|
||||
confidence: float = 0.5,
|
||||
iou: float = 0.45
|
||||
) -> Dict:
|
||||
start_time = time.time()
|
||||
|
||||
model = await self.model_service.load_model(model_id)
|
||||
if not model:
|
||||
return {
|
||||
'success': False,
|
||||
'message': f'模型加载失败: {model_id}',
|
||||
'detections': [],
|
||||
'stats': None
|
||||
}
|
||||
|
||||
try:
|
||||
results = model(image, conf=confidence, iou=iou, verbose=False)
|
||||
|
||||
detections = []
|
||||
for result in results:
|
||||
boxes = result.boxes
|
||||
for box in boxes:
|
||||
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
|
||||
conf = float(box.conf[0].cpu().numpy())
|
||||
cls = int(box.cls[0].cpu().numpy())
|
||||
class_name = result.names[cls]
|
||||
|
||||
label_map = self.model_service.model_configs[model_id]['labels']
|
||||
label = label_map.get(class_name, class_name)
|
||||
|
||||
detections.append({
|
||||
'class': class_name,
|
||||
'label': label,
|
||||
'confidence': round(conf, 3),
|
||||
'bbox': [int(x1), int(y1), int(x2), int(y2)]
|
||||
})
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
avg_confidence = sum(d['confidence'] for d in detections) / len(detections) if detections else 0
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'message': '检测完成',
|
||||
'detections': detections,
|
||||
'stats': {
|
||||
'total_detections': len(detections),
|
||||
'avg_confidence': round(avg_confidence, 3),
|
||||
'processing_time': round(processing_time, 3),
|
||||
'model_used': model_id
|
||||
}
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"图片检测失败: {e}")
|
||||
return {
|
||||
'success': False,
|
||||
'message': f'检测失败: {str(e)}',
|
||||
'detections': [],
|
||||
'stats': None
|
||||
}
|
||||
|
||||
async def detect_frame(
|
||||
self,
|
||||
frame: np.ndarray,
|
||||
model_id: str,
|
||||
confidence: float = 0.5,
|
||||
iou: float = 0.45,
|
||||
draw: bool = True
|
||||
) -> tuple:
|
||||
start_time = time.time()
|
||||
|
||||
model = await self.model_service.load_model(model_id)
|
||||
if not model:
|
||||
return frame, {
|
||||
'success': False,
|
||||
'detections': [],
|
||||
'stats': None
|
||||
}
|
||||
|
||||
try:
|
||||
results = model(frame, conf=confidence, iou=iou, verbose=False)
|
||||
|
||||
detections = []
|
||||
for result in results:
|
||||
boxes = result.boxes
|
||||
for box in boxes:
|
||||
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
|
||||
conf = float(box.conf[0].cpu().numpy())
|
||||
cls = int(box.cls[0].cpu().numpy())
|
||||
class_name = result.names[cls]
|
||||
|
||||
label_map = self.model_service.model_configs[model_id]['labels']
|
||||
label = label_map.get(class_name, class_name)
|
||||
|
||||
detections.append({
|
||||
'class': class_name,
|
||||
'label': label,
|
||||
'confidence': round(conf, 3),
|
||||
'bbox': [int(x1), int(y1), int(x2), int(y2)]
|
||||
})
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
fps = 1.0 / processing_time if processing_time > 0 else 0
|
||||
avg_confidence = sum(d['confidence'] for d in detections) / len(detections) if detections else 0
|
||||
|
||||
result_data = {
|
||||
'success': True,
|
||||
'detections': detections,
|
||||
'stats': {
|
||||
'total_detections': len(detections),
|
||||
'avg_confidence': round(avg_confidence, 3),
|
||||
'processing_time': round(processing_time, 3),
|
||||
'fps': round(fps, 2),
|
||||
'model_used': model_id
|
||||
}
|
||||
}
|
||||
|
||||
if draw:
|
||||
frame = self.draw_detections(frame, detections, fps)
|
||||
|
||||
return frame, result_data
|
||||
except Exception as e:
|
||||
logger.error(f"帧检测失败: {e}")
|
||||
return frame, {
|
||||
'success': False,
|
||||
'detections': [],
|
||||
'stats': None
|
||||
}
|
||||
115
apps/server/services/model_service.py
Normal file
115
apps/server/services/model_service.py
Normal file
@@ -0,0 +1,115 @@
|
||||
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] = {}
|
||||
# 基础路径：从 apps/server/services/model_service.py 到 jc-video-web 根目录
|
||||
base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
|
||||
|
||||
self.model_configs = {
|
||||
'fire_detection': {
|
||||
'path': os.path.join(base_dir, 'models', 'fire_detection', 'best.pt'),
|
||||
'type': 'yolov10',
|
||||
'classes': ['Fire', 'Smoke'],
|
||||
'labels': {'Fire': '火焰', 'Smoke': '烟雾'},
|
||||
'size': '61MB',
|
||||
'description': '基于YOLOv10的火灾烟雾检测模型',
|
||||
'name': '火灾检测'
|
||||
},
|
||||
'helmet_detection': {
|
||||
'path': os.path.join(base_dir, 'models', 'helmet_detection', 'yolov8n.pt'),
|
||||
'type': 'yolov8',
|
||||
'classes': ['person', 'helmet'],
|
||||
'labels': {'person': '人员', 'helmet': '安全帽'},
|
||||
'size': '6MB',
|
||||
'description': '基于YOLOv8的安全帽检测模型',
|
||||
'name': '安全帽检测'
|
||||
},
|
||||
'crowd_detection': {
|
||||
'path': os.path.join(base_dir, 'models', 'crowd_detection', 'yolov8l.pt'),
|
||||
'type': 'yolov8',
|
||||
'classes': ['person'],
|
||||
'labels': {'person': '人员'},
|
||||
'size': '100MB',
|
||||
'description': '基于YOLOv8的人群聚集检测模型',
|
||||
'name': '人群检测'
|
||||
},
|
||||
'smoking_detection': {
|
||||
'path': os.path.join(base_dir, 'models', 'smoking_detection', 'smoking_yolov8n.pt'),
|
||||
'type': 'yolov8',
|
||||
'classes': ['cigarette', 'smoke'],
|
||||
'labels': {'cigarette': '香烟', 'smoke': '烟雾'},
|
||||
'size': '6MB',
|
||||
'description': '基于YOLOv8的抽烟检测模型',
|
||||
'name': '抽烟检测'
|
||||
},
|
||||
'loitering_detection': {
|
||||
'path': os.path.join(base_dir, 'models', 'loitering_detection', 'yolov8n.pt'),
|
||||
'type': 'yolov8',
|
||||
'classes': ['person'],
|
||||
'labels': {'person': '人员'},
|
||||
'size': '6MB',
|
||||
'description': '基于YOLOv8的徘徊检测模型',
|
||||
'name': '徘徊检测'
|
||||
}
|
||||
}
|
||||
|
||||
def get_available_models(self) -> List[Dict]:
|
||||
available_models = []
|
||||
for model_id, config in self.model_configs.items():
|
||||
if 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) -> Optional[YOLO]:
|
||||
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]
|
||||
|
||||
# 处理 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) -> Optional[YOLO]:
|
||||
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
|
||||
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
|
||||
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