""" 객체 감지 모듈 YOLOv8을 사용하여 자율주행에 필요한 다양한 객체를 감지합니다. - 신호등 (빨간불) - 정지 표지판 - 보행자 - 차량 (앞차 거리 유지) - 자전거/오토바이 """ import cv2 import numpy as np # ultralytics(YOLO)는 선택 의존성 — 없으면 객체 감지만 비활성화하고 나머지는 동작 try: from ultralytics import YOLO ULTRALYTICS_AVAILABLE = True except ImportError: YOLO = None ULTRALYTICS_AVAILABLE = False class ObjectDetector: """YOLO 기반 자율주행 객체 감지기""" # COCO 데이터셋 클래스 ID PERSON_ID = 0 BICYCLE_ID = 1 CAR_ID = 2 MOTORCYCLE_ID = 3 BUS_ID = 5 TRUCK_ID = 7 TRAFFIC_LIGHT_ID = 9 STOP_SIGN_ID = 11 def __init__(self, model_path="yolov8n.pt", confidence=0.5, traffic_light_model_path=None): """ Args: model_path (str): YOLO 모델 파일 경로 (일반 객체용) confidence (float): 감지 신뢰도 임계값 (0.0-1.0) traffic_light_model_path (str): 신호등 특화 모델 경로 (선택적) """ if not ULTRALYTICS_AVAILABLE: raise ImportError( "ultralytics가 설치되어 있지 않습니다. " "객체 감지(YOLO)를 사용하려면 'pip install ultralytics'로 설치하세요." ) self.model = YOLO(model_path) self.confidence = confidence # 신호등 특화 모델 (있으면 사용) self.traffic_light_model = None self.use_traffic_light_model = False if traffic_light_model_path: try: import os if os.path.exists(traffic_light_model_path): self.traffic_light_model = YOLO(traffic_light_model_path) self.use_traffic_light_model = True print(f"신호등 특화 모델 로드 성공: {traffic_light_model_path}") else: print(f"신호등 특화 모델을 찾을 수 없습니다: {traffic_light_model_path}") print("기본 HSV 색상 감지를 사용합니다.") except Exception as e: print(f"신호등 특화 모델 로드 실패: {e}") print("기본 HSV 색상 감지를 사용합니다.") # 빨간색 HSV 범위 (신호등용, 백업) # 빨간색은 HSV에서 두 구간으로 나뉨 (0-10, 170-180) self.red_lower1 = np.array([0, 50, 50], dtype=np.uint8) self.red_upper1 = np.array([10, 255, 255], dtype=np.uint8) self.red_lower2 = np.array([170, 50, 50], dtype=np.uint8) self.red_upper2 = np.array([180, 255, 255], dtype=np.uint8) # 초록색 HSV 범위 (신호등용) self.green_lower = np.array([40, 50, 50], dtype=np.uint8) self.green_upper = np.array([90, 255, 255], dtype=np.uint8) # 거리 임계값 (bbox 높이 비율 기반) self.vehicle_close_threshold = 0.4 # 차량이 화면의 40% 이상 차지하면 가까움 self.vehicle_near_threshold = 0.25 # 25% 이상이면 근접 self.person_threshold = 0.15 # 보행자/자전거 15% 이상이면 위험 def detect_all_objects(self, frame): """ 자율주행에 필요한 모든 객체 감지 (통합 버전) Args: frame (numpy.ndarray): 입력 프레임 (BGR) Returns: dict: { 'should_stop': bool, # 완전 정지 필요 'should_slow': bool, # 감속 필요 'speed_factor': float, # 속도 계수 (0.0-1.0, 1.0=정상속도) 'reason': str, # 정지/감속 이유 'detections': list, # 감지된 모든 객체 'red_light': bool, # 빨간 신호등 'stop_sign': bool, # 정지 표지판 'person': bool, # 보행자 'vehicle': bool, # 차량 'bicycle': bool # 자전거/오토바이 } """ if frame is None: return self._empty_result() height, width = frame.shape[:2] # 감지 플래그 초기화 red_light_detected = False green_light_detected = False stop_sign_detected = False person_detected = False vehicle_detected = False bicycle_detected = False detections = [] closest_vehicle_ratio = 0.0 # 신호등 특화 모델 사용 (우선) if self.use_traffic_light_model and self.traffic_light_model is not None: traffic_results = self.traffic_light_model(frame, conf=self.confidence, verbose=False) for result in traffic_results: boxes = result.boxes for box in boxes: class_id = int(box.cls[0]) confidence = float(box.conf[0]) x1, y1, x2, y2 = map(int, box.xyxy[0]) # 클래스 이름 가져오기 class_name = self.traffic_light_model.names[class_id].lower() # 빨간불 감지 (다양한 클래스명 처리) if 'red' in class_name: red_light_detected = True detections.append({ 'type': 'red_light', 'confidence': confidence, 'bbox': (x1, y1, x2, y2), 'priority': 1 }) # 초록불 감지 elif 'green' in class_name: green_light_detected = True detections.append({ 'type': 'green_light', 'confidence': confidence, 'bbox': (x1, y1, x2, y2), 'priority': 3 }) # 일반 객체 감지 (COCO 모델) results = self.model(frame, conf=self.confidence, verbose=False) # 감지된 객체 분석 for result in results: boxes = result.boxes for box in boxes: class_id = int(box.cls[0]) confidence = float(box.conf[0]) x1, y1, x2, y2 = map(int, box.xyxy[0]) bbox_height = y2 - y1 bbox_width = x2 - x1 height_ratio = bbox_height / height # 신호등 감지 (특화 모델 없을 때만 HSV 사용) if not self.use_traffic_light_model and class_id == self.TRAFFIC_LIGHT_ID: traffic_light_roi = frame[y1:y2, x1:x2] # 빨간불 체크 if self._is_red_light(traffic_light_roi): red_light_detected = True detections.append({ 'type': 'red_light', 'confidence': confidence, 'bbox': (x1, y1, x2, y2), 'priority': 1 }) # 초록불 체크 elif self._is_green_light(traffic_light_roi): green_light_detected = True detections.append({ 'type': 'green_light', 'confidence': confidence, 'bbox': (x1, y1, x2, y2), 'priority': 3 }) # 정지 표지판 감지 elif class_id == self.STOP_SIGN_ID: stop_sign_detected = True detections.append({ 'type': 'stop_sign', 'confidence': confidence, 'bbox': (x1, y1, x2, y2), 'priority': 1 }) # 보행자 감지 elif class_id == self.PERSON_ID: person_detected = True is_close = height_ratio > self.person_threshold detections.append({ 'type': 'person', 'confidence': confidence, 'bbox': (x1, y1, x2, y2), 'height_ratio': height_ratio, 'is_close': is_close, 'priority': 0 # 최우선 }) # 자전거/오토바이 감지 elif class_id in [self.BICYCLE_ID, self.MOTORCYCLE_ID]: bicycle_detected = True obj_type = 'bicycle' if class_id == self.BICYCLE_ID else 'motorcycle' is_close = height_ratio > self.person_threshold detections.append({ 'type': obj_type, 'confidence': confidence, 'bbox': (x1, y1, x2, y2), 'height_ratio': height_ratio, 'is_close': is_close, 'priority': 0 # 최우선 }) # 차량 감지 (car, truck, bus) elif class_id in [self.CAR_ID, self.TRUCK_ID, self.BUS_ID]: vehicle_detected = True if class_id == self.CAR_ID: obj_type = 'car' elif class_id == self.TRUCK_ID: obj_type = 'truck' else: obj_type = 'bus' # 가장 가까운 차량 추적 if height_ratio > closest_vehicle_ratio: closest_vehicle_ratio = height_ratio detections.append({ 'type': obj_type, 'confidence': confidence, 'bbox': (x1, y1, x2, y2), 'height_ratio': height_ratio, 'priority': 2 }) # 우선순위에 따른 제어 결정 should_stop = False should_slow = False speed_factor = 1.0 reason = "Normal" # 1. 최우선: 보행자/자전거 (즉시 정지) for det in detections: if det['type'] in ['person', 'bicycle', 'motorcycle'] and det.get('is_close', False): should_stop = True speed_factor = 0.0 reason = f"{det['type']} - STOP" break # 2. 신호등/정지 표지판 if not should_stop: if red_light_detected: should_stop = True speed_factor = 0.0 reason = "Red Light" elif stop_sign_detected: should_stop = True speed_factor = 0.0 reason = "Stop Sign" # 3. 앞차 거리 기반 속도 조절 if not should_stop and vehicle_detected: if closest_vehicle_ratio > self.vehicle_close_threshold: # 매우 가까움: 정지 should_stop = True speed_factor = 0.0 reason = "Vehicle too close" elif closest_vehicle_ratio > self.vehicle_near_threshold: # 근접: 감속 should_slow = True # 거리 비율에 따라 속도 조절 (0.25~0.4 구간을 0.5~0.8로 매핑) speed_factor = 0.5 + (self.vehicle_close_threshold - closest_vehicle_ratio) / \ (self.vehicle_close_threshold - self.vehicle_near_threshold) * 0.3 speed_factor = max(0.5, min(0.8, speed_factor)) reason = f"Vehicle near ({int(speed_factor*100)}%)" return { 'should_stop': should_stop, 'should_slow': should_slow, 'speed_factor': speed_factor, 'reason': reason, 'detections': detections, 'red_light': red_light_detected, 'green_light': green_light_detected, 'stop_sign': stop_sign_detected, 'person': person_detected, 'vehicle': vehicle_detected, 'bicycle': bicycle_detected } def detect_stop_objects(self, frame): """ 신호등(빨간색)과 정지 표지판 감지 (하위 호환성용) Args: frame (numpy.ndarray): 입력 프레임 (BGR) Returns: dict: { 'should_stop': bool, # 정지해야 하는지 여부 'red_light': bool, # 빨간 신호등 감지 여부 'stop_sign': bool, # 정지 표지판 감지 여부 'detections': list # 감지된 객체 정보 리스트 } """ # 새로운 detect_all_objects 호출하여 호환성 유지 result = self.detect_all_objects(frame) # 신호등/정지표지판만 필터링 filtered_detections = [d for d in result['detections'] if d['type'] in ['red_light', 'stop_sign']] return { 'should_stop': result['red_light'] or result['stop_sign'], 'red_light': result['red_light'], 'stop_sign': result['stop_sign'], 'detections': filtered_detections } def _empty_result(self): """빈 결과 반환""" return { 'should_stop': False, 'should_slow': False, 'speed_factor': 1.0, 'reason': "No frame", 'detections': [], 'red_light': False, 'green_light': False, 'stop_sign': False, 'person': False, 'vehicle': False, 'bicycle': False } def _is_red_light(self, roi): """ 신호등 영역이 빨간색인지 판별 Args: roi (numpy.ndarray): 신호등 영역 이미지 Returns: bool: 빨간색 신호등이면 True """ if roi is None or roi.size == 0: return False # HSV로 변환 hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV) # 빨간색 마스크 (두 구간) mask1 = cv2.inRange(hsv, self.red_lower1, self.red_upper1) mask2 = cv2.inRange(hsv, self.red_lower2, self.red_upper2) red_mask = cv2.bitwise_or(mask1, mask2) # 빨간색 픽셀 비율 계산 red_pixels = cv2.countNonZero(red_mask) total_pixels = roi.shape[0] * roi.shape[1] red_ratio = red_pixels / total_pixels if total_pixels > 0 else 0 # 빨간색 픽셀이 10% 이상이면 빨간 신호등으로 판정 return red_ratio > 0.1 def _is_green_light(self, roi): """ 신호등 영역이 초록색인지 판별 Args: roi (numpy.ndarray): 신호등 영역 이미지 Returns: bool: 초록색 신호등이면 True """ if roi is None or roi.size == 0: return False # HSV로 변환 hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV) # 초록색 마스크 green_mask = cv2.inRange(hsv, self.green_lower, self.green_upper) # 초록색 픽셀 비율 계산 green_pixels = cv2.countNonZero(green_mask) total_pixels = roi.shape[0] * roi.shape[1] green_ratio = green_pixels / total_pixels if total_pixels > 0 else 0 # 초록색 픽셀이 10% 이상이면 초록 신호등으로 판정 return green_ratio > 0.1 def draw_detections(self, frame, detection_result): """ 감지 결과를 프레임에 그리기 (모든 객체 타입 지원) Args: frame (numpy.ndarray): 입력 프레임 detection_result (dict): detect_all_objects() 또는 detect_stop_objects() 결과 Returns: numpy.ndarray: 시각화된 프레임 """ result_frame = frame.copy() # 객체 타입별 색상 및 레이블 정의 type_colors = { 'red_light': ((0, 0, 255), "RED LIGHT"), # 빨간색 'stop_sign': ((0, 165, 255), "STOP SIGN"), # 주황색 'person': ((255, 0, 255), "PERSON"), # 보라색 (위험) 'bicycle': ((255, 0, 200), "BICYCLE"), # 분홍색 'motorcycle': ((255, 0, 150), "MOTORCYCLE"), # 핑크 'car': ((255, 255, 0), "CAR"), # 시안색 'truck': ((200, 255, 0), "TRUCK"), # 연두색 'bus': ((150, 255, 0), "BUS") # 녹색 } for detection in detection_result['detections']: x1, y1, x2, y2 = detection['bbox'] obj_type = detection['type'] confidence = detection['confidence'] # 객체 타입에 따른 색상 및 레이블 if obj_type in type_colors: color, label_prefix = type_colors[obj_type] else: color = (128, 128, 128) # 회색 (기타) label_prefix = obj_type.upper() # 레이블 생성 label = f"{label_prefix} {confidence:.2f}" # 거리 정보 추가 (해당되는 경우) if 'height_ratio' in detection: distance_pct = int(detection['height_ratio'] * 100) label += f" ({distance_pct}%)" # 가까운 객체는 굵은 선으로 표시 is_close = detection.get('is_close', False) thickness = 3 if is_close else 2 # 바운딩 박스 그리기 cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, thickness) # 레이블 배경 label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2) cv2.rectangle(result_frame, (x1, y1 - label_size[1] - 10), (x1 + label_size[0], y1), color, -1) # 레이블 텍스트 cv2.putText(result_frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2) # 경고 메시지 표시 h, w = result_frame.shape[:2] if detection_result.get('should_stop', False): # 정지 경고 cv2.putText(result_frame, "!!! STOP !!!", (w // 2 - 80, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3) elif detection_result.get('should_slow', False): # 감속 경고 speed_pct = int(detection_result.get('speed_factor', 1.0) * 100) cv2.putText(result_frame, f"SLOW DOWN ({speed_pct}%)", (w // 2 - 120, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 165, 255), 3) # 정지/감속 이유 표시 if 'reason' in detection_result and detection_result['reason'] != "Normal": cv2.putText(result_frame, detection_result['reason'], (10, h - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) return result_frame # ===== 테스트 코드 ===== if __name__ == "__main__": import os model_path = os.path.join(os.path.dirname(__file__), "yolov8n.pt") # 모델 파일 확인 if not os.path.exists(model_path): print(f"YOLO 모델 파일을 찾을 수 없습니다: {model_path}") print("yolov8n.pt 파일을 python 폴더에 넣어주세요.") exit(1) print("=== 객체 감지기 테스트 ===") print(f"모델: {model_path}") detector = ObjectDetector(model_path) # 웹캠 테스트 cap = cv2.VideoCapture(0) if not cap.isOpened(): print("웹캠을 열 수 없습니다.") exit(1) print("\n웹캠에서 신호등과 정지 표지판을 감지합니다.") print("'q' 키를 눌러 종료\n") while True: ret, frame = cap.read() if not ret: break # 객체 감지 result = detector.detect_stop_objects(frame) # 결과 시각화 display_frame = detector.draw_detections(frame, result) # 상태 표시 status_text = [] if result['red_light']: status_text.append("RED LIGHT") if result['stop_sign']: status_text.append("STOP SIGN") if status_text: print(f"감지: {', '.join(status_text)}") cv2.imshow('Object Detection Test', display_frame) if cv2.waitKey(1) & 0xFF == ord('q'): break cap.release() cv2.destroyAllWindows() print("테스트 종료")