火灾检测模型由基于YOLOv10的火灾烟雾检测模型改为复合模型[基于YOLOv8的火灾检测模型(单火焰检测)+YOLOv10-M,专用火灾烟雾模型]

This commit is contained in:
2026-06-10 14:18:43 +08:00
parent 30ea6eb0fb
commit 0e011dacfd
6 changed files with 206 additions and 21 deletions

View File

@@ -235,13 +235,15 @@ class CameraService:
confidence = config.get('confidence', 0.5)
iou = config.get('iou', 0.45)
draw = True
composite = config.get('composite', False)
processed_frame, result = await detection_service.detect_frame(
frame,
model_id=model_id,
confidence=confidence,
iou=iou,
draw=draw
draw=draw,
composite=composite
)
if result['success']:

View File

@@ -4,10 +4,11 @@ import numpy as np
import time
import uuid
import logging
import torch
from typing import Dict, List, Optional
from PIL import Image, ImageDraw, ImageFont
import torch
from .loitering_service import get_loitering_service
logger = logging.getLogger(__name__)
@@ -148,19 +149,39 @@ class DetectionService:
model_id: str,
confidence: float = 0.5,
iou: float = 0.45,
draw: bool = True
draw: bool = True,
composite: bool = False
) -> 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:
# 如果是火灾检测模型且启用了复合检测
if composite and model_id == 'fire_detection':
result_data = await self.detect_fire_composite(frame, confidence=confidence, iou=iou)
if result_data['success']:
detections = result_data['detections']
processing_time = time.time() - start_time
fps = 1.0 / processing_time if processing_time > 0 else 0
# 更新 stats 中的 fps 和处理时间
result_data['stats']['fps'] = round(fps, 2)
result_data['stats']['processing_time'] = round(processing_time, 3)
if draw:
frame = self.draw_detections(frame, detections, fps)
return frame, result_data
# 普通单模型检测
model = await self.model_service.load_model(model_id)
if not model:
return frame, {
'success': False,
'detections': [],
'stats': None
}
results = model(frame, conf=confidence, iou=iou, verbose=False)
detections = []
@@ -282,6 +303,115 @@ class DetectionService:
'stats': None
}
async def detect_fire_composite(
self,
image: np.ndarray,
confidence: float = 0.1,
iou: float = 0.45
) -> Dict:
"""
复合火灾检测:同时检测火焰和烟雾
"""
start_time = time.time()
try:
# 1. 检测火焰
fire_model = await self.model_service.load_model('fire_detection')
fire_results = fire_model(image, conf=confidence, iou=iou, verbose=False)
fire_detections = []
for result in fire_results:
for box in result.boxes:
try:
xyxy_values = box.xyxy.squeeze().tolist()
if len(xyxy_values) >= 4:
x1, y1, x2, y2 = float(xyxy_values[0]), float(xyxy_values[1]), float(xyxy_values[2]), float(xyxy_values[3])
else:
continue
conf = float(box.conf[0]) if hasattr(box.conf, '__getitem__') else float(box.conf)
cls = int(box.cls[0]) if hasattr(box.cls, '__getitem__') else int(box.cls)
class_name = result.names[cls]
if class_name == 'Fire':
fire_detections.append({
'class': 'Fire',
'label': '火焰',
'confidence': round(conf, 3),
'bbox': [int(x1), int(y1), int(x2), int(y2)],
'type': 'fire'
})
except Exception as e:
logger.error(f"火焰检测解析失败: {e}")
continue
# 2. 检测烟雾使用YOLOv10-M专用火灾烟雾模型只保留 Smoke 类别)
smoke_model = await self.model_service.load_model('smoke_detection')
smoke_results = smoke_model(image, conf=confidence, iou=iou, verbose=False)
smoke_detections = []
for result in smoke_results:
for box in result.boxes:
try:
xyxy_values = box.xyxy.squeeze().tolist()
if len(xyxy_values) >= 4:
x1, y1, x2, y2 = float(xyxy_values[0]), float(xyxy_values[1]), float(xyxy_values[2]), float(xyxy_values[3])
else:
continue
conf = float(box.conf[0]) if hasattr(box.conf, '__getitem__') else float(box.conf)
cls = int(box.cls[0]) if hasattr(box.cls, '__getitem__') else int(box.cls)
class_name = result.names[cls]
# 只保留 Smoke 类别(注意模型输出为大写 S
if class_name == 'Smoke':
smoke_detections.append({
'class': 'Smoke',
'label': '烟雾',
'confidence': round(conf, 3),
'bbox': [int(x1), int(y1), int(x2), int(y2)],
'type': 'smoke'
})
except Exception as e:
logger.error(f"烟雾检测解析失败: {e}")
continue
# 3. 合并所有检测
all_detections = fire_detections + smoke_detections
# 4. 判定是否疑似火灾(火焰或烟雾任一检测到即判定为是)
suspected_fire = len(fire_detections) > 0 or len(smoke_detections) > 0
processing_time = time.time() - start_time
avg_confidence = sum(d['confidence'] for d in all_detections) / len(all_detections) if all_detections else 0
return {
'success': True,
'message': '复合火灾检测完成',
'detections': all_detections,
'stats': {
'total_detections': len(all_detections),
'fire_count': len(fire_detections),
'smoke_count': len(smoke_detections),
'avg_confidence': round(avg_confidence, 3),
'suspected_fire': suspected_fire,
'suspected_fire_label': '' if suspected_fire else '',
'processing_time': round(processing_time, 3),
'model_used': 'fire_composite'
}
}
except Exception as e:
logger.error(f"复合火灾检测失败: {e}")
import traceback
logger.error(f"错误堆栈: {traceback.format_exc()}")
return {
'success': False,
'message': f'复合火灾检测失败: {str(e)}',
'detections': [],
'stats': None
}
def _apply_behavior_analysis(
self,
result_data: Dict,