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
| from ultralytics import YOLO import cv2 import numpy as np
class HandDetector: """基于YOLO26的手部检测器(用于手机使用检测)""" def __init__(self, model_size='n'): """ 初始化手部检测器 Args: model_size: 模型大小 ('n', 's', 'm', 'l') """ self.model = YOLO(f'yolo26{model_size}-hand.pt') self.classes = { 0: 'hand', 1: 'phone', 2: 'cup', 3: 'cigarette' } def detect(self, frame): """ 检测手部和手机 Args: frame: BGR图像 Returns: detections: 检测结果列表 """ results = self.model(frame, verbose=False) detections = [] for r in results: boxes = r.boxes for i in range(len(boxes)): box = boxes.xyxyn[i].cpu().numpy() cls = int(boxes.cls[i].cpu().numpy()) conf = float(boxes.conf[i].cpu().numpy()) detections.append({ 'class': self.classes.get(cls, 'unknown'), 'confidence': conf, 'bbox': box, 'class_id': cls }) return detections def detect_phone_usage(self, frame, driver_roi): """ 检测手机使用行为 Args: frame: 图像 driver_roi: 驾驶员区域 (x1, y1, x2, y2) Returns: is_phone_use: 是否使用手机 details: 详细信息 """ detections = self.detect(frame) hands_in_roi = [] phones_in_roi = [] for det in detections: cx = (det['bbox'][0] + det['bbox'][2]) / 2 cy = (det['bbox'][1] + det['bbox'][3]) / 2 if (driver_roi[0] <= cx <= driver_roi[2] and driver_roi[1] <= cy <= driver_roi[3]): if det['class'] == 'hand': hands_in_roi.append(det) elif det['class'] == 'phone': phones_in_roi.append(det) is_phone_use = len(hands_in_roi) > 0 and len(phones_in_roi) > 0 phone_position = None if is_phone_use: phone_bbox = phones_in_roi[0]['bbox'] phone_position = self.classify_phone_position(phone_bbox, frame.shape) return { 'is_phone_use': is_phone_use, 'phone_position': phone_position, 'confidence': max([d['confidence'] for d in detections]) if detections else 0 } def classify_phone_position(self, bbox, frame_shape): """分类手机位置""" cy = (bbox[1] + bbox[3]) / 2 h = frame_shape[0] if cy < 0.3 * h: return 'EAR' elif cy < 0.6 * h: return 'CHEST' else: return 'LAP'
if __name__ == "__main__": detector = HandDetector(model_size='n') frame = np.random.randint(0, 255, (720, 1280, 3), dtype=np.uint8) driver_roi = (0.3, 0.2, 0.7, 0.8) result = detector.detect_phone_usage(frame, driver_roi) print(f"手机使用: {result['is_phone_use']}, 位置: {result['phone_position']}")
|