dsh/dual_camera_system/safety_detector.py

"""
施工现场安全行为检测模块
使用 YOLO11 模型检测人员、安全帽、反光衣
判断是否存在违规行为（未戴安全帽、未穿反光衣）
"""

import cv2
import numpy as np
from typing import Optional, List, Tuple, Dict, Any
from dataclasses import dataclass
from enum import Enum


class SafetyViolationType(Enum):
    """安全违规类型"""
    NO_HELMET = "未戴安全帽"           # 未戴安全帽
    NO_SAFETY_VEST = "未穿反光衣"      # 未穿反光衣
    NO_BOTH = "反光衣和安全帽都没戴"   # 都没有


@dataclass
class SafetyDetection:
    """安全检测结果"""
    # 基础信息
    class_id: int                        # 类别ID
    class_name: str                      # 类别名称
    confidence: float                    # 置信度
    bbox: Tuple[int, int, int, int]      # 边界框 (x1, y1, x2, y2)
    center: Tuple[int, int]              # 中心点坐标
    track_id: Optional[int] = None       # 跟踪ID


@dataclass
class PersonSafetyStatus:
    """人员安全状态"""
    track_id: int                              # 跟踪ID
    person_bbox: Tuple[int, int, int, int]     # 人体边界框
    person_conf: float                         # 人体置信度
    has_helmet: bool = False                   # 是否戴安全帽
    helmet_conf: float = 0.0                   # 安全帽置信度
    has_safety_vest: bool = False              # 是否穿反光衣
    vest_conf: float = 0.0                     # 反光衣置信度
    is_violation: bool = False                 # 是否违规
    violation_types: List[SafetyViolationType] = None  # 违规类型列表
    
    def __post_init__(self):
        if self.violation_types is None:
            self.violation_types = []
    
    def check_violation(self) -> bool:
        """检查是否违规"""
        self.violation_types = []
        
        if not self.has_helmet and not self.has_safety_vest:
            self.violation_types.append(SafetyViolationType.NO_BOTH)
        elif not self.has_helmet:
            self.violation_types.append(SafetyViolationType.NO_HELMET)
        elif not self.has_safety_vest:
            self.violation_types.append(SafetyViolationType.NO_SAFETY_VEST)
        
        self.is_violation = len(self.violation_types) > 0
        return self.is_violation
    
    def get_violation_desc(self) -> str:
        """获取违规描述"""
        if not self.is_violation:
            return ""
        
        if SafetyViolationType.NO_BOTH in self.violation_types:
            return "反光衣和安全帽都没戴"
        elif SafetyViolationType.NO_HELMET in self.violation_types:
            return "未戴安全帽"
        elif SafetyViolationType.NO_SAFETY_VEST in self.violation_types:
            return "未穿反光衣"
        return ""


class SafetyDetector:
    """
    施工现场安全检测器
    使用 YOLO11 检测人员、安全帽、反光衣
    """
    
    # 类别映射 (根据 yolo11m_safety.pt 模型的训练标签)
    # 0: 安全帽, 3: 人, 4: 安全衣/反光衣
    CLASS_MAP = {
        0: '安全帽',
        3: '人',
        4: '反光衣'
    }
    
    # 反向映射
    CLASS_ID_MAP = {
        'helmet': 0,
        'person': 3,
        'safety_vest': 4
    }
    
    def __init__(self, model_path: str = None, use_gpu: bool = True, 
                 conf_threshold: float = 0.5, person_threshold: float = 0.8):
        """
        初始化安全检测器
        
        Args:
            model_path: 模型路径，默认使用 yolo11m_safety.pt
            use_gpu: 是否使用 GPU
            conf_threshold: 一般物品置信度阈值 (安全帽、反光衣)
            person_threshold: 人员检测置信度阈值
        """
        self.model = None
        self.model_path = model_path or '/home/wen/dsh/yolo/yolo11m_safety.pt'
        self.use_gpu = use_gpu
        self.device = 'cuda:0' if use_gpu else 'cpu'
        
        # 置信度阈值
        self.conf_threshold = conf_threshold
        self.person_threshold = person_threshold
        
        # 加载模型
        self._load_model()
    
    def _load_model(self):
        """加载 YOLO11 安全检测模型"""
        try:
            from ultralytics import YOLO
            
            self.model = YOLO(self.model_path)
            
            # 预热模型
            dummy = np.zeros((640, 640, 3), dtype=np.uint8)
            self.model(dummy, device=self.device, verbose=False)
            
            print(f"安全检测模型加载成功: {self.model_path} (device={self.device})")
        except ImportError:
            raise ImportError("未安装 ultralytics，请运行: pip install ultralytics")
        except Exception as e:
            raise RuntimeError(f"加载模型失败: {e}")
    
    def detect(self, frame: np.ndarray) -> List[SafetyDetection]:
        """
        检测画面中的安全相关对象
        
        Args:
            frame: 输入图像
            
        Returns:
            检测结果列表
        """
        if self.model is None or frame is None:
            return []
        
        results = []
        
        try:
            detections = self.model(frame, device=self.device, verbose=False)
            
            for det in detections:
                boxes = det.boxes
                if boxes is None:
                    continue
                
                for i in range(len(boxes)):
                    # 获取类别
                    cls_id = int(boxes.cls[i])
                    
                    # 只处理我们关心的类别
                    if cls_id not in self.CLASS_MAP:
                        continue
                    
                    cls_name = self.CLASS_MAP[cls_id]
                    conf = float(boxes.conf[i])
                    
                    # 根据类别设置不同的置信度阈值
                    threshold = self.person_threshold if cls_id == 3 else self.conf_threshold
                    if conf < threshold:
                        continue
                    
                    # 获取边界框
                    xyxy = boxes.xyxy[i].cpu().numpy()
                    x1, y1, x2, y2 = map(int, xyxy)
                    
                    # 过滤过小的检测框
                    width = x2 - x1
                    height = y2 - y1
                    if width < 10 or height < 10:
                        continue
                    
                    # 计算中心点
                    center_x = (x1 + x2) // 2
                    center_y = (y1 + y2) // 2
                    
                    detection = SafetyDetection(
                        class_id=cls_id,
                        class_name=cls_name,
                        confidence=conf,
                        bbox=(x1, y1, x2, y2),
                        center=(center_x, center_y)
                    )
                    results.append(detection)
                    
        except Exception as e:
            print(f"检测错误: {e}")
        
        return results
    
    def check_safety(self, frame: np.ndarray, 
                     detections: List[SafetyDetection] = None) -> List[PersonSafetyStatus]:
        """
        检查人员安全状态
        
        Args:
            frame: 输入图像
            detections: 检测结果，如果为 None 则自动检测
            
        Returns:
            人员安全状态列表
        """
        if detections is None:
            detections = self.detect(frame)
        
        # 分类检测结果
        persons = []
        helmets = []
        vests = []
        
        for det in detections:
            if det.class_id == 3:  # 人
                persons.append(det)
            elif det.class_id == 0:  # 安全帽
                helmets.append(det)
            elif det.class_id == 4:  # 反光衣
                vests.append(det)
        
        # 检查每个人员的安全状态
        results = []
        
        for person in persons:
            status = PersonSafetyStatus(
                track_id=person.track_id or 0,
                person_bbox=person.bbox,
                person_conf=person.confidence
            )
            
            px1, py1, px2, py2 = person.bbox
            
            # 检查是否戴安全帽
            # 安全帽应该在人体上方区域（头部附近）
            for helmet in helmets:
                hx1, hy1, hx2, hy2 = helmet.bbox
                # 检查安全帽是否在人体框内
                helmet_center_x = (hx1 + hx2) / 2
                helmet_center_y = (hy1 + hy2) / 2
                
                # 安全帽中心在人体框内，且在人体上半部分
                if (hx1 >= px1 and hx2 <= px2 and 
                    helmet_center_y >= py1 and 
                    helmet_center_y <= py1 + (py2 - py1) * 0.5):
                    status.has_helmet = True
                    status.helmet_conf = helmet.confidence
                    break
            
            # 检查是否穿反光衣
            # 反光衣应该与人体有重叠
            for vest in vests:
                vx1, vy1, vx2, vy2 = vest.bbox
                
                # 计算重叠区域
                overlap_x1 = max(px1, vx1)
                overlap_y1 = max(py1, vy1)
                overlap_x2 = min(px2, vx2)
                overlap_y2 = min(py2, vy2)
                
                # 如果有重叠
                if overlap_x1 < overlap_x2 and overlap_y1 < overlap_y2:
                    # 计算重叠面积占比
                    overlap_area = (overlap_x2 - overlap_x1) * (overlap_y2 - overlap_y1)
                    vest_area = (vx2 - vx1) * (vy2 - vy1)
                    overlap_ratio = overlap_area / vest_area if vest_area > 0 else 0
                    
                    # 重叠比例超过30%认为穿了反光衣
                    if overlap_ratio > 0.3:
                        status.has_safety_vest = True
                        status.vest_conf = vest.confidence
                        break
            
            # 检查是否违规
            status.check_violation()
            results.append(status)
        
        return results
    
    def detect_with_tracking(self, frame: np.ndarray, 
                             prev_tracks: Dict[int, Tuple[int, int]] = None,
                             max_disappeared: int = 30) -> Tuple[List[SafetyDetection], Dict[int, Tuple[int, int]]]:
        """
        带跟踪的检测
        
        Args:
            frame: 输入图像
            prev_tracks: 上一帧的跟踪状态 {track_id: center}
            max_disappeared: 最大消失帧数
            
        Returns:
            (检测结果列表, 当前跟踪状态)
        """
        detections = self.detect(frame)
        
        if prev_tracks is None:
            prev_tracks = {}
        
        # 简单的质心跟踪
        # 这里只对人体进行跟踪
        persons = [d for d in detections if d.class_id == 3]
        
        # 分配跟踪ID
        if len(persons) > 0:
            if len(prev_tracks) == 0:
                # 初始化
                for i, person in enumerate(persons):
                    person.track_id = i + 1
            else:
                # 匹配
                used_ids = set()
                for person in persons:
                    # 找最近的已跟踪对象
                    min_dist = float('inf')
                    best_id = None
                    
                    for track_id, center in prev_tracks.items():
                        if track_id in used_ids:
                            continue
                        
                        dist = np.sqrt((person.center[0] - center[0])**2 + 
                                      (person.center[1] - center[1])**2)
                        if dist < min_dist:
                            min_dist = dist
                            best_id = track_id
                    
                    if best_id is not None and min_dist < 100:  # 距离阈值
                        person.track_id = best_id
                        used_ids.add(best_id)
                    else:
                        # 新ID
                        new_id = max(prev_tracks.keys(), default=0) + 1
                        person.track_id = new_id
        
        # 更新跟踪状态
        new_tracks = {}
        for person in persons:
            if person.track_id is not None:
                new_tracks[person.track_id] = person.center
        
        return detections, new_tracks


def draw_safety_result(frame: np.ndarray, 
                       detections: List[SafetyDetection],
                       status_list: List[PersonSafetyStatus]) -> np.ndarray:
    """
    在图像上绘制安全检测结果
    
    Args:
        frame: 输入图像
        detections: 检测结果
        status_list: 人员安全状态
        
    Returns:
        绘制后的图像
    """
    result = frame.copy()
    
    # 绘制检测框
    for det in detections:
        x1, y1, x2, y2 = det.bbox
        
        # 根据类别选择颜色
        if det.class_id == 3:  # 人
            color = (0, 255, 0)  # 绿色
        elif det.class_id == 0:  # 安全帽
            color = (255, 165, 0)  # 橙色
        elif det.class_id == 4:  # 反光衣
            color = (0, 165, 255)  # 黄色
        else:
            color = (255, 255, 255)
        
        cv2.rectangle(result, (x1, y1), (x2, y2), color, 2)
        
        # 绘制标签
        label = f"{det.class_name}: {det.conf:.2f}"
        cv2.putText(result, label, (x1, y1 - 5), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
    
    # 绘制安全状态
    for status in status_list:
        x1, y1, x2, y2 = status.person_bbox
        
        if status.is_violation:
            # 违规 - 红色警告
            color = (0, 0, 255)
            text = status.get_violation_desc()
            cv2.rectangle(result, (x1, y1), (x2, y2), color, 3)
            cv2.putText(result, text, (x1, y2 + 20), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
        else:
            # 正常 - 显示安全标识
            color = (0, 255, 0)
            text = "安全装备齐全"
            cv2.putText(result, text, (x1, y2 + 20), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
    
    return result


class LLMSafetyDetector:
    """
    基于大模型的安全检测器
    结合 YOLO 检测和大模型判断
    """
    
    def __init__(self, yolo_model_path: str = None, 
                 llm_config: Dict[str, Any] = None,
                 use_gpu: bool = True,
                 use_llm: bool = True):
        """
        初始化检测器
        
        Args:
            yolo_model_path: YOLO 模型路径
            llm_config: 大模型配置
            use_gpu: 是否使用 GPU
            use_llm: 是否使用大模型判断
        """
        # YOLO 检测器
        self.yolo_detector = SafetyDetector(
            model_path=yolo_model_path,
            use_gpu=use_gpu
        )
        
        # 大模型分析器
        self.use_llm = use_llm
        self.llm_analyzer = None
        
        if use_llm:
            try:
                from llm_service import SafetyAnalyzer, NumberRecognizer
                self.llm_analyzer = SafetyAnalyzer(llm_config)
                self.number_recognizer = NumberRecognizer(llm_config)
                print("大模型安全分析器初始化成功")
            except ImportError:
                print("未找到 llm_service 模块，将使用规则判断")
                self.use_llm = False
            except Exception as e:
                print(f"大模型初始化失败: {e}，将使用规则判断")
                self.use_llm = False
    
    def detect(self, frame: np.ndarray) -> List[SafetyDetection]:
        """
        YOLO 检测
        
        Args:
            frame: 输入图像
            
        Returns:
            检测结果列表
        """
        return self.yolo_detector.detect(frame)
    
    def check_safety(self, frame: np.ndarray,
                     detections: List[SafetyDetection] = None,
                     use_llm: bool = None) -> List[PersonSafetyStatus]:
        """
        检查人员安全状态
        
        Args:
            frame: 输入图像
            detections: YOLO 检测结果
            use_llm: 是否使用大模型（覆盖默认设置）
            
        Returns:
            人员安全状态列表
        """
        # 先用 YOLO 检测
        if detections is None:
            detections = self.yolo_detector.detect(frame)
        
        # 规则判断
        rule_status_list = self.yolo_detector.check_safety(frame, detections)
        
        # 如果不使用大模型，直接返回规则判断结果
        should_use_llm = use_llm if use_llm is not None else self.use_llm
        if not should_use_llm or self.llm_analyzer is None:
            return rule_status_list
        
        # 使用大模型对每个人员进行判断
        llm_status_list = []
        
        for status in rule_status_list:
            # 裁剪人员区域
            x1, y1, x2, y2 = status.person_bbox
            margin = 10
            x1 = max(0, x1 - margin)
            y1 = max(0, y1 - margin)
            x2 = min(frame.shape[1], x2 + margin)
            y2 = min(frame.shape[0], y2 + margin)
            
            person_image = frame[y1:y2, x1:x2]
            
            # 调用大模型分析
            try:
                llm_result = self.llm_analyzer.check_person_safety(person_image)
                
                # 更新状态
                if llm_result.get('success', False):
                    status.has_helmet = llm_result.get('has_helmet', False)
                    status.has_safety_vest = llm_result.get('has_vest', False)
                    
                    # 重新检查违规
                    status.check_violation()
                    
                    # 如果大模型判断有违规，使用大模型的描述
                    if status.is_violation and llm_result.get('violation_desc'):
                        # 更新违规类型
                        desc = llm_result.get('violation_desc', '')
                        if '安全帽' in desc and '反光' in desc:
                            status.violation_types = [SafetyViolationType.NO_BOTH]
                        elif '安全帽' in desc:
                            status.violation_types = [SafetyViolationType.NO_HELMET]
                        elif '反光' in desc:
                            status.violation_types = [SafetyViolationType.NO_SAFETY_VEST]
            except Exception as e:
                print(f"大模型分析失败: {e}")
            
            llm_status_list.append(status)
        
        return llm_status_list
    
    def recognize_number(self, frame: np.ndarray,
                         person_bbox: Tuple[int, int, int, int]) -> Dict[str, Any]:
        """
        识别人员编号
        
        Args:
            frame: 输入图像
            person_bbox: 人员边界框
            
        Returns:
            编号识别结果
        """
        if self.number_recognizer is None:
            return {'number': None, 'success': False}
        
        # 裁剪人员区域
        x1, y1, x2, y2 = person_bbox
        person_image = frame[y1:y2, x1:x2]
        
        return self.number_recognizer.recognize_person_number(person_image)
    
    def detect_with_tracking(self, frame: np.ndarray,
                             prev_tracks: Dict[int, Tuple[int, int]] = None,
                             max_disappeared: int = 30) -> Tuple[List[SafetyDetection], Dict[int, Tuple[int, int]]]:
        """
        带跟踪的检测
        
        Args:
            frame: 输入图像
            prev_tracks: 上一帧的跟踪状态
            max_disappeared: 最大消失帧数
            
        Returns:
            (检测结果列表, 当前跟踪状态)
        """
        return self.yolo_detector.detect_with_tracking(frame, prev_tracks, max_disappeared)