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/models/__init__.py
Normal file
0
apps/server/models/__init__.py
Normal file
39
apps/server/models/schemas.py
Normal file
39
apps/server/models/schemas.py
Normal file
@@ -0,0 +1,39 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Dict, Any
|
||||
from enum import Enum
|
||||
|
||||
class ModelInfo(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
description: str
|
||||
classes: List[str]
|
||||
labels: Dict[str, str]
|
||||
size: str
|
||||
type: str
|
||||
|
||||
class Detection(BaseModel):
|
||||
class_name: str
|
||||
label: str
|
||||
confidence: float
|
||||
bbox: List[int]
|
||||
|
||||
class DetectionStats(BaseModel):
|
||||
total_detections: int
|
||||
avg_confidence: float
|
||||
processing_time: float
|
||||
model_used: str
|
||||
|
||||
class ImageDetectionResult(BaseModel):
|
||||
success: bool
|
||||
message: str
|
||||
data: Dict[str, Any]
|
||||
|
||||
class VideoDetectionRequest(BaseModel):
|
||||
model_id: str
|
||||
confidence: float = Field(default=0.5, ge=0.1, le=1.0)
|
||||
iou: float = Field(default=0.45, ge=0.1, le=0.9)
|
||||
|
||||
class DetectionConfig(BaseModel):
|
||||
model_id: str
|
||||
confidence: float = Field(default=0.5, ge=0.1, le=1.0)
|
||||
iou: float = Field(default=0.45, ge=0.1, le=0.9)
|
||||
357
apps/server/models/smoking_yolo_adapter.py
Normal file
357
apps/server/models/smoking_yolo_adapter.py
Normal file
@@ -0,0 +1,357 @@
|
||||
"""
|
||||
YOLO 格式的抽烟检测模型适配器
|
||||
将 PaddleDetection 模型包装为 YOLO 接口
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import tempfile
|
||||
import logging
|
||||
from typing import List, Dict, Union
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SmokingDetectionYOLO:
|
||||
"""
|
||||
模拟 YOLO 接口的抽烟检测模型
|
||||
底层使用 PaddleDetection Docker 容器
|
||||
"""
|
||||
|
||||
def __init__(self, model_path=None):
|
||||
"""
|
||||
初始化模型
|
||||
|
||||
Args:
|
||||
model_path: 模型路径（可选，仅用于兼容性）
|
||||
"""
|
||||
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
|
||||
|
||||
# YOLO 兼容属性
|
||||
self.names = {0: 'cigarette'}
|
||||
self.model = self # 自我引用，保持与 YOLO 相同的接口
|
||||
|
||||
# 检查 Docker
|
||||
self._check_docker()
|
||||
|
||||
logger.info(f"抽烟检测模型初始化完成，Docker可用: {self.available}")
|
||||
|
||||
def _check_docker(self):
|
||||
"""检查 Docker 环境"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "info"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
if self.available:
|
||||
# 检查镜像
|
||||
result = subprocess.run(
|
||||
["docker", "image", "inspect", self.docker_image],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Docker 检查失败: {e}")
|
||||
self.available = False
|
||||
|
||||
def __call__(self, source, conf=0.1, iou=0.45, verbose=False, stream=False):
|
||||
"""
|
||||
模拟 YOLO 模型的调用接口
|
||||
|
||||
Args:
|
||||
source: 图片路径、OpenCV 图片、或图片列表
|
||||
conf: 置信度阈值
|
||||
iou: IoU 阈值（PaddleDetection 不支持，仅用于兼容）
|
||||
verbose: 是否输出详细信息
|
||||
stream: 是否流式输出（仅用于兼容）
|
||||
|
||||
Returns:
|
||||
YOLOResult 对象列表
|
||||
"""
|
||||
if not self.available:
|
||||
logger.error("Docker 不可用，无法运行检测")
|
||||
return [YOLOResult([])]
|
||||
|
||||
# 处理不同类型的输入
|
||||
if isinstance(source, str):
|
||||
# 图片路径
|
||||
image = cv2.imread(source)
|
||||
if image is None:
|
||||
logger.error(f"无法读取图片: {source}")
|
||||
return [YOLOResult([])]
|
||||
return self._detect_single(image, conf, verbose)
|
||||
|
||||
elif isinstance(source, np.ndarray):
|
||||
# OpenCV 图片
|
||||
return self._detect_single(source, conf, verbose)
|
||||
|
||||
elif isinstance(source, list):
|
||||
# 图片列表
|
||||
results = []
|
||||
for img in source:
|
||||
if isinstance(img, str):
|
||||
img = cv2.imread(img)
|
||||
if img is not None:
|
||||
results.extend(self._detect_single(img, conf, verbose))
|
||||
return results
|
||||
|
||||
else:
|
||||
logger.error(f"不支持的输入类型: {type(source)}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def _detect_single(self, image: np.ndarray, conf: float, verbose: bool) -> List['YOLOResult']:
|
||||
"""检测单张图片"""
|
||||
|
||||
try:
|
||||
# 创建临时文件
|
||||
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as f:
|
||||
temp_input = f.name
|
||||
|
||||
# 保存输入图片
|
||||
cv2.imwrite(temp_input, image)
|
||||
|
||||
if verbose:
|
||||
logger.info(f"正在检测: {temp_input}")
|
||||
|
||||
# 构建 Docker 命令
|
||||
cmd = [
|
||||
"docker", "run", "--rm",
|
||||
"-v", f"{temp_input}:/workspace/input.jpg",
|
||||
self.docker_image,
|
||||
"python", "deploy/python/infer.py",
|
||||
f"--model_dir={self.model_dir}",
|
||||
"--image_file=/workspace/input.jpg",
|
||||
"--device=CPU",
|
||||
"--output_dir=/workspace",
|
||||
f"--threshold={conf}"
|
||||
]
|
||||
|
||||
# 执行检测
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60
|
||||
)
|
||||
|
||||
if verbose:
|
||||
logger.info(f"检测完成，返回码: {result.returncode}")
|
||||
|
||||
# 解析结果
|
||||
detections = self._parse_output(result.stdout)
|
||||
|
||||
if verbose and detections:
|
||||
logger.info(f"检测到 {len(detections)} 个目标")
|
||||
|
||||
# 清理临时文件
|
||||
try:
|
||||
os.remove(temp_input)
|
||||
except:
|
||||
pass
|
||||
|
||||
return [YOLOResult(detections)]
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.error("检测超时")
|
||||
return [YOLOResult([])]
|
||||
except Exception as e:
|
||||
logger.error(f"检测失败: {e}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def _parse_output(self, output: str) -> List[Dict]:
|
||||
"""解析检测输出"""
|
||||
detections = []
|
||||
import re
|
||||
|
||||
# 使用正则表达式匹配检测行
|
||||
# 格式: class_id:0, confidence:0.8921, left_top:[268.66,231.64],right_bottom:[351.87,258.66]
|
||||
pattern = r'class_id:\d+,\s*confidence:([\d.]+),\s*left_top:\[([\d.]+),\s*([\d.]+)\],\s*right_bottom:\[([\d.]+),\s*([\d.]+)\]'
|
||||
|
||||
for line in output.split('\n'):
|
||||
match = re.search(pattern, line)
|
||||
if match:
|
||||
try:
|
||||
confidence = float(match.group(1))
|
||||
x1 = float(match.group(2))
|
||||
y1 = float(match.group(3))
|
||||
x2 = float(match.group(4))
|
||||
y2 = float(match.group(5))
|
||||
|
||||
detections.append({
|
||||
'bbox': [int(x1), int(y1), int(x2), int(y2)],
|
||||
'confidence': confidence,
|
||||
'class': 0,
|
||||
'name': 'cigarette'
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"解析检测结果失败: {e}, line: {line}")
|
||||
continue
|
||||
|
||||
return detections
|
||||
|
||||
def predict(self, source, **kwargs):
|
||||
"""兼容 predict 方法"""
|
||||
return self.__call__(source, **kwargs)
|
||||
|
||||
|
||||
class YOLOResult:
|
||||
"""
|
||||
模拟 YOLO 检测结果对象
|
||||
提供与 ultralytics YOLO 结果相同的接口
|
||||
"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
self.names = {0: 'cigarette'}
|
||||
|
||||
# 创建 boxes 对象
|
||||
self.boxes = Boxes(detections)
|
||||
|
||||
# 其他 YOLO 结果属性
|
||||
self.probs = None
|
||||
self.keypoints = None
|
||||
self.obb = None
|
||||
self.speed = {'preprocess': 0, 'inference': 0, 'postprocess': 0}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
"""支持索引访问"""
|
||||
if idx < len(self.detections):
|
||||
return YOLOResult([self.detections[idx]])
|
||||
return YOLOResult([])
|
||||
|
||||
def plot(self, **kwargs):
|
||||
"""绘制检测结果（兼容方法）"""
|
||||
return None
|
||||
|
||||
|
||||
class Boxes:
|
||||
"""
|
||||
模拟 YOLO boxes 对象
|
||||
提供 xyxy, conf, cls 等属性
|
||||
"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
|
||||
# 尝试使用 torch，如果没有则使用 numpy
|
||||
try:
|
||||
import torch
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = torch.tensor(xyxy_list, dtype=torch.float32)
|
||||
self.conf = torch.tensor(conf_list, dtype=torch.float32)
|
||||
self.cls = torch.tensor(cls_list, dtype=torch.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = torch.empty((0, 4))
|
||||
self.conf = torch.empty((0, 1))
|
||||
self.cls = torch.empty((0, 1), dtype=torch.int64)
|
||||
self.id = None
|
||||
|
||||
except ImportError:
|
||||
# 如果没有 torch，使用 numpy
|
||||
import numpy as np
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = np.array(xyxy_list, dtype=np.float32)
|
||||
self.conf = np.array(conf_list, dtype=np.float32)
|
||||
self.cls = np.array(cls_list, dtype=np.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = np.empty((0, 4), dtype=np.float32)
|
||||
self.conf = np.empty((0, 1), dtype=np.float32)
|
||||
self.cls = np.empty((0, 1), dtype=np.int64)
|
||||
self.id = None
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __iter__(self):
|
||||
"""使 Boxes 可迭代"""
|
||||
for i in range(len(self.detections)):
|
||||
yield Box(self, i)
|
||||
|
||||
def cpu(self):
|
||||
"""兼容方法"""
|
||||
return self
|
||||
|
||||
def numpy(self):
|
||||
"""转换为 numpy"""
|
||||
if hasattr(self.xyxy, 'numpy'):
|
||||
return type('Boxes', (), {
|
||||
'xyxy': self.xyxy.numpy(),
|
||||
'conf': self.conf.numpy(),
|
||||
'cls': self.cls.numpy(),
|
||||
'id': self.id
|
||||
})()
|
||||
return self
|
||||
|
||||
|
||||
class Box:
|
||||
"""
|
||||
模拟单个检测框对象
|
||||
"""
|
||||
|
||||
def __init__(self, boxes: Boxes, index: int):
|
||||
self._boxes = boxes
|
||||
self._index = index
|
||||
|
||||
@property
|
||||
def xyxy(self):
|
||||
"""返回 xyxy 坐标 (1, 4) 形状"""
|
||||
import torch
|
||||
import numpy as np
|
||||
coords = self._boxes.xyxy[self._index]
|
||||
if isinstance(coords, torch.Tensor):
|
||||
return coords.unsqueeze(0)
|
||||
else:
|
||||
return np.array([coords])
|
||||
|
||||
@property
|
||||
def conf(self):
|
||||
"""返回置信度 (1,) 形状 - 与 YOLO 兼容"""
|
||||
import torch
|
||||
import numpy as np
|
||||
conf_val = self._boxes.conf[self._index]
|
||||
# 返回 (1,) 形状，与 YOLO 一致
|
||||
if isinstance(conf_val, torch.Tensor):
|
||||
return conf_val.view(1)
|
||||
else:
|
||||
return np.array([conf_val])
|
||||
|
||||
@property
|
||||
def cls(self):
|
||||
"""返回类别 (1,) 形状 - 与 YOLO 兼容"""
|
||||
import torch
|
||||
import numpy as np
|
||||
cls_val = self._boxes.cls[self._index]
|
||||
# 返回 (1,) 形状，与 YOLO 一致
|
||||
if isinstance(cls_val, torch.Tensor):
|
||||
return cls_val.view(1)
|
||||
else:
|
||||
return np.array([cls_val])
|
||||
359
apps/server/models/smoking_yolo_adapter_fast.py
Normal file
359
apps/server/models/smoking_yolo_adapter_fast.py
Normal file
@@ -0,0 +1,359 @@
|
||||
"""
|
||||
YOLO 格式的抽烟检测模型适配器（快速版）
|
||||
使用 HTTP API 与常驻 Docker 容器通信
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import tempfile
|
||||
import logging
|
||||
import time
|
||||
import json
|
||||
import requests
|
||||
from typing import List, Dict, Optional
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SmokingDetectionYOLO:
|
||||
"""
|
||||
模拟 YOLO 接口的抽烟检测模型（快速版）
|
||||
使用 HTTP API 与常驻 Docker 容器通信
|
||||
"""
|
||||
|
||||
_container_name = "smoking-detection-daemon"
|
||||
_initialized = False
|
||||
_server_url = "http://localhost:8080"
|
||||
|
||||
def __init__(self, model_path=None):
|
||||
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
|
||||
|
||||
# YOLO 兼容属性
|
||||
self.names = {0: 'cigarette'}
|
||||
self.model = self
|
||||
|
||||
# 检查 Docker 并启动常驻服务器
|
||||
self._check_docker()
|
||||
if self.available:
|
||||
self._start_daemon()
|
||||
|
||||
logger.info(f"抽烟检测模型快速版初始化完成，Docker可用: {self.available}")
|
||||
|
||||
def _check_docker(self):
|
||||
"""检查 Docker 环境"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "info"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
if self.available:
|
||||
result = subprocess.run(
|
||||
["docker", "image", "inspect", self.docker_image],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Docker 检查失败: {e}")
|
||||
self.available = False
|
||||
|
||||
def _start_daemon(self):
|
||||
"""启动常驻服务器"""
|
||||
try:
|
||||
# 检查容器是否已在运行
|
||||
result = subprocess.run(
|
||||
["docker", "ps", "-q", "-f", f"name={self._container_name}"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
|
||||
if result.stdout.strip():
|
||||
logger.info(f"常驻服务器已在运行: {self._container_name}")
|
||||
SmokingDetectionYOLO._initialized = True
|
||||
return
|
||||
|
||||
# 检查容器是否存在但已停止
|
||||
result = subprocess.run(
|
||||
["docker", "ps", "-aq", "-f", f"name={self._container_name}"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
|
||||
if result.stdout.strip():
|
||||
# 删除旧容器
|
||||
logger.info("删除旧容器")
|
||||
subprocess.run(
|
||||
["docker", "rm", "-f", self._container_name],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
|
||||
# 创建新容器并启动常驻服务器
|
||||
logger.info("启动常驻服务器...")
|
||||
subprocess.run(
|
||||
[
|
||||
"docker", "run", "-d",
|
||||
"--name", self._container_name,
|
||||
"-p", "8080:8080",
|
||||
"-v", "/tmp:/workspace/input",
|
||||
"-v", "/Users/wwh/project/video-model/PaddlePaddle/PaddleDetection-release-2.9/smoking_server_daemon.py:/workspace/PaddleDetection/smoking_server_daemon.py",
|
||||
"-w", "/workspace/PaddleDetection",
|
||||
self.docker_image,
|
||||
"python", "smoking_server_daemon.py"
|
||||
],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
|
||||
# 等待服务器启动
|
||||
logger.info("等待服务器启动...")
|
||||
time.sleep(5)
|
||||
|
||||
SmokingDetectionYOLO._initialized = True
|
||||
logger.info("常驻服务器启动成功")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"启动常驻服务器失败: {e}")
|
||||
SmokingDetectionYOLO._initialized = False
|
||||
|
||||
def __call__(self, source, conf=0.1, iou=0.45, verbose=False, stream=False):
|
||||
"""模拟 YOLO 模型的调用接口"""
|
||||
if not self.available:
|
||||
logger.error("Docker 不可用，无法运行检测")
|
||||
return [YOLOResult([])]
|
||||
|
||||
# 处理不同类型的输入
|
||||
if isinstance(source, str):
|
||||
image = cv2.imread(source)
|
||||
if image is None:
|
||||
logger.error(f"无法读取图片: {source}")
|
||||
return [YOLOResult([])]
|
||||
return self._detect_single(image, conf, verbose)
|
||||
|
||||
elif isinstance(source, np.ndarray):
|
||||
return self._detect_single(source, conf, verbose)
|
||||
|
||||
elif isinstance(source, list):
|
||||
results = []
|
||||
for img in source:
|
||||
if isinstance(img, str):
|
||||
img = cv2.imread(img)
|
||||
if img is not None:
|
||||
results.extend(self._detect_single(img, conf, verbose))
|
||||
return results
|
||||
|
||||
else:
|
||||
logger.error(f"不支持的输入类型: {type(source)}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def _detect_single(self, image: np.ndarray, conf: float, verbose: bool) -> List['YOLOResult']:
|
||||
"""检测单张图片（使用 HTTP API）"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 创建临时文件
|
||||
input_filename = f"smoking_fast_{int(time.time()*1000)}.jpg"
|
||||
temp_input = f"/tmp/{input_filename}"
|
||||
|
||||
# 保存输入图片
|
||||
cv2.imwrite(temp_input, image)
|
||||
|
||||
if verbose:
|
||||
logger.info(f"正在检测: {temp_input}")
|
||||
|
||||
# 发送 HTTP 请求
|
||||
request = {
|
||||
'image_path': f'/workspace/input/{input_filename}',
|
||||
'threshold': conf
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
f"{self._server_url}/detect",
|
||||
json=request,
|
||||
timeout=30
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
if verbose:
|
||||
logger.info(f"检测完成，耗时: {elapsed:.2f}秒")
|
||||
|
||||
# 解析结果
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
if data.get('success'):
|
||||
detections = data.get('detections', [])
|
||||
else:
|
||||
logger.error(f"检测失败: {data.get('error')}")
|
||||
detections = []
|
||||
else:
|
||||
logger.error(f"HTTP 错误: {response.status_code}")
|
||||
detections = []
|
||||
|
||||
# 清理临时文件
|
||||
try:
|
||||
os.remove(temp_input)
|
||||
except:
|
||||
pass
|
||||
|
||||
return [YOLOResult(detections)]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"检测失败: {e}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def predict(self, source, **kwargs):
|
||||
"""兼容 predict 方法"""
|
||||
return self.__call__(source, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def stop_daemon(cls):
|
||||
"""停止常驻服务器"""
|
||||
try:
|
||||
subprocess.run(
|
||||
["docker", "stop", cls._container_name],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
logger.info("常驻服务器已停止")
|
||||
except Exception as e:
|
||||
logger.error(f"停止常驻服务器失败: {e}")
|
||||
|
||||
|
||||
# YOLOResult, Boxes, Box 类（与之前相同）
|
||||
class YOLOResult:
|
||||
"""模拟 YOLO 检测结果对象"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
self.names = {0: 'cigarette'}
|
||||
self.boxes = Boxes(detections)
|
||||
self.probs = None
|
||||
self.keypoints = None
|
||||
self.obb = None
|
||||
self.speed = {'preprocess': 0, 'inference': 0, 'postprocess': 0}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if idx < len(self.detections):
|
||||
return YOLOResult([self.detections[idx]])
|
||||
return YOLOResult([])
|
||||
|
||||
def plot(self, **kwargs):
|
||||
return None
|
||||
|
||||
|
||||
class Boxes:
|
||||
"""模拟 YOLO boxes 对象"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = torch.tensor(xyxy_list, dtype=torch.float32)
|
||||
self.conf = torch.tensor(conf_list, dtype=torch.float32)
|
||||
self.cls = torch.tensor(cls_list, dtype=torch.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = torch.empty((0, 4))
|
||||
self.conf = torch.empty((0, 1))
|
||||
self.cls = torch.empty((0, 1), dtype=torch.int64)
|
||||
self.id = None
|
||||
|
||||
except ImportError:
|
||||
import numpy as np
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = np.array(xyxy_list, dtype=np.float32)
|
||||
self.conf = np.array(conf_list, dtype=np.float32)
|
||||
self.cls = np.array(cls_list, dtype=np.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = np.empty((0, 4), dtype=np.float32)
|
||||
self.conf = np.empty((0, 1), dtype=np.float32)
|
||||
self.cls = np.empty((0, 1), dtype=np.int64)
|
||||
self.id = None
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __iter__(self):
|
||||
for i in range(len(self.detections)):
|
||||
yield Box(self, i)
|
||||
|
||||
def cpu(self):
|
||||
return self
|
||||
|
||||
def numpy(self):
|
||||
if hasattr(self.xyxy, 'numpy'):
|
||||
return type('Boxes', (), {
|
||||
'xyxy': self.xyxy.numpy(),
|
||||
'conf': self.conf.numpy(),
|
||||
'cls': self.cls.numpy(),
|
||||
'id': self.id
|
||||
})()
|
||||
return self
|
||||
|
||||
|
||||
class Box:
|
||||
"""模拟单个检测框对象"""
|
||||
|
||||
def __init__(self, boxes: Boxes, index: int):
|
||||
self._boxes = boxes
|
||||
self._index = index
|
||||
|
||||
@property
|
||||
def xyxy(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
coords = self._boxes.xyxy[self._index]
|
||||
if isinstance(coords, torch.Tensor):
|
||||
return coords.unsqueeze(0)
|
||||
else:
|
||||
return np.array([coords])
|
||||
|
||||
@property
|
||||
def conf(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
conf_val = self._boxes.conf[self._index]
|
||||
if isinstance(conf_val, torch.Tensor):
|
||||
return conf_val.view(1)
|
||||
else:
|
||||
return np.array([conf_val])
|
||||
|
||||
@property
|
||||
def cls(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
cls_val = self._boxes.cls[self._index]
|
||||
if isinstance(cls_val, torch.Tensor):
|
||||
return cls_val.view(1)
|
||||
else:
|
||||
return np.array([cls_val])
|
||||
399
apps/server/models/smoking_yolo_adapter_optimized.py
Normal file
399
apps/server/models/smoking_yolo_adapter_optimized.py
Normal file
@@ -0,0 +1,399 @@
|
||||
"""
|
||||
YOLO 格式的抽烟检测模型适配器（优化版）
|
||||
使用后台 Docker 容器，避免重复启动
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import tempfile
|
||||
import logging
|
||||
import time
|
||||
from typing import List, Dict, Optional
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SmokingDetectionYOLO:
|
||||
"""
|
||||
模拟 YOLO 接口的抽烟检测模型（优化版）
|
||||
使用后台 Docker 容器，避免每次检测都启动新容器
|
||||
"""
|
||||
|
||||
_container_name = "smoking-detection-server"
|
||||
_container_started = False
|
||||
|
||||
def __init__(self, model_path=None):
|
||||
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
|
||||
|
||||
# YOLO 兼容属性
|
||||
self.names = {0: 'cigarette'}
|
||||
self.model = self
|
||||
|
||||
# 检查 Docker 并启动后台容器
|
||||
self._check_docker()
|
||||
if self.available:
|
||||
self._start_background_container()
|
||||
|
||||
logger.info(f"抽烟检测模型初始化完成，Docker可用: {self.available}")
|
||||
|
||||
def _check_docker(self):
|
||||
"""检查 Docker 环境"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "info"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
if self.available:
|
||||
result = subprocess.run(
|
||||
["docker", "image", "inspect", self.docker_image],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Docker 检查失败: {e}")
|
||||
self.available = False
|
||||
|
||||
def _start_background_container(self):
|
||||
"""启动后台容器"""
|
||||
try:
|
||||
# 检查容器是否已在运行
|
||||
result = subprocess.run(
|
||||
["docker", "ps", "-q", "-f", f"name={self._container_name}"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
|
||||
if result.stdout.strip():
|
||||
logger.info(f"后台容器已在运行: {self._container_name}")
|
||||
SmokingDetectionYOLO._container_started = True
|
||||
return
|
||||
|
||||
# 检查容器是否存在但已停止
|
||||
result = subprocess.run(
|
||||
["docker", "ps", "-aq", "-f", f"name={self._container_name}"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
|
||||
if result.stdout.strip():
|
||||
# 启动已存在的容器
|
||||
logger.info(f"启动已存在的容器: {self._container_name}")
|
||||
subprocess.run(
|
||||
["docker", "start", self._container_name],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
else:
|
||||
# 创建新容器
|
||||
logger.info(f"创建后台容器: {self._container_name}")
|
||||
subprocess.run(
|
||||
[
|
||||
"docker", "run", "-d",
|
||||
"--name", self._container_name,
|
||||
"-v", "/tmp:/workspace/input",
|
||||
"-v", "/tmp:/workspace/output",
|
||||
self.docker_image,
|
||||
"tail", "-f", "/dev/null"
|
||||
],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
|
||||
SmokingDetectionYOLO._container_started = True
|
||||
logger.info("后台容器启动成功")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"启动后台容器失败: {e}")
|
||||
SmokingDetectionYOLO._container_started = False
|
||||
|
||||
def __call__(self, source, conf=0.1, iou=0.45, verbose=False, stream=False):
|
||||
"""模拟 YOLO 模型的调用接口"""
|
||||
if not self.available:
|
||||
logger.error("Docker 不可用，无法运行检测")
|
||||
return [YOLOResult([])]
|
||||
|
||||
# 处理不同类型的输入
|
||||
if isinstance(source, str):
|
||||
image = cv2.imread(source)
|
||||
if image is None:
|
||||
logger.error(f"无法读取图片: {source}")
|
||||
return [YOLOResult([])]
|
||||
return self._detect_single(image, conf, verbose)
|
||||
|
||||
elif isinstance(source, np.ndarray):
|
||||
return self._detect_single(source, conf, verbose)
|
||||
|
||||
elif isinstance(source, list):
|
||||
results = []
|
||||
for img in source:
|
||||
if isinstance(img, str):
|
||||
img = cv2.imread(img)
|
||||
if img is not None:
|
||||
results.extend(self._detect_single(img, conf, verbose))
|
||||
return results
|
||||
|
||||
else:
|
||||
logger.error(f"不支持的输入类型: {type(source)}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def _detect_single(self, image: np.ndarray, conf: float, verbose: bool) -> List['YOLOResult']:
|
||||
"""检测单张图片（使用后台容器）"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 创建临时文件
|
||||
input_filename = f"smoking_input_{int(time.time()*1000)}.jpg"
|
||||
output_filename = f"smoking_output_{int(time.time()*1000)}.jpg"
|
||||
temp_input = f"/tmp/{input_filename}"
|
||||
temp_output = f"/tmp/{output_filename}"
|
||||
|
||||
# 保存输入图片
|
||||
cv2.imwrite(temp_input, image)
|
||||
|
||||
if verbose:
|
||||
logger.info(f"正在检测: {temp_input}")
|
||||
|
||||
# 使用后台容器执行检测
|
||||
if SmokingDetectionYOLO._container_started:
|
||||
# 使用 exec 在运行中的容器内执行
|
||||
cmd = [
|
||||
"docker", "exec",
|
||||
self._container_name,
|
||||
"python", "deploy/python/infer.py",
|
||||
f"--model_dir={self.model_dir}",
|
||||
f"--image_file=/workspace/input/{input_filename}",
|
||||
"--device=CPU",
|
||||
f"--output_dir=/workspace/output",
|
||||
f"--threshold={conf}"
|
||||
]
|
||||
else:
|
||||
# 回退到原来的方式
|
||||
cmd = [
|
||||
"docker", "run", "--rm",
|
||||
"-v", f"{temp_input}:/workspace/input.jpg",
|
||||
self.docker_image,
|
||||
"python", "deploy/python/infer.py",
|
||||
f"--model_dir={self.model_dir}",
|
||||
"--image_file=/workspace/input.jpg",
|
||||
"--device=CPU",
|
||||
"--output_dir=/workspace",
|
||||
f"--threshold={conf}"
|
||||
]
|
||||
|
||||
# 执行检测
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
if verbose:
|
||||
logger.info(f"检测完成，耗时: {elapsed:.2f}秒")
|
||||
|
||||
# 解析结果
|
||||
detections = self._parse_output(result.stdout)
|
||||
|
||||
# 清理临时文件
|
||||
try:
|
||||
os.remove(temp_input)
|
||||
except:
|
||||
pass
|
||||
|
||||
return [YOLOResult(detections)]
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.error("检测超时")
|
||||
return [YOLOResult([])]
|
||||
except Exception as e:
|
||||
logger.error(f"检测失败: {e}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def _parse_output(self, output: str) -> List[Dict]:
|
||||
"""解析检测输出"""
|
||||
detections = []
|
||||
import re
|
||||
|
||||
pattern = r'class_id:\d+,\s*confidence:([\d.]+),\s*left_top:\[([\d.]+),\s*([\d.]+)\],\s*right_bottom:\[([\d.]+),\s*([\d.]+)\]'
|
||||
|
||||
for line in output.split('\n'):
|
||||
match = re.search(pattern, line)
|
||||
if match:
|
||||
try:
|
||||
confidence = float(match.group(1))
|
||||
x1 = float(match.group(2))
|
||||
y1 = float(match.group(3))
|
||||
x2 = float(match.group(4))
|
||||
y2 = float(match.group(5))
|
||||
|
||||
detections.append({
|
||||
'bbox': [int(x1), int(y1), int(x2), int(y2)],
|
||||
'confidence': confidence,
|
||||
'class': 0,
|
||||
'name': 'cigarette'
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"解析检测结果失败: {e}, line: {line}")
|
||||
continue
|
||||
|
||||
return detections
|
||||
|
||||
def predict(self, source, **kwargs):
|
||||
"""兼容 predict 方法"""
|
||||
return self.__call__(source, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def stop_background_container(cls):
|
||||
"""停止后台容器"""
|
||||
try:
|
||||
subprocess.run(
|
||||
["docker", "stop", cls._container_name],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
logger.info("后台容器已停止")
|
||||
except Exception as e:
|
||||
logger.error(f"停止后台容器失败: {e}")
|
||||
|
||||
|
||||
# 复制 YOLOResult, Boxes, Box 类（与原版相同）
|
||||
class YOLOResult:
|
||||
"""模拟 YOLO 检测结果对象"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
self.names = {0: 'cigarette'}
|
||||
self.boxes = Boxes(detections)
|
||||
self.probs = None
|
||||
self.keypoints = None
|
||||
self.obb = None
|
||||
self.speed = {'preprocess': 0, 'inference': 0, 'postprocess': 0}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if idx < len(self.detections):
|
||||
return YOLOResult([self.detections[idx]])
|
||||
return YOLOResult([])
|
||||
|
||||
def plot(self, **kwargs):
|
||||
return None
|
||||
|
||||
|
||||
class Boxes:
|
||||
"""模拟 YOLO boxes 对象"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = torch.tensor(xyxy_list, dtype=torch.float32)
|
||||
self.conf = torch.tensor(conf_list, dtype=torch.float32)
|
||||
self.cls = torch.tensor(cls_list, dtype=torch.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = torch.empty((0, 4))
|
||||
self.conf = torch.empty((0, 1))
|
||||
self.cls = torch.empty((0, 1), dtype=torch.int64)
|
||||
self.id = None
|
||||
|
||||
except ImportError:
|
||||
import numpy as np
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = np.array(xyxy_list, dtype=np.float32)
|
||||
self.conf = np.array(conf_list, dtype=np.float32)
|
||||
self.cls = np.array(cls_list, dtype=np.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = np.empty((0, 4), dtype=np.float32)
|
||||
self.conf = np.empty((0, 1), dtype=np.float32)
|
||||
self.cls = np.empty((0, 1), dtype=np.int64)
|
||||
self.id = None
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __iter__(self):
|
||||
for i in range(len(self.detections)):
|
||||
yield Box(self, i)
|
||||
|
||||
def cpu(self):
|
||||
return self
|
||||
|
||||
def numpy(self):
|
||||
if hasattr(self.xyxy, 'numpy'):
|
||||
return type('Boxes', (), {
|
||||
'xyxy': self.xyxy.numpy(),
|
||||
'conf': self.conf.numpy(),
|
||||
'cls': self.cls.numpy(),
|
||||
'id': self.id
|
||||
})()
|
||||
return self
|
||||
|
||||
|
||||
class Box:
|
||||
"""模拟单个检测框对象"""
|
||||
|
||||
def __init__(self, boxes: Boxes, index: int):
|
||||
self._boxes = boxes
|
||||
self._index = index
|
||||
|
||||
@property
|
||||
def xyxy(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
coords = self._boxes.xyxy[self._index]
|
||||
if isinstance(coords, torch.Tensor):
|
||||
return coords.unsqueeze(0)
|
||||
else:
|
||||
return np.array([coords])
|
||||
|
||||
@property
|
||||
def conf(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
conf_val = self._boxes.conf[self._index]
|
||||
if isinstance(conf_val, torch.Tensor):
|
||||
return conf_val.view(1)
|
||||
else:
|
||||
return np.array([conf_val])
|
||||
|
||||
@property
|
||||
def cls(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
cls_val = self._boxes.cls[self._index]
|
||||
if isinstance(cls_val, torch.Tensor):
|
||||
return cls_val.view(1)
|
||||
else:
|
||||
return np.array([cls_val])
|
||||
379
apps/server/models/smoking_yolo_adapter_v2.py
Normal file
379
apps/server/models/smoking_yolo_adapter_v2.py
Normal file
@@ -0,0 +1,379 @@
|
||||
"""
|
||||
YOLO 格式的抽烟检测模型适配器（V2 - 常驻进程版）
|
||||
使用 Docker 容器内的常驻 Python 进程，避免每次检测都启动新进程
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import tempfile
|
||||
import logging
|
||||
import time
|
||||
import json
|
||||
from typing import List, Dict, Optional
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SmokingDetectionYOLO:
|
||||
"""
|
||||
模拟 YOLO 接口的抽烟检测模型（V2 - 常驻进程版）
|
||||
使用 Docker 容器内的常驻 Python 进程
|
||||
"""
|
||||
|
||||
_container_name = "smoking-detection-v2"
|
||||
_process = None
|
||||
_initialized = False
|
||||
|
||||
def __init__(self, model_path=None):
|
||||
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
|
||||
|
||||
# YOLO 兼容属性
|
||||
self.names = {0: 'cigarette'}
|
||||
self.model = self
|
||||
|
||||
# 检查 Docker 并启动常驻进程
|
||||
self._check_docker()
|
||||
if self.available:
|
||||
self._start_server()
|
||||
|
||||
logger.info(f"抽烟检测模型 V2 初始化完成，Docker可用: {self.available}")
|
||||
|
||||
def _check_docker(self):
|
||||
"""检查 Docker 环境"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "info"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
if self.available:
|
||||
result = subprocess.run(
|
||||
["docker", "image", "inspect", self.docker_image],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Docker 检查失败: {e}")
|
||||
self.available = False
|
||||
|
||||
def _start_server(self):
|
||||
"""启动常驻服务器进程"""
|
||||
try:
|
||||
# 检查是否已有进程在运行
|
||||
if SmokingDetectionYOLO._process is not None:
|
||||
logger.info("常驻进程已在运行")
|
||||
return
|
||||
|
||||
# 检查容器是否已存在
|
||||
result = subprocess.run(
|
||||
["docker", "ps", "-aq", "-f", f"name={self._container_name}"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
|
||||
if result.stdout.strip():
|
||||
# 删除旧容器
|
||||
logger.info("删除旧容器")
|
||||
subprocess.run(
|
||||
["docker", "rm", "-f", self._container_name],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
|
||||
# 启动新容器并运行服务器
|
||||
logger.info("启动常驻服务器...")
|
||||
|
||||
# 获取 smoking_server.py 的绝对路径
|
||||
server_script_path = "/Users/wwh/project/video-model/PaddlePaddle/PaddleDetection-release-2.9/smoking_server.py"
|
||||
|
||||
# 使用 Popen 保持进程运行，挂载 server 脚本
|
||||
SmokingDetectionYOLO._process = subprocess.Popen(
|
||||
[
|
||||
"docker", "run", "-i", "--rm",
|
||||
"--name", self._container_name,
|
||||
"-v", "/tmp:/workspace/input",
|
||||
"-v", f"{server_script_path}:/workspace/PaddleDetection/smoking_server.py",
|
||||
"-w", "/workspace/PaddleDetection",
|
||||
self.docker_image,
|
||||
"python", "smoking_server.py"
|
||||
],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
bufsize=1
|
||||
)
|
||||
|
||||
# 等待服务器启动（读取模型加载完成的消息）
|
||||
logger.info("等待服务器启动...")
|
||||
start_wait = time.time()
|
||||
while time.time() - start_wait < 30: # 最多等待30秒
|
||||
if SmokingDetectionYOLO._process.poll() is not None:
|
||||
# 进程已退出
|
||||
stderr = SmokingDetectionYOLO._process.stderr.read()
|
||||
logger.error(f"服务器启动失败: {stderr}")
|
||||
SmokingDetectionYOLO._process = None
|
||||
return
|
||||
|
||||
# 尝试读取 stderr 看是否加载完成
|
||||
import select
|
||||
if SmokingDetectionYOLO._process.stderr:
|
||||
ready, _, _ = select.select([SmokingDetectionYOLO._process.stderr], [], [], 0.5)
|
||||
if ready:
|
||||
line = SmokingDetectionYOLO._process.stderr.readline()
|
||||
if line:
|
||||
logger.info(f"Server: {line.strip()}")
|
||||
if "模型加载完成" in line:
|
||||
break
|
||||
time.sleep(0.1)
|
||||
|
||||
# 检查进程是否还在运行
|
||||
if SmokingDetectionYOLO._process.poll() is None:
|
||||
SmokingDetectionYOLO._initialized = True
|
||||
logger.info("常驻服务器启动成功")
|
||||
else:
|
||||
stderr = SmokingDetectionYOLO._process.stderr.read()
|
||||
logger.error(f"服务器启动失败: {stderr}")
|
||||
SmokingDetectionYOLO._process = None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"启动常驻服务器失败: {e}")
|
||||
SmokingDetectionYOLO._process = None
|
||||
|
||||
def __call__(self, source, conf=0.1, iou=0.45, verbose=False, stream=False):
|
||||
"""模拟 YOLO 模型的调用接口"""
|
||||
if not self.available:
|
||||
logger.error("Docker 不可用，无法运行检测")
|
||||
return [YOLOResult([])]
|
||||
|
||||
if not SmokingDetectionYOLO._initialized:
|
||||
logger.error("常驻服务器未初始化")
|
||||
return [YOLOResult([])]
|
||||
|
||||
# 处理不同类型的输入
|
||||
if isinstance(source, str):
|
||||
image = cv2.imread(source)
|
||||
if image is None:
|
||||
logger.error(f"无法读取图片: {source}")
|
||||
return [YOLOResult([])]
|
||||
return self._detect_single(image, conf, verbose)
|
||||
|
||||
elif isinstance(source, np.ndarray):
|
||||
return self._detect_single(source, conf, verbose)
|
||||
|
||||
elif isinstance(source, list):
|
||||
results = []
|
||||
for img in source:
|
||||
if isinstance(img, str):
|
||||
img = cv2.imread(img)
|
||||
if img is not None:
|
||||
results.extend(self._detect_single(img, conf, verbose))
|
||||
return results
|
||||
|
||||
else:
|
||||
logger.error(f"不支持的输入类型: {type(source)}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def _detect_single(self, image: np.ndarray, conf: float, verbose: bool) -> List['YOLOResult']:
|
||||
"""检测单张图片（使用常驻进程）"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 创建临时文件
|
||||
input_filename = f"smoking_v2_{int(time.time()*1000)}.jpg"
|
||||
temp_input = f"/tmp/{input_filename}"
|
||||
|
||||
# 保存输入图片
|
||||
cv2.imwrite(temp_input, image)
|
||||
|
||||
if verbose:
|
||||
logger.info(f"正在检测: {temp_input}")
|
||||
|
||||
# 发送请求到常驻进程
|
||||
request = {
|
||||
'image_path': f'/workspace/input/{input_filename}',
|
||||
'threshold': conf
|
||||
}
|
||||
|
||||
SmokingDetectionYOLO._process.stdin.write(json.dumps(request) + '\n')
|
||||
SmokingDetectionYOLO._process.stdin.flush()
|
||||
|
||||
# 读取响应
|
||||
response_line = SmokingDetectionYOLO._process.stdout.readline()
|
||||
response = json.loads(response_line)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
if verbose:
|
||||
logger.info(f"检测完成，耗时: {elapsed:.2f}秒")
|
||||
|
||||
# 解析结果
|
||||
if response.get('success'):
|
||||
detections = response.get('detections', [])
|
||||
else:
|
||||
logger.error(f"检测失败: {response.get('error')}")
|
||||
detections = []
|
||||
|
||||
# 清理临时文件
|
||||
try:
|
||||
os.remove(temp_input)
|
||||
except:
|
||||
pass
|
||||
|
||||
return [YOLOResult(detections)]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"检测失败: {e}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def predict(self, source, **kwargs):
|
||||
"""兼容 predict 方法"""
|
||||
return self.__call__(source, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def stop_server(cls):
|
||||
"""停止常驻服务器"""
|
||||
if cls._process is not None:
|
||||
cls._process.terminate()
|
||||
cls._process.wait()
|
||||
cls._process = None
|
||||
cls._initialized = False
|
||||
logger.info("常驻服务器已停止")
|
||||
|
||||
|
||||
# YOLOResult, Boxes, Box 类（与之前相同）
|
||||
class YOLOResult:
|
||||
"""模拟 YOLO 检测结果对象"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
self.names = {0: 'cigarette'}
|
||||
self.boxes = Boxes(detections)
|
||||
self.probs = None
|
||||
self.keypoints = None
|
||||
self.obb = None
|
||||
self.speed = {'preprocess': 0, 'inference': 0, 'postprocess': 0}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if idx < len(self.detections):
|
||||
return YOLOResult([self.detections[idx]])
|
||||
return YOLOResult([])
|
||||
|
||||
def plot(self, **kwargs):
|
||||
return None
|
||||
|
||||
|
||||
class Boxes:
|
||||
"""模拟 YOLO boxes 对象"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = torch.tensor(xyxy_list, dtype=torch.float32)
|
||||
self.conf = torch.tensor(conf_list, dtype=torch.float32)
|
||||
self.cls = torch.tensor(cls_list, dtype=torch.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = torch.empty((0, 4))
|
||||
self.conf = torch.empty((0, 1))
|
||||
self.cls = torch.empty((0, 1), dtype=torch.int64)
|
||||
self.id = None
|
||||
|
||||
except ImportError:
|
||||
import numpy as np
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = np.array(xyxy_list, dtype=np.float32)
|
||||
self.conf = np.array(conf_list, dtype=np.float32)
|
||||
self.cls = np.array(cls_list, dtype=np.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = np.empty((0, 4), dtype=np.float32)
|
||||
self.conf = np.empty((0, 1), dtype=np.float32)
|
||||
self.cls = np.empty((0, 1), dtype=np.int64)
|
||||
self.id = None
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __iter__(self):
|
||||
for i in range(len(self.detections)):
|
||||
yield Box(self, i)
|
||||
|
||||
def cpu(self):
|
||||
return self
|
||||
|
||||
def numpy(self):
|
||||
if hasattr(self.xyxy, 'numpy'):
|
||||
return type('Boxes', (), {
|
||||
'xyxy': self.xyxy.numpy(),
|
||||
'conf': self.conf.numpy(),
|
||||
'cls': self.cls.numpy(),
|
||||
'id': self.id
|
||||
})()
|
||||
return self
|
||||
|
||||
|
||||
class Box:
|
||||
"""模拟单个检测框对象"""
|
||||
|
||||
def __init__(self, boxes: Boxes, index: int):
|
||||
self._boxes = boxes
|
||||
self._index = index
|
||||
|
||||
@property
|
||||
def xyxy(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
coords = self._boxes.xyxy[self._index]
|
||||
if isinstance(coords, torch.Tensor):
|
||||
return coords.unsqueeze(0)
|
||||
else:
|
||||
return np.array([coords])
|
||||
|
||||
@property
|
||||
def conf(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
conf_val = self._boxes.conf[self._index]
|
||||
if isinstance(conf_val, torch.Tensor):
|
||||
return conf_val.view(1)
|
||||
else:
|
||||
return np.array([conf_val])
|
||||
|
||||
@property
|
||||
def cls(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
cls_val = self._boxes.cls[self._index]
|
||||
if isinstance(cls_val, torch.Tensor):
|
||||
return cls_val.view(1)
|
||||
else:
|
||||
return np.array([cls_val])
|
||||
377
apps/server/models/smoking_yolo_adapter_v2_simple.py
Normal file
377
apps/server/models/smoking_yolo_adapter_v2_simple.py
Normal file
@@ -0,0 +1,377 @@
|
||||
"""
|
||||
YOLO 格式的抽烟检测模型适配器（V2 简化版）
|
||||
使用 Docker exec 在后台容器中执行检测
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import tempfile
|
||||
import logging
|
||||
import time
|
||||
import json
|
||||
from typing import List, Dict, Optional
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SmokingDetectionYOLO:
|
||||
"""
|
||||
模拟 YOLO 接口的抽烟检测模型（V2 简化版）
|
||||
使用 Docker exec 在后台容器中执行检测
|
||||
"""
|
||||
|
||||
_container_name = "smoking-detection-server"
|
||||
_initialized = False
|
||||
|
||||
def __init__(self, model_path=None):
|
||||
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
|
||||
|
||||
# YOLO 兼容属性
|
||||
self.names = {0: 'cigarette'}
|
||||
self.model = self
|
||||
|
||||
# 检查 Docker 并启动后台容器
|
||||
self._check_docker()
|
||||
if self.available:
|
||||
self._start_background_container()
|
||||
|
||||
logger.info(f"抽烟检测模型 V2 简化版初始化完成，Docker可用: {self.available}")
|
||||
|
||||
def _check_docker(self):
|
||||
"""检查 Docker 环境"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "info"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
if self.available:
|
||||
result = subprocess.run(
|
||||
["docker", "image", "inspect", self.docker_image],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
self.available = result.returncode == 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Docker 检查失败: {e}")
|
||||
self.available = False
|
||||
|
||||
def _start_background_container(self):
|
||||
"""启动后台容器"""
|
||||
try:
|
||||
# 检查容器是否已在运行
|
||||
result = subprocess.run(
|
||||
["docker", "ps", "-q", "-f", f"name={self._container_name}"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
|
||||
if result.stdout.strip():
|
||||
logger.info(f"后台容器已在运行: {self._container_name}")
|
||||
SmokingDetectionYOLO._initialized = True
|
||||
return
|
||||
|
||||
# 检查容器是否存在但已停止
|
||||
result = subprocess.run(
|
||||
["docker", "ps", "-aq", "-f", f"name={self._container_name}"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
|
||||
if result.stdout.strip():
|
||||
# 启动已存在的容器
|
||||
logger.info(f"启动已存在的容器: {self._container_name}")
|
||||
subprocess.run(
|
||||
["docker", "start", self._container_name],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
else:
|
||||
# 创建新容器（保持运行）
|
||||
logger.info(f"创建后台容器: {self._container_name}")
|
||||
subprocess.run(
|
||||
[
|
||||
"docker", "run", "-d",
|
||||
"--name", self._container_name,
|
||||
"-v", "/tmp:/workspace/input",
|
||||
"-v", "/Users/wwh/project/video-model/PaddlePaddle/PaddleDetection-release-2.9/smoking_server.py:/workspace/PaddleDetection/smoking_server.py",
|
||||
"-w", "/workspace/PaddleDetection",
|
||||
self.docker_image,
|
||||
"tail", "-f", "/dev/null"
|
||||
],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
|
||||
SmokingDetectionYOLO._initialized = True
|
||||
logger.info("后台容器启动成功")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"启动后台容器失败: {e}")
|
||||
SmokingDetectionYOLO._initialized = False
|
||||
|
||||
def __call__(self, source, conf=0.1, iou=0.45, verbose=False, stream=False):
|
||||
"""模拟 YOLO 模型的调用接口"""
|
||||
if not self.available:
|
||||
logger.error("Docker 不可用，无法运行检测")
|
||||
return [YOLOResult([])]
|
||||
|
||||
# 处理不同类型的输入
|
||||
if isinstance(source, str):
|
||||
image = cv2.imread(source)
|
||||
if image is None:
|
||||
logger.error(f"无法读取图片: {source}")
|
||||
return [YOLOResult([])]
|
||||
return self._detect_single(image, conf, verbose)
|
||||
|
||||
elif isinstance(source, np.ndarray):
|
||||
return self._detect_single(source, conf, verbose)
|
||||
|
||||
elif isinstance(source, list):
|
||||
results = []
|
||||
for img in source:
|
||||
if isinstance(img, str):
|
||||
img = cv2.imread(img)
|
||||
if img is not None:
|
||||
results.extend(self._detect_single(img, conf, verbose))
|
||||
return results
|
||||
|
||||
else:
|
||||
logger.error(f"不支持的输入类型: {type(source)}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def _detect_single(self, image: np.ndarray, conf: float, verbose: bool) -> List['YOLOResult']:
|
||||
"""检测单张图片（使用 docker exec）"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 创建临时文件
|
||||
input_filename = f"smoking_v2_{int(time.time()*1000)}.jpg"
|
||||
temp_input = f"/tmp/{input_filename}"
|
||||
|
||||
# 保存输入图片
|
||||
cv2.imwrite(temp_input, image)
|
||||
|
||||
if verbose:
|
||||
logger.info(f"正在检测: {temp_input}")
|
||||
|
||||
# 使用 docker exec 执行检测
|
||||
cmd = [
|
||||
"docker", "exec", "-i",
|
||||
self._container_name,
|
||||
"python", "smoking_server.py"
|
||||
]
|
||||
|
||||
# 发送请求
|
||||
request = {
|
||||
'image_path': f'/workspace/input/{input_filename}',
|
||||
'threshold': conf
|
||||
}
|
||||
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
input=json.dumps(request) + '\n',
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
if verbose:
|
||||
logger.info(f"检测完成，耗时: {elapsed:.2f}秒")
|
||||
|
||||
# 解析结果
|
||||
try:
|
||||
# 找到最后一行 JSON 输出
|
||||
lines = result.stdout.strip().split('\n')
|
||||
json_line = None
|
||||
for line in reversed(lines):
|
||||
line = line.strip()
|
||||
if line.startswith('{'):
|
||||
json_line = line
|
||||
break
|
||||
|
||||
if json_line:
|
||||
response = json.loads(json_line)
|
||||
if response.get('success'):
|
||||
detections = response.get('detections', [])
|
||||
else:
|
||||
logger.error(f"检测失败: {response.get('error')}")
|
||||
detections = []
|
||||
else:
|
||||
logger.error(f"无法解析输出: {result.stdout}")
|
||||
detections = []
|
||||
except Exception as e:
|
||||
logger.error(f"解析结果失败: {e}, stdout: {result.stdout}")
|
||||
detections = []
|
||||
|
||||
# 清理临时文件
|
||||
try:
|
||||
os.remove(temp_input)
|
||||
except:
|
||||
pass
|
||||
|
||||
return [YOLOResult(detections)]
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.error("检测超时")
|
||||
return [YOLOResult([])]
|
||||
except Exception as e:
|
||||
logger.error(f"检测失败: {e}")
|
||||
return [YOLOResult([])]
|
||||
|
||||
def predict(self, source, **kwargs):
|
||||
"""兼容 predict 方法"""
|
||||
return self.__call__(source, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def stop_server(cls):
|
||||
"""停止后台容器"""
|
||||
try:
|
||||
subprocess.run(
|
||||
["docker", "stop", cls._container_name],
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
)
|
||||
logger.info("后台容器已停止")
|
||||
except Exception as e:
|
||||
logger.error(f"停止后台容器失败: {e}")
|
||||
|
||||
|
||||
# YOLOResult, Boxes, Box 类（与之前相同）
|
||||
class YOLOResult:
|
||||
"""模拟 YOLO 检测结果对象"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
self.names = {0: 'cigarette'}
|
||||
self.boxes = Boxes(detections)
|
||||
self.probs = None
|
||||
self.keypoints = None
|
||||
self.obb = None
|
||||
self.speed = {'preprocess': 0, 'inference': 0, 'postprocess': 0}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if idx < len(self.detections):
|
||||
return YOLOResult([self.detections[idx]])
|
||||
return YOLOResult([])
|
||||
|
||||
def plot(self, **kwargs):
|
||||
return None
|
||||
|
||||
|
||||
class Boxes:
|
||||
"""模拟 YOLO boxes 对象"""
|
||||
|
||||
def __init__(self, detections: List[Dict]):
|
||||
self.detections = detections
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = torch.tensor(xyxy_list, dtype=torch.float32)
|
||||
self.conf = torch.tensor(conf_list, dtype=torch.float32)
|
||||
self.cls = torch.tensor(cls_list, dtype=torch.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = torch.empty((0, 4))
|
||||
self.conf = torch.empty((0, 1))
|
||||
self.cls = torch.empty((0, 1), dtype=torch.int64)
|
||||
self.id = None
|
||||
|
||||
except ImportError:
|
||||
import numpy as np
|
||||
|
||||
if detections:
|
||||
xyxy_list = [[d['bbox'][0], d['bbox'][1], d['bbox'][2], d['bbox'][3]] for d in detections]
|
||||
conf_list = [[d['confidence']] for d in detections]
|
||||
cls_list = [[d['class']] for d in detections]
|
||||
|
||||
self.xyxy = np.array(xyxy_list, dtype=np.float32)
|
||||
self.conf = np.array(conf_list, dtype=np.float32)
|
||||
self.cls = np.array(cls_list, dtype=np.int64)
|
||||
self.id = None
|
||||
else:
|
||||
self.xyxy = np.empty((0, 4), dtype=np.float32)
|
||||
self.conf = np.empty((0, 1), dtype=np.float32)
|
||||
self.cls = np.empty((0, 1), dtype=np.int64)
|
||||
self.id = None
|
||||
|
||||
def __len__(self):
|
||||
return len(self.detections)
|
||||
|
||||
def __iter__(self):
|
||||
for i in range(len(self.detections)):
|
||||
yield Box(self, i)
|
||||
|
||||
def cpu(self):
|
||||
return self
|
||||
|
||||
def numpy(self):
|
||||
if hasattr(self.xyxy, 'numpy'):
|
||||
return type('Boxes', (), {
|
||||
'xyxy': self.xyxy.numpy(),
|
||||
'conf': self.conf.numpy(),
|
||||
'cls': self.cls.numpy(),
|
||||
'id': self.id
|
||||
})()
|
||||
return self
|
||||
|
||||
|
||||
class Box:
|
||||
"""模拟单个检测框对象"""
|
||||
|
||||
def __init__(self, boxes: Boxes, index: int):
|
||||
self._boxes = boxes
|
||||
self._index = index
|
||||
|
||||
@property
|
||||
def xyxy(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
coords = self._boxes.xyxy[self._index]
|
||||
if isinstance(coords, torch.Tensor):
|
||||
return coords.unsqueeze(0)
|
||||
else:
|
||||
return np.array([coords])
|
||||
|
||||
@property
|
||||
def conf(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
conf_val = self._boxes.conf[self._index]
|
||||
if isinstance(conf_val, torch.Tensor):
|
||||
return conf_val.view(1)
|
||||
else:
|
||||
return np.array([conf_val])
|
||||
|
||||
@property
|
||||
def cls(self):
|
||||
import torch
|
||||
import numpy as np
|
||||
cls_val = self._boxes.cls[self._index]
|
||||
if isinstance(cls_val, torch.Tensor):
|
||||
return cls_val.view(1)
|
||||
else:
|
||||
return np.array([cls_val])
|
||||
Reference in New Issue
Block a user