{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "%% [Cell 1] - 라이브러리 임포트 및 전역 설정 (Configuration & Global Variables)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# 기본 라이브러리 및 외부 패키지 임포트\n", "import os\n", "import sys\n", "import time\n", "import glob\n", "import cv2\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "from types import SimpleNamespace\n", "import torch\n", "import torch.nn as nn\n", "from tqdm import tqdm\n", "\n", "# MMDetection 관련 (mmdet/mmdet2의 API 등)\n", "import mmcv\n", "from mmengine.registry import init_default_scope\n", "from mmdet.apis import inference_detector, init_detector\n", "from mmdet.utils import register_all_modules\n", "\n", "# YOLOX Tracker (ByteTracker)\n", "from yolox.tracker.byte_tracker import BYTETracker\n", "from yolox.tracking_utils.timer import Timer\n", "\n", "# PIL for 이미지 처리 (폰트 사용)\n", "from PIL import ImageFont, ImageDraw, Image\n", "\n", "# ---------------------------\n", "# 전역 경로 및 저장 폴더 설정\n", "BASE_DIR = \"/DATA2/ltb\"\n", "VIDEO_DIR = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/it1_현대_추론용_폴더별분리/250509_신호수테스트\"\n", "SAVE_RESULT_BASE = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/it1_현대_추론용_폴더별분리/250509_신호수테스트\"\n", "SAVE_SUBFOLDER = \"exp_polygon_alarm_v5\"\n", "SAVE_RESULT = os.path.join(SAVE_RESULT_BASE, SAVE_SUBFOLDER)\n", "os.makedirs(SAVE_RESULT, exist_ok=True)\n", "print(\"✅ SAVE_RESULT =\", SAVE_RESULT)\n", "\n", "# ---------------------------\n", "# 전역 on/off 스위치 (각 단계 기능을 제어)\n", "# (아래 스위치가 off이면 해당 단계의 모든 처리는 스킵됩니다.)\n", "ENABLE_INFERENCE = True # 모델 인퍼런스 실행 여부\n", "ENABLE_VISUALIZATION = True # 모델 결과 시각화 여부\n", "ENABLE_POSITIVE_ENSEMBLE = True # (예: worker 3-way 병합)\n", "ENABLE_NEGATIVE_ENSEMBLE = True # (예: helmet / head 중복 제거)\n", "ENABLE_IOU_VIS_OVERRIDE = True # 지정 클래스의 IOU 조건에 따른 시각화 옵션 변경\n", "ENABLE_DETECT_ZONE_FILTER = False # Detect Zone 내부에 있는지 확인 (시각화 포함)\n", "ENABLE_DANGER_ZONE_WARNING = False # Danger Zone 내 존재 시 경고 문구 표시\n", "\n", "# ---------------------------\n", "# 전역 시각화 및 폰트 설정\n", "VIS_STYLE = {\n", " \"box_color\": (0, 255, 0),\n", " \"box_thickness\": 2,\n", " \"font_color\": (0, 255, 0),\n", " \"font_thickness\": 2,\n", " \"font_size\": 15, # 포인트\n", " \"font_path\": \"/usr/share/fonts/truetype/nanum/NanumGothic.ttf\",\n", "}\n", "# 텍스트 위치 offset (픽셀)\n", "TEXT_VERTICAL_OFFSET = 25\n", "\n", "# ---------------------------\n", "# 트래커 관련 설정\n", "TRACK_THRESH = 0.5\n", "TRACK_BUFFER = 120\n", "MATCH_THRESH = 0.6\n", "TRACKING_CORRECTION_DELAY = 3\n", "\n", "# ---------------------------\n", "# Detect & Danger Zone 좌표 (예시)\n", "# Detect Zone는 영상 전체 혹은 일부 (작업 구역)\n", "WORKZONE_POLYGON = [(0, 0), (1600, 0), (1600, 1080), (0, 1080)]\n", "# Danger Zone는 경고를 위한 구역 (예시)\n", "DANGER_POLYGON = [(600,800), (1400,800), (1600, 300), (800,300)]\n", "\n", "# Danger Zone 경고 문구 옵션\n", "DANGER_WARNING_TEXT = \"WARNING: Worker in restricted area!\"\n", "DANGER_WARNING_COLOR = (0, 0, 255) # 빨간색\n", "DANGER_WARNING_FONT_SCALE = 1.5\n", "DANGER_WARNING_FONT_THICKNESS = 2\n", "\n", "# ---------------------------\n", "# 앙상블 및 병합 관련 전역 변수 (예시)\n", "AGGREGATE_WORKER_COLOR = (0, 255, 0) # 초록색\n", "AGGREGATE_WORKER_3WAY_ON = True\n", "AGGREGATE_WORKER_WHHB_W = 0.8\n", "AGGREGATE_WORKER_FALL_W = 0.3\n", "AGGREGATE_WORKER_HARNESS_W = 0.5\n", "WORKER_MERGE_IOU_THRESHOLD = 0.3\n", "\n", "# 헬멧 관련 후처리 옵션\n", "HELMET_WEARING_IOU_THRESHOLD = 0.2 # 작업자와의 IOU 기준\n", "HELMET_ALERT_TEXT = \"안전모 착용\" # 헬멧이 없을 경우 텍스트로 나타남\n", "\n", "# ---------------------------\n", "# 모델별 설정 (MODELS_CONFIG)\n", "# 각 모델이 어떤 클래스(라벨)를 갖고, 어떤 임계값, 인퍼런스, 시각화를 사용할 것인지 설정\n", "MODELS_CONFIG = {\n", " \"whhb\": {\n", " \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/whh/it1_250326_1819_4class_rota_resize_0p5_1p0_2p0_5p0.py\",\n", " \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/whh/best_coco_bbox_mAP_epoch_12.pth\",\n", " \"class_names\": [\"worker\", \"helmet\", \"head\", \"background\"],\n", " \"do_inference\": False,\n", " \"vis_on\": False,\n", " \"class_params\": {\n", " \"default\": {\"detect_threshold\": 0.0, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.1},\n", " \"worker\": {\"detect_threshold\": 0.6, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.3},\n", " \"helmet\": {\"detect_threshold\": 0.3, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.3}\n", " },\n", " \"class_vis\": {\"worker\": True, \"helmet\": True, \"head\": True, \"background\": False}\n", " },\n", " \"con8\": {\n", " \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/con8/yolo8_large_960size_con_equip_8class.py\",\n", " \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/con8/best_coco_bbox_mAP_epoch_12.pth\",\n", " \"class_names\": [\"doger\", \"scissor_lift\", \"evacavator\", \"dump_truck\",\n", " \"mixer_truck\", \"crane_mobile\", \"cargo_truck\", \"forklift\"],\n", " \"do_inference\": False,\n", " \"vis_on\": False,\n", " \"class_params\": {\"default\": {\"detect_threshold\": 0.05, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.2}},\n", " \"class_vis\": {}\n", " },\n", " \"fall\": {\n", " \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/fall전용/yolov8_l_960_custom_fall.py\",\n", " \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/fall전용/best_coco_bbox_mAP_epoch_15.pth\",\n", " \"class_names\": [\"worker\", \"background\"],\n", " \"do_inference\": False,\n", " \"vis_on\": False,\n", " \"class_params\": {\n", " \"default\": {\"detect_threshold\": 0.0, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.1},\n", " \"worker\": {\"detect_threshold\": 0.4, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.3}\n", " },\n", " \"class_vis\": {}\n", " },\n", " \"harness\": {\n", " \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/dl_all_combined_signal_harness_6class_01/dl_all_combined_signal_harness_6class.py\",\n", " \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/dl_all_combined_signal_harness_6class_01/best_coco_bbox_mAP_epoch_13.pth\",\n", " \"class_names\": [\"worker\", \"신호수\", \"helmet\", \"안전벨트 착용\", \"mixer_truck\", \"excavator\"],\n", " \"do_inference\": True,\n", " \"vis_on\": True,\n", " \"class_params\": {\"default\": {\"detect_threshold\": 0.1, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.1},\n", " \"신호수\": {\"detect_threshold\": 0.0, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.1},\n", " \"안전벨트 착용\": {\"detect_threshold\": 0.3, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.3}\n", " \n", " },\n", " \"class_vis\": {\n", " # \"worker\": False, \"helmet\": False, \"mixer_truck\": False, \"excavator\": False, \n", " # \"안전벨트 착용\": False\n", " }\n", " },\n", "}\n", "\n", "# ---------------------------\n", "# 디바이스 설정 (GPU 사용 등)\n", "DEVICE = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n", "\n", "# MMDet 초기화: 경로 등록 및 모듈 등록 (mmyolo 경로 포함)\n", "mmyolo_dir = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo\"\n", "if mmyolo_dir not in sys.path:\n", " sys.path.append(mmyolo_dir)\n", "register_all_modules()\n", "init_default_scope('mmdet')\n", "\n", "print(\"===== [Cell 1] 전역 설정 완료 =====\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "%% [Cell 2] - 유틸리티 함수 모듈 (Utility Functions)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ---------------------------\n", "# 기본 기하학 계산 함수들\n", "def iou_calc(bbox_a, bbox_b):\n", " \"\"\"\n", " 두 bbox(bbox_a, bbox_b)의 IoU를 계산하는 함수.\n", " bbox 형식: [x1, y1, x2, y2]\n", " \"\"\"\n", " inter_x1 = max(bbox_a[0], bbox_b[0])\n", " inter_y1 = max(bbox_a[1], bbox_b[1])\n", " inter_x2 = min(bbox_a[2], bbox_b[2])\n", " inter_y2 = min(bbox_a[3], bbox_b[3])\n", " inter_w = max(0, inter_x2 - inter_x1)\n", " inter_h = max(0, inter_y2 - inter_y1)\n", " inter_area = inter_w * inter_h\n", " area_a = (bbox_a[2] - bbox_a[0]) * (bbox_a[3] - bbox_a[1])\n", " area_b = (bbox_b[2] - bbox_b[0]) * (bbox_b[3] - bbox_b[1])\n", " union_area = area_a + area_b - inter_area\n", " return inter_area / union_area if union_area > 0 else 0.0\n", "\n", "def is_point_in_polygon(point, polygon):\n", " \"\"\"\n", " 주어진 점(point)이 다각형(polygon) 내부에 있는지를 체크.\n", " point 형식: (x, y), polygon: [(x1,y1), (x2,y2), ...]\n", " \"\"\"\n", " pts = np.array(polygon, dtype=np.int32)\n", " return cv2.pointPolygonTest(pts, point, False) >= 0\n", "\n", "# ---------------------------\n", "# 색상 결정 함수\n", "# (각 클래스 및 모델별 고정 색상이나 팔레트 할당)\n", "FIXED_CLASS_COLORS = {\"worker\": (0, 255, 0), \"helmet\": (42,42,165), \"안전모 미착용\": (0,0,255), \"신호수\": (0, 255, 255)}\n", "COLOR_PALETTE = [\n", " (255,0,0), (211,0,148), (255,255,0),\n", " (0,255,255), (255,0,255), (100,100,100),\n", " (128,128,0), (128,0,128), (0,128,128),\n", " (128,128,255), (128,255,128), (255,128,128),\n", " (192,192,192), (0,0,128), (128,0,0),\n", " (128,255,255), (255,128,255)\n", "]\n", "_ASSIGNED_CLASS_COLORS = {} # 내부 저장용\n", "\n", "def get_color_for_class(class_name, model_key=None):\n", " \"\"\"\n", " 클래스별 색상을 결정.\n", " - 특정 모델과 클래스에 대해 고정 색상(FIXED_CLASS_COLORS) 적용\n", " - con8 모델은 CON8_OVERRIDE_COLOR 옵션에 따라 보라색 등 고정\n", " - 그 외는 팔레트에서 할당\n", " \"\"\"\n", " # worker나 작업자이면 VIS_STYLE의 box_color 우선\n", " if class_name in [\"worker\", \"작업자\"]:\n", " return VIS_STYLE[\"box_color\"]\n", " # whhb 혹은 harness에서의 고정 색상\n", " if model_key in [\"whhb\", \"harness\"] and class_name in FIXED_CLASS_COLORS:\n", " return FIXED_CLASS_COLORS[class_name]\n", " # con8 모델: 보라색 고정 (예시)\n", " if model_key == \"con8\":\n", " return (128, 0, 128)\n", " # 이미 할당된 색상 반환\n", " if class_name in _ASSIGNED_CLASS_COLORS:\n", " return _ASSIGNED_CLASS_COLORS[class_name]\n", " # 새 색상 할당: 팔레트에서 순차적으로\n", " for col in COLOR_PALETTE:\n", " if col not in _ASSIGNED_CLASS_COLORS.values():\n", " _ASSIGNED_CLASS_COLORS[class_name] = col\n", " return col\n", " # 팔레트가 소진되면 반복\n", " idx = len(_ASSIGNED_CLASS_COLORS) % len(COLOR_PALETTE)\n", " col = COLOR_PALETTE[idx]\n", " _ASSIGNED_CLASS_COLORS[class_name] = col\n", " return col\n", "\n", "# ---------------------------\n", "# 텍스트 시각화 관련 함수들\n", "def pt_to_px(pt, dpi=96):\n", " \"\"\"포인트(pt)를 픽셀(px)로 변환\"\"\"\n", " return int(pt * dpi / 72)\n", "\n", "def get_font_size_px(frame_height):\n", " \"\"\"현재 영상 프레임 높이에 맞게 폰트 크기를 절대값(px)로 변환\"\"\"\n", " base_px = pt_to_px(VIS_STYLE[\"font_size\"])\n", " BASE_FRAME_REF = 1080\n", " return max(10, int(base_px * frame_height / BASE_FRAME_REF))\n", "\n", "def draw_custom_text(img, text, position, frame_height, color):\n", " \"\"\"\n", " img에 텍스트를 그리는 함수.\n", " OpenCV freetype 모듈을 사용하거나, PIL을 fallback으로 사용.\n", " \"\"\"\n", " font_px = get_font_size_px(img.shape[0])\n", " thickness = VIS_STYLE[\"font_thickness\"]\n", " try:\n", " # freetype 사용 (opencv-contrib-python 필요)\n", " ft = cv2.freetype.createFreeType2()\n", " ft.loadFontData(fontFileName=VIS_STYLE[\"font_path\"], id=0)\n", " ft.putText(img, text, position, font_px, color, thickness, cv2.LINE_AA, True)\n", " except Exception as e:\n", " # PIL로 대체\n", " pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n", " draw = ImageDraw.Draw(pil_img)\n", " font = ImageFont.truetype(VIS_STYLE[\"font_path\"], font_px)\n", " draw.text(position, text, font=font, fill=color[::-1])\n", " img[:] = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "%% [Cell 3] - 모델 인퍼런스 및 각 모델별 탐지 딕셔너리 생성 (Detection Extraction)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_model_detections(frame, model_key, model_obj):\n", " \"\"\"\n", " 주어진 프레임에 대해 모델(model_key)의 인퍼런스를 실행하고,\n", " 각 클래스별 탐지 결과(딕셔너리)를 리스트로 반환.\n", " \n", " 반환 딕셔너리 형태:\n", " {\n", " \"bbox\": [x1,y1,x2,y2],\n", " \"score\": float,\n", " \"label\": class_name,\n", " \"model_key\": model_key\n", " }\n", " 단, 모델별 설정 (do_inference, vis_on, class_params) 조건에 맞추어\n", " detect_threshold 이상의 항목만 포함합니다.\n", " \"\"\"\n", " detections = []\n", " # 인퍼런스 실행 여부 및 vis_on 체크\n", " model_cfg = MODELS_CONFIG.get(model_key, {})\n", " if not (model_cfg.get(\"do_inference\", False) and model_cfg.get(\"vis_on\", False)):\n", " return detections\n", "\n", " result = inference_detector(model_obj, frame)\n", " # mmdet v2+의 경우 pred_instances를 사용\n", " if not hasattr(result, \"pred_instances\"):\n", " return detections\n", "\n", " bboxes = result.pred_instances.bboxes\n", " scores = result.pred_instances.scores\n", " labels = result.pred_instances.labels\n", " class_names = model_cfg[\"class_names\"]\n", " for i in range(len(bboxes)):\n", " cid = int(labels[i])\n", " # 잘못된 인덱스 방지\n", " if cid >= len(class_names):\n", " continue\n", " cname = class_names[cid]\n", " score = float(scores[i])\n", " bbox = bboxes[i].tolist()\n", " # 해당 클래스의 detect_threshold 가져오기 (없으면 0.0 사용)\n", " thr = model_cfg.get(\"class_params\", {}).get(cname, {}).get(\"detect_threshold\", \n", " model_cfg.get(\"class_params\", {}).get(\"default\", {}).get(\"detect_threshold\", 0.0))\n", " if score >= thr:\n", " detections.append({\n", " \"bbox\": bbox,\n", " \"score\": score,\n", " \"label\": cname,\n", " \"model_key\": model_key\n", " })\n", " return detections\n", "\n", "# 모델 로드 함수 (각 모델별 초기화)\n", "LOADED_MODELS = {}\n", "def load_all_models():\n", " loaded = {}\n", " for mk, cfg in MODELS_CONFIG.items():\n", " if cfg.get(\"do_inference\", False):\n", " print(f\"[모델 로드] {mk}\")\n", " model = init_detector(cfg[\"config\"], cfg[\"checkpoint\"], device=DEVICE)\n", " model.dataset_meta = {\"CLASSES\": cfg[\"class_names\"]}\n", " loaded[mk] = model\n", " else:\n", " print(f\"[모델 스킵] {mk}\")\n", " return loaded\n", "\n", "LOADED_MODELS = load_all_models()\n", "print(\"===== [Cell 3] 모델 인퍼런스 및 탐지 딕셔너리 준비 완료 =====\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "%% [Cell 4] - 탐지 결과 결합 및 필터링 (Detection Aggregation & Filtering)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def aggregate_detections(frame):\n", " \"\"\"\n", " 모든 모델(들)에 대해 인퍼런스 실행 후, 탐지 결과를 하나의 리스트로 합침.\n", " 추후 각 단계별(on/off) 필터링 및 앙상블 처리 전에\n", " 모델별로 분리되어 있었던 결과들을 통합하여 전달합니다.\n", " \"\"\"\n", " all_dets = []\n", " for model_key, model_obj in LOADED_MODELS.items():\n", " dets = get_model_detections(frame, model_key, model_obj)\n", " all_dets.extend(dets)\n", " return all_dets\n", "\n", "# ---------------------------\n", "# 각 단계 on/off 조건에 따라 결과를 필터링하는 함수들\n", "\n", "def filter_by_class_vis(detections):\n", " \"\"\"\n", " 각 탐지 딕셔너리에 대해 모델별 설정 class_vis 값을 참조하여\n", " 시각화 여부가 True인 항목만 남김.\n", " \"\"\"\n", " filtered = []\n", " for det in detections:\n", " model_key = det[\"model_key\"]\n", " cname = det[\"label\"]\n", " model_cfg = MODELS_CONFIG.get(model_key, {})\n", " # class_vis가 존재하면 해당 클래스가 True여야 함; 없으면 True로 간주\n", " if model_cfg.get(\"class_vis\", {}).get(cname, True):\n", " filtered.append(det)\n", " return filtered\n", "\n", "def apply_tracking_filter(detections, tracker, frame_shape, fps):\n", " \"\"\"\n", " 간단한 트래킹 처리 적용:\n", " - ByteTracker를 이용하여 detections에 트랙 ID를 부여하고,\n", " - 일정 안정 시간을 충족한 항목만 필터링 (예: required_pass_ratio)\n", " 이 함수는 매우 단순화된 예시입니다.\n", " \"\"\"\n", " # 각 detection의 bbox를 numpy array로 변환: [x1, y1, x2, y2, score]\n", " boxes = []\n", " for det in detections:\n", " x1, y1, x2, y2 = det[\"bbox\"]\n", " boxes.append([x1, y1, x2, y2, det[\"score\"]])\n", " boxes_arr = np.array(boxes, dtype=np.float32) if boxes else np.zeros((0,5), dtype=np.float32)\n", " online_tracks = tracker.update(boxes_arr, frame_shape)\n", " filtered = []\n", " for t in online_tracks:\n", " # 단순히 트랙 ID가 부여된 것만 사용 (실제 조건은 추가 가능)\n", " tid = t.track_id\n", " x1, y1, w, h = t.tlwh\n", " filtered.append({\n", " \"bbox\": [x1, y1, x1+w, y1+h],\n", " \"score\": t.score,\n", " \"label\": None, # 추후 할당 (가장 근접한 detection에서 결정)\n", " \"track_id\": tid,\n", " \"model_key\": \"tracked\" # 추적 결과임을 명시\n", " })\n", " return filtered\n", "\n", "def assign_class_to_tracks(tracks, detections):\n", " \"\"\"\n", " 단순하게 각 트랙에 대해, 겹치는 detection 중 score가 가장 높은 것의 클래스, 모델정보를 할당.\n", " \"\"\"\n", " for tr in tracks:\n", " best_iou = 0\n", " best_det = None\n", " for det in detections:\n", " iou_val = iou_calc(tr[\"bbox\"], det[\"bbox\"])\n", " if iou_val > best_iou:\n", " best_iou = iou_val\n", " best_det = det\n", " if best_det and best_iou > 0.5:\n", " tr[\"label\"] = best_det[\"label\"]\n", " tr[\"model_key\"] = best_det[\"model_key\"]\n", " return tracks\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "%% [Cell 5] - 앙상블 및 병합 함수 (Ensemble & Merge Processing)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def apply_positive_ensemble(detections, iou_thresh=0.3):\n", " \"\"\"\n", " [Positive Ensemble]\n", " 예를 들어, worker 클래스의 경우\n", " whhb, fall, harness 모델 각각에서 나온 worker 탐지들 중\n", " 같은 박스 (IOU ≥ iou_thresh)인 경우, score가 높은 항목만 남기고\n", " 라벨은 \"작업자\"로 통일하며, 색상은 전역 AGGREGATE_WORKER_COLOR로 지정.\n", " 이 함수는 on/off 스위치(ENABLE_POSITIVE_ENSEMBLE)가 on일 때 호출됩니다.\n", " \"\"\"\n", " used = [False] * len(detections)\n", " final = []\n", " for i in range(len(detections)):\n", " if used[i]:\n", " continue\n", " det_i = detections[i]\n", " # 대상 클래스 검사 (예시: worker, 혹은 모델과 상관없이 \"worker\")\n", " if det_i[\"label\"] in [\"worker\"]:\n", " group = [det_i]\n", " used[i] = True\n", " # 모델별 가중치 부여 (예시)\n", " weight_i = AGGREGATE_WORKER_WHHB_W if det_i[\"model_key\"]==\"whhb\" else \\\n", " (AGGREGATE_WORKER_FALL_W if det_i[\"model_key\"]==\"fall\" else AGGREGATE_WORKER_HARNESS_W)\n", " weights = [weight_i]\n", " for j in range(i+1, len(detections)):\n", " if used[j]:\n", " continue\n", " det_j = detections[j]\n", " if det_j[\"label\"] in [\"worker\"]:\n", " if iou_calc(det_i[\"bbox\"], det_j[\"bbox\"]) >= iou_thresh:\n", " group.append(det_j)\n", " used[j] = True\n", " weight_j = AGGREGATE_WORKER_WHHB_W if det_j[\"model_key\"]==\"whhb\" else \\\n", " (AGGREGATE_WORKER_FALL_W if det_j[\"model_key\"]==\"fall\" else AGGREGATE_WORKER_HARNESS_W)\n", " weights.append(weight_j)\n", " # 그룹 내 가중 평균 계산\n", " boxes = np.array([g[\"bbox\"] for g in group], dtype=np.float32)\n", " agg_bbox = np.average(boxes, axis=0, weights=weights).tolist()\n", " scores = [g[\"score\"] for g in group]\n", " agg_score = np.average(scores, weights=weights)\n", " # 새로운 앙상블 결과 딕셔너리: 라벨은 \"작업자\", 모델은 \"aggregated_worker\"\n", " aggregated = {\n", " \"bbox\": agg_bbox,\n", " \"score\": agg_score,\n", " \"label\": \"작업자\",\n", " \"model_key\": \"aggregated_worker\",\n", " \"override_color\": AGGREGATE_WORKER_COLOR\n", " }\n", " final.append(aggregated)\n", " else:\n", " # 대상 클래스가 아니면 그대로 남김.\n", " final.append(det_i)\n", " return final\n", "\n", "def apply_negative_ensemble(detections, iou_thresh=0.2):\n", " \"\"\"\n", " [Negative Ensemble]\n", " 예를 들어, whhb 모델의 \"helmet\"과 \"head\" (또는 유사 클래스) 간에 겹침이 발생한 경우,\n", " 각 탐지별로 가중치(score에 가중치)를 곱해 비교하고, 더 높은 신뢰도의 항목만 남깁니다.\n", " 이 함수는 ENABLE_NEGATIVE_ENSEMBLE 스위치가 on일 때 사용.\n", " \"\"\"\n", " used = [False] * len(detections)\n", " final = []\n", " for i in range(len(detections)):\n", " if used[i]:\n", " continue\n", " det_i = detections[i]\n", " # negative 대상: 예시로 \"helmet\"과 \"head\"\n", " if det_i[\"label\"] in [\"helmet\", \"head\"]:\n", " for j in range(i+1, len(detections)):\n", " if used[j]:\n", " continue\n", " det_j = detections[j]\n", " if {det_i[\"label\"], det_j[\"label\"]} == {\"helmet\", \"head\"}:\n", " # 가중치 조건: 예를 들어 helmet에는 0.8, head에는 0.3 등\n", " weight_i = 0.8 if det_i[\"label\"]==\"helmet\" else 0.3\n", " weight_j = 0.8 if det_j[\"label\"]==\"helmet\" else 0.3\n", " if det_i[\"score\"] * weight_i < det_j[\"score\"] * weight_j:\n", " used[i] = True\n", " break\n", " else:\n", " used[j] = True\n", " if not used[i]:\n", " final.append(det_i)\n", " else:\n", " final.append(det_i)\n", " return final\n", "\n", "# ---------------------------\n", "# 구역(Zone) 필터: Detect Zone와 Danger Zone 적용 함수\n", "def apply_detect_zone_filter(detections, zone_polygon, target_labels):\n", " \"\"\"\n", " target_labels에 해당하는 탐지들만, bbox 중심이 zone_polygon 내에 있을 때만 남김.\n", " \"\"\"\n", " filtered = []\n", " for det in detections:\n", " if det[\"label\"] in target_labels:\n", " x1, y1, x2, y2 = det[\"bbox\"]\n", " cx, cy = (x1+x2)/2, (y1+y2)/2\n", " if is_point_in_polygon((cx, cy), zone_polygon):\n", " filtered.append(det)\n", " else:\n", " filtered.append(det)\n", " return filtered\n", "\n", "def apply_danger_zone_warning(detections, danger_polygon, target_labels):\n", " \"\"\"\n", " danger_polygon 내부에 target_labels의 탐지가 하나라도 있으면,\n", " 각 탐지 딕셔너리에 \"alarm\": True 플래그를 추가.\n", " \"\"\"\n", " danger_triggered = False\n", " for det in detections:\n", " if det[\"label\"] in target_labels:\n", " x1, y1, x2, y2 = det[\"bbox\"]\n", " cx, cy = (x1+x2)/2, (y1+y2)/2\n", " if is_point_in_polygon((cx, cy), danger_polygon):\n", " det[\"alarm\"] = True\n", " danger_triggered = True\n", " return detections, danger_triggered\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "%% [Cell 6] - 최종 시각화를 위한 draw 함수 (Final Drawing Function)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def draw_final_detections(frame, detections):\n", " \"\"\"\n", " 최종으로 화면에 박스 및 텍스트를 그리는 함수.\n", " 전달된 detections 리스트에 대해 각 조건(예: IOU에 따른 색상 오버라이드, 트랙 정보 등)\n", " 를 반영하여 draw 후 최종 프레임 이미지를 반환합니다.\n", " \"\"\"\n", " out_img = frame.copy()\n", " # 먼저, Danger Zone이 활성화된 경우, 경고 문구를 좌측 상단에 출력\n", " _, danger_triggered = apply_danger_zone_warning(detections, DANGER_POLYGON, target_labels=[\"worker\", \"signalman\", \"작업자\"])\n", " if danger_triggered and ENABLE_DANGER_ZONE_WARNING:\n", " cv2.putText(out_img, DANGER_WARNING_TEXT, (50,50), cv2.FONT_HERSHEY_SIMPLEX, \n", " DANGER_WARNING_FONT_SCALE, DANGER_WARNING_COLOR, DANGER_WARNING_FONT_THICKNESS, cv2.LINE_AA)\n", " \n", " # 각 탐지에 대해 박스 및 텍스트 그리기\n", " for det in detections:\n", " box = det[\"bbox\"]\n", " x1, y1, x2, y2 = map(int, box)\n", " model_key = det.get(\"model_key\", \"\")\n", " # override_color가 있으면 사용, 아니면 기본 get_color_for_class 사용\n", " color = det.get(\"override_color\", get_color_for_class(det[\"label\"], model_key))\n", " cv2.rectangle(out_img, (x1, y1), (x2, y2), color, VIS_STYLE[\"box_thickness\"])\n", " # 텍스트 (클래스명 및 confidence)\n", " txt = f\"{det['label']} {det['score']:.2f}\"\n", " draw_custom_text(out_img, txt, (x1, max(y1 - TEXT_VERTICAL_OFFSET, 0)), out_img.shape[0], color)\n", " return out_img\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "%% [Cell 7] - 메인 프로세싱 함수 (Pipeline Process)\n", "python\n", "복사\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def process_frame(frame):\n", " \"\"\"\n", " 한 프레임에 대해 전체 파이프라인 실행:\n", " 1. 각 모델별 인퍼런스 → 탐지 딕셔너리 생성 (aggregate_detections)\n", " 2. class_vis 필터 적용\n", " 3. (선택) 트래킹 적용 → 간단 트래킹은 별도 처리 (여기서는 생략 가능)\n", " 4. on/off 스위치에 따라 앙상블 (positive, negative) 적용\n", " 5. on/off 스위치에 따라 Detect Zone 필터 적용\n", " 6. 최종 draw 함수 호출 → 최종 프레임 리턴\n", " \"\"\"\n", " # 1. 모델별 탐지 결과 합치기\n", " dets = aggregate_detections(frame)\n", " # 2. class_vis 필터 적용 (각 모델별 설정에 의해 숨김이 있다면 제거)\n", " dets = filter_by_class_vis(dets)\n", " \n", " # 3. (트래킹 적용 가능하면 여기서 적용 -- 본 예시는 단순 합산)\n", " # tracker = BYTETracker(SimpleNamespace(track_thresh=TRACK_THRESH, track_buffer=TRACK_BUFFER, match_thresh=MATCH_THRESH))\n", " # dets_tracked = apply_tracking_filter(dets, tracker, frame.shape[:2], fps=30)\n", " # dets = assign_class_to_tracks(dets_tracked, dets) # 단순한 트랙 할당 예시\n", " \n", " # 4. positive 앙상블 (예: worker 클래스 병합)\n", " if ENABLE_POSITIVE_ENSEMBLE:\n", " dets = apply_positive_ensemble(dets, iou_thresh=WORKER_MERGE_IOU_THRESHOLD)\n", " \n", " # 5. negative 앙상블 (예: helmet / head 중복)\n", " if ENABLE_NEGATIVE_ENSEMBLE:\n", " dets = apply_negative_ensemble(dets, iou_thresh=0.2)\n", " \n", " # 6. IOU 기반 시각화 옵션 오버라이드 (예: worker, helmet 등)\n", " # (여기서는 조건에 따라 override_color나 텍스트를 수정하는 추가 로직을 넣을 수 있음)\n", " if ENABLE_IOU_VIS_OVERRIDE:\n", " # 예시: worker와 helmet 간 IOU가 높으면 helmet의 텍스트를 빈 문자열로 변경\n", " for det in dets:\n", " if det[\"label\"] == \"helmet\":\n", " # worker와의 IOU 체크\n", " for other in dets:\n", " if other[\"label\"] in [\"worker\", \"작업자\"]:\n", " if iou_calc(det[\"bbox\"], other[\"bbox\"]) >= HELMET_WEARING_IOU_THRESHOLD:\n", " det[\"helmet_text\"] = \"\"\n", " det[\"override_color\"] = (0, 255, 0) # 초록색\n", " break\n", " \n", " # 7. Detect Zone 필터 적용 (지정 클래스에 한해서)\n", " if ENABLE_DETECT_ZONE_FILTER:\n", " # 예시로 worker, signalman, 작업자만 detect zone 반영 (zone 내에 있으면 통과)\n", " dets = apply_detect_zone_filter(dets, WORKZONE_POLYGON, target_labels=[\"worker\", \"작업자\", \"signalman\"])\n", " \n", " # 최종 결과를 시각화 전 최종 목록(final_vis)이라 가정하고 draw 호출\n", " final_frame = draw_final_detections(frame, dets)\n", " return final_frame\n", "\n", "# 예시: 단일 프레임 디버깅\n", "if __name__ == \"__main__\":\n", " test_video = glob.glob(os.path.join(VIDEO_DIR, \"*.mp4\"))\n", " if test_video:\n", " cap = cv2.VideoCapture(test_video[0])\n", " ret, test_frame = cap.read()\n", " cap.release()\n", " if ret:\n", " processed = process_frame(test_frame)\n", " plt.figure(figsize=(8,6))\n", " plt.imshow(cv2.cvtColor(processed, cv2.COLOR_BGR2RGB))\n", " plt.title(\"최종 시각화 결과 (단일 프레임)\")\n", " plt.axis(\"off\")\n", " plt.show()\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "%% [Cell 8] - 메인 실행부 (Main Video Processing Loop)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def process_video(video_path, resize_factor=1.0):\n", " \"\"\"\n", " 전체 비디오에 대해 각 프레임을 처리하고, 최종 결과 영상을 저장하는 함수.\n", " 각 프레임은 process_frame() 함수의 결과를 writer에 기록합니다.\n", " \"\"\"\n", " base_name = os.path.basename(video_path)\n", " name, ext = os.path.splitext(base_name)\n", " video_save_folder = os.path.join(SAVE_RESULT, name)\n", " os.makedirs(video_save_folder, exist_ok=True)\n", " output_video_path = os.path.join(video_save_folder, f\"detected_{name}.mp4\")\n", " \n", " cap = cv2.VideoCapture(video_path)\n", " if not cap.isOpened():\n", " print(\"[오류] 영상 열기 실패 =>\", video_path)\n", " return\n", " fps = cap.get(cv2.CAP_PROP_FPS)\n", " if fps <= 0:\n", " fps = 30\n", " w_org = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n", " h_org = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n", " \n", " fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n", " writer = cv2.VideoWriter(output_video_path, fourcc, fps, (int(w_org*resize_factor), int(h_org*resize_factor)))\n", " if not writer.isOpened():\n", " print(\"[경고] VideoWriter 생성 실패 =>\", output_video_path)\n", " writer = None\n", " \n", " frame_idx = 0\n", " total_time = time.time()\n", " while True:\n", " ret, frame = cap.read()\n", " if not ret:\n", " break\n", " proc_frame = process_frame(frame)\n", " if resize_factor != 1.0:\n", " proc_frame = cv2.resize(proc_frame, (int(w_org*resize_factor), int(h_org*resize_factor)))\n", " if writer:\n", " writer.write(proc_frame)\n", " \n", " if frame_idx % 100 == 0:\n", " debug_path = os.path.join(video_save_folder, f\"{name}_frame_{frame_idx:04d}.jpg\")\n", " cv2.imwrite(debug_path, proc_frame)\n", " print(f\"[DEBUG] Frame {frame_idx} 저장: {debug_path}\")\n", " frame_idx += 1\n", " cap.release()\n", " if writer:\n", " writer.release()\n", " print(f\"[INFO] 영상 처리 완료: {output_video_path}, 총 처리 시간: {time.time()-total_time:.1f}초\")\n", "\n", "def main():\n", " video_files = [f for f in os.listdir(VIDEO_DIR) if os.path.splitext(f)[1].lower() in [\".mp4\", \".avi\", \".mov\", \".mkv\"]]\n", " for vf in sorted(video_files):\n", " vp = os.path.join(VIDEO_DIR, vf)\n", " print(f\"처리 시작: {vp}\")\n", " process_video(vp, resize_factor=1.0)\n", " print(\"===== 전체 영상 처리 완료 =====\")\n", "\n", "if __name__ == \"__main__\":\n", " main()\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "%% [Cell 9] - 추가 디버깅/유틸 함수 (예: 임의 이미지 디버깅 등)\n", "python\n", "복사\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def debug_arbitrary_image(img_path):\n", " \"\"\"\n", " 지정한 이미지 파일에 대해 각 모델별 인퍼런스를 실행한 결과를\n", " 최종 시각화한 이미지를 저장합니다.\n", " \"\"\"\n", " if not os.path.exists(img_path):\n", " print(\"[오류] 이미지 파일이 존재하지 않습니다:\", img_path)\n", " return\n", " img = cv2.imread(img_path)\n", " if img is None:\n", " print(\"[오류] 이미지 읽기 실패:\", img_path)\n", " return\n", " # 단일 모델 예시: whhb만 사용\n", " if \"whhb\" not in LOADED_MODELS:\n", " print(\"[스킵] whhb 모델 없음\")\n", " return\n", " dets = get_model_detections(img, \"whhb\", LOADED_MODELS[\"whhb\"])\n", " vis_img = draw_final_detections(img, dets)\n", " debug_folder = os.path.join(SAVE_RESULT, \"debug_arbitrary\")\n", " os.makedirs(debug_folder, exist_ok=True)\n", " debug_path = os.path.join(debug_folder, \"debug_arbitrary_image.jpg\")\n", " cv2.imwrite(debug_path, vis_img)\n", " print(\"[디버그] 임의 이미지 디텍션 결과 저장:\", debug_path)\n", "\n", "# 테스트할 이미지 경로 지정 후 호출\n", "# debug_arbitrary_image(\"/path/to/test_image.jpg\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "전체 코드 요약\n", "Cell 1: 전역 설정 – 저장 경로, on/off 스위치, 폰트 및 시각화 기본값, 각 모델별 설정(MODELS_CONFIG) 등이 정의되어 있다.\n", "\n", "Cell 2: 유틸 함수 – IoU 계산, 점 영역 내 포함 여부, 색상 결정, 텍스트 그리기 함수 등을 정의하여 공용으로 사용한다.\n", "\n", "Cell 3: 모델 인퍼런스 – 각 모델에 대해 인퍼런스를 실행하고 threshold 조건에 맞는 탐지 딕셔너리를 생성한다.\n", "\n", "Cell 4: 탐지 결과 결합 및 기본 필터 – 각 모델에서 나온 탐지 결과를 하나로 합치고 class_vis 조건 등을 적용한다.\n", "\n", "Cell 5: 앙상블 및 병합 – positive(예, worker 3-way 병합)와 negative(예, helmet/head 중복) 앙상블 함수 및 zone 필터 함수를 정의한다.\n", "\n", "Cell 6: 최종 시각화 – 조건들을 모두 통과한 탐지 딕셔너리로 프레임에 박스와 텍스트를 그린다.\n", "\n", "Cell 7: 전체 파이프라인 – 한 프레임마다 위의 단계를 순차적으로 적용하도록 구성한 process_frame() 함수.\n", "\n", "Cell 8: 메인 실행부 – 비디오 파일을 순회하면서 각 프레임을 처리하고 결과 영상을 저장한다.\n", "\n", "Cell 9: 추가 디버깅 – 임의 이미지에 대해 디텍션 결과를 확인할 수 있는 함수.\n", "\n", "각 단계는 on/off 스위치를 통해 개별 기능을 쉽게 활성/비활성할 수 있으므로, 후에 조건이나 기능을 수정하거나 추가할 때도 각 함수의 역할이 명확하게 분리되어 있어 디버깅 및 유지보수가 용이합니다.\n", "\n", "현재 대화의 컨텍스트 사용량: 약 70% 정도.\n", "\n", "위 코드를 복사하여 각 셀 단위로 실행하면, 조건별로 최종 시각화되는 클래스(탐지 결과)가 어떤 조건을 거쳐서 출력되는지 명확하게 확인할 수 있습니다. 필요에 따라 각 함수에 추가 주석과 print 디버그문을 삽입해 검사해 보시기 바랍니다." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.20" } }, "nbformat": 4, "nbformat_minor": 4 }