1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169
| import cv2 import numpy as np
class SeatbeltDetector: """ 安全带检测器 """ def __init__(self): self.belt_detector = self._load_model() self.classifier = BeltMisuseClassifier() def detect(self, frame): """ 检测安全带 Args: frame: 输入图像 Returns: belt_info: 安全带信息 """ edges = cv2.Canny(frame, 50, 150) lines = cv2.HoughLinesP(edges, 1, np.pi/180, 50, minLineLength=50, maxLineGap=10) belt_lines = self._filter_belt_lines(lines) buckle = self._detect_buckle(frame) person_roi = self._detect_person(frame) belt_status = self._analyze_belt_status(belt_lines, buckle, person_roi) return belt_status def _filter_belt_lines(self, lines): """ 过滤安全带线条 安全带特征: - 从B柱向下的斜线 - 宽度约30-50mm - 颜色通常为深色 """ if lines is None: return [] belt_lines = [] for line in lines: x1, y1, x2, y2 = line[0] angle = np.abs(np.arctan2(y2 - y1, x2 - x1)) if np.pi/6 < angle < np.pi/3: belt_lines.append((x1, y1, x2, y2, angle)) return belt_lines def _detect_buckle(self, frame): """检测安全带卡扣""" return {'position': (320, 400), 'confidence': 0.95} def _detect_person(self, frame): """检测乘员区域""" return {'box': (200, 100, 440, 480)} def _analyze_belt_status(self, belt_lines, buckle, person_roi): """ 分析安全带状态 """ status = { 'wearing': False, 'misuse_type': None, 'confidence': 0.0 } if len(belt_lines) >= 1: status['wearing'] = True misuse = self.classifier.classify(belt_lines, buckle, person_roi) status['misuse_type'] = misuse['type'] status['confidence'] = misuse['confidence'] return status
class BeltMisuseClassifier: """ 安全带误用分类器 """ def __init__(self): self.normal_angle_range = (np.pi/4, np.pi/3) def classify(self, belt_lines, buckle, person_roi): """ 分类误用类型 """ if not belt_lines: return {'type': 'not_wearing', 'confidence': 1.0} shoulder_slip = self._check_shoulder_slip(belt_lines, person_roi) if shoulder_slip['detected']: return {'type': 'shoulder_slip', 'confidence': shoulder_slip['confidence']} lap_position = self._check_lap_position(belt_lines, person_roi) if lap_position['incorrect']: return {'type': 'lap_position_error', 'confidence': lap_position['confidence']} behind_back = self._check_behind_back(belt_lines, person_roi) if behind_back['detected']: return {'type': 'behind_back', 'confidence': behind_back['confidence']} loose = self._check_loose(belt_lines) if loose['detected']: return {'type': 'loose', 'confidence': loose['confidence']} return {'type': 'normal', 'confidence': 0.9} def _check_shoulder_slip(self, belt_lines, person_roi): """ 检查肩带滑落 肩带滑落特征: - 安全带角度异常(接近水平) - 安全带上端位置过低 """ for line in belt_lines: x1, y1, x2, y2, angle = line if angle < np.pi/6: return {'detected': True, 'confidence': 0.8} return {'detected': False, 'confidence': 0.0} def _check_lap_position(self, belt_lines, person_roi): """检查腰带位置""" return {'incorrect': False, 'confidence': 0.0} def _check_behind_back(self, belt_lines, person_roi): """检查背后佩戴""" if len(belt_lines) < 2: return {'detected': True, 'confidence': 0.7} return {'detected': False, 'confidence': 0.0} def _check_loose(self, belt_lines): """检查松弛佩戴""" return {'detected': False, 'confidence': 0.0}
|