{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "클래스별 색상 지정 개선\n", "\n", "기존에 worker, helmet, head는 고정 색상을 사용하고, 나머지 클래스는 전역 COLOR_PALETTE에서 순차적으로 할당하도록 수정하였습니다.\n", "\n", "전역 변수 FIXED_CLASS_COLORS와 ASSIGNED_CLASS_COLORS를 추가하여, 고정 색상이 아닌 클래스에 대해 자동으로 서로 겹치지 않는 색이 할당되도록 구현했습니다.\n", "\n", "위험구역/작업구역 알고리즘 수정\n", "\n", "경고 알람은 최종으로 시각화되는 객체(final_vis)에 대해서만 적용하도록 수정하였습니다.\n", "\n", "즉, bbox가 실제로 그려지는 대상(조건을 통과한 객체)만 대상으로 위험구역 내 여부를 체크합니다.\n", "\n", "앙상블 로직 확인 및 분리\n", "\n", "앙상블 로직에서는 whhb:worker와 fall:worker의 경우 두 모델의 worker 검출이 IoU 조건을 만족하면 가중 평균하여 병합하도록 하고, helmet과 head는 IoU 조건 만족 시 가중치에 따라 신뢰도가 높은 쪽만 남기는 방식으로 구현되어 있습니다.\n", "\n", "코드 내 주석을 통해 각 조건의 작동 방식을 상세히 설명해두었습니다.\n", "\n", "낮은 신뢰도 검출 기록 및 분석 기능 추가 (on/off 토글)\n", "\n", "전역 변수 ENABLE_LOW_CONF_LOGGING과 LOW_CONF_THRESHOLD를 추가하여, 설정한 임계값(예: 0.5) 이하의 검출 객체에 대해 영상 이름, 프레임(타임라인), 객체 이름과 함께 간단한 분석(객체 크기, 형태 등) 결과를 로그 파일로 저장하도록 하였습니다.\n", "\n", "검출 함수 detect_objects_single_model()를 수정하여, 클래스별 임계값 통과 여부와 관계없이 낮은 신뢰도 검출(LOW_CONF_THRESHOLD 미만)을 별도 리스트로 기록합니다.\n", "\n", "신뢰도 개선 및 오탐 제거에 대한 코멘트 출력\n", "\n", "영상별로 낮은 신뢰도 로그를 기반으로 각 클래스에 대한 총평(예: “검출된 low confidence 객체 수, 객체 크기 및 형태를 확인하고 threshold 조정 또는 데이터 보강을 고려하세요”)을 산출하여 별도 추천 파일에 저장하도록 하였습니다." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "주요 변경 및 추가 사항\n", "모델·클래스별 시각화 on/off 옵션\n", "\n", "각 모델 설정에 추가로 \"class_vis\" 항목을 넣어, 예를 들어 whhb 모델의 경우\n", "\n", "python\n", "복사\n", "\"class_vis\": {\n", " \"worker\": True,\n", " \"helmet\": True,\n", " \"head\": True,\n", " \"background\": False\n", "}\n", "와 같이 설정하면, 해당 클래스의 검출은 최종 시각화(final_vis)에 포함됩니다.\n", "\n", "process_video()에서 최종 시각화 대상(final_vis)을 구성할 때 각 detection의 모델 및 클래스에 대해 해당 옵션을 확인하도록 수정했습니다.\n", "\n", "고정 색상 사용 범위 제한\n", "\n", "고정 색상(FIXED_CLASS_COLORS)은 오직 whhb 모델에 한해 적용되고, 다른 모델의 동일 클래스(예: worker)는 자동 할당(ASSIGNED_CLASS_COLORS)하도록 수정했습니다.\n", "\n", "이를 위해, get_color_for_class() 함수를 모델 키(model_key)도 인자로 받아 처리하도록 변경하고, 호출 시에도 detection의 model_key를 함께 전달합니다.\n", "\n", "겹치는 bbox 조정 기능\n", "\n", "최종 시각화 대상(final_vis) 중 서로 IOU가 0.6 이상인 두 bbox에 대해, 후자의 박스 크기를 임의(예, scale factor 0.9)로 조절하여 겹침을 완화합니다.\n", "\n", "함수 adjust_overlapping_boxes()를 추가하였으며, process_video()에서 final_vis를 구성한 후 호출하여 bbox를 조정합니다.\n", "\n", "헬멧 착용 여부 판단 및 경고 기능\n", "\n", "헬멧 착용 판단을 위한 별도의 슬라이딩 윈도우(기본 3초, fps에 따라 자동 계산)를 사용합니다.\n", "\n", "whhb 모델에서 worker 검출에 대해, 같은 프레임에서 동일 모델(whhb)의 helmet 검출과의 IOU를 계산하고, 해당 track의 최근 window 내 평균 IOU가 낮으면(예, 0.1 이하) 경고 문구(설정 가능한 텍스트, 위치, 색상 등)를 출력합니다.\n", "\n", "전역 변수로 HELMET_WINDOW_DURATION, HELMET_IOU_LOWER, HELMET_WARNING_ENABLED, HELMET_WARNING_TEXT 등도 추가했습니다.\n", "\n", "process_video() 내에서 각 프레임마다 worker 검출에 대해 헬멧 IOU를 업데이트하고, 조건에 부합하면 경고 문구를 표시합니다.\n", "\n", "아래 전체 코드를 확인하시고, 필요에 따라 전역 변수(경고 문구, 윈도우 길이, IOU threshold 등)를 조절하시면 됩니다." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 1] 라이브러리 임포트 + 최상위 전역 설정\n", "import os\n", "import sys\n", "import time\n", "import glob\n", "import cv2\n", "import numpy as np\n", "import matplotlib\n", "import matplotlib.pyplot as plt\n", "from loguru import logger\n", "from types import SimpleNamespace\n", "import mmcv\n", "import gc\n", "import torch\n", "import torch.nn as nn\n", "from tqdm import tqdm\n", "\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 ByteTracker\n", "from yolox.tracker.byte_tracker import BYTETracker\n", "from yolox.tracking_utils.timer import Timer\n", "\n", "########################################\n", "# (!!!) 자주 바꿀 수 있는 글로벌 파라미터들\n", "########################################\n", "\n", "# 1) 경로/저장 관련\n", "BASE_DIR = \"/DATA2/ltb\"\n", "VIDEO_DIR = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/(박스)it1_현대_테스트할영상_편집본_모음\"\n", "SAVE_RESULT_BASE = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/output_val_samples/(박스)it1_현대_테스트할영상_편집본_모음_위험구역_전체테스트_11_하네스_fall클래스몇개끔_윈도우3초로\"\n", "SAVE_SUBFOLDER = \"exp_polygon_alarm_v3\"\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", "# 2) 모델 목록 (예: 4개 whhb, con8, harness, fall)\n", "# 각 모델 설정에 'class_vis' 추가 (True/False로 시각화 여부 지정)\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\": True,\n", " \"vis_on\": True,\n", " \"class_params\": {\n", " \"default\": {\"detect_threshold\": 0.0, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.3},\n", " \"worker\": {\"detect_threshold\": 0.1, \"stable_time_window\": 3.0, \"required_pass_ratio\": 0.7},\n", " \"helmet\": {\"detect_threshold\": 0.1, \"stable_time_window\": 2.0, \"required_pass_ratio\": 0.3}\n", " },\n", " \"class_vis\": { # whhb의 경우, 고정 색상 사용 (아래 설정에서 worker, helmet, head 고정)\n", " \"worker\": True,\n", " \"helmet\": True,\n", " \"head\": True,\n", " \"background\": False\n", " }\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\": True,\n", " \"vis_on\": True,\n", " \"class_params\": {\n", " \"default\": {\"detect_threshold\": 0.05, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.2}\n", " },\n", " \"class_vis\": {} # 빈 dict는 모두 True(기본)\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\": True,\n", " \"vis_on\": True,\n", " \"class_params\": {\n", " \"default\": {\"detect_threshold\": 0.0, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.1},\n", " \"worker\": {\"detect_threshold\": 0.15, \"stable_time_window\": 3.0, \"required_pass_ratio\": 0.6}\n", " },\n", " \"class_vis\": {} # 기본 True\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', 'signalman', 'helmet', 'harness', 'mixer_truck', 'excavator'],\n", " \"do_inference\": True,\n", " \"vis_on\": True,\n", " \"class_params\": {\n", " \"default\": {\"detect_threshold\": 0.0, \"stable_time_window\": 1.0, \"required_pass_ratio\": 0.4},\n", " \"harness\": {\"detect_threshold\": 0.4, \"stable_time_window\": 3.0, \"required_pass_ratio\": 0.3},\n", " \"signalman\": {\"detect_threshold\": 0.4, \"stable_time_window\": 3.0, \"required_pass_ratio\": 0.3}\n", " },\n", " \"class_vis\": {\n", " \"mixer_truck\": False,\n", " \"excavator\": False,\n", " \"worker\": False,\n", " \"helmet\": False,\n", " }\n", " },\n", "}\n", "\n", "# 3) 앙상블 모드: \"union\" or \"priority\"\n", "ENSEMBLE_STRATEGY = \"union\"\n", "\n", "# 4) ByteTracker 파라미터\n", "TRACK_THRESH = 0.05\n", "TRACK_BUFFER = 300\n", "MATCH_THRESH = 0.95\n", "\n", "##############\n", "DEVICE = 'cuda:0'\n", "\n", "# 5) 시각화 폰트 & BBox 라인 굵기\n", "FONT_SCALE_DEFAULT = 1\n", "FONT_THICKNESS_DEFAULT = 1\n", "BBOX_LINE_THICKNESS = 2\n", "\n", "# OpenCV 폰트 종류\n", "USE_CUSTOM_FONT_FACE = True\n", "CUSTOM_FONT_FACE = cv2.FONT_HERSHEY_DUPLEX\n", "\n", "# [수정] 클래스별 색상 지정\n", "# 단, whhb 모델은 고정 색상(FIXED_CLASS_COLORS)을 사용하고, 나머지 모델은 자동 할당\n", "COLOR_PALETTE = [\n", " (42,42,165), (211,0,148), \n", " (255,255,0), (0,255,255), (255,0,255), (100,100,100), (128,128,0),\n", " (128,0,128), (0,128,128), (128,128,255), (128,255,128), (255,128,128),\n", " (192,192,192), (0,0,128), (128,0,0), (128,255,255), (255,128,255)\n", "]\n", "# 고정할 클래스 색상(whhb 전용)\n", "FIXED_CLASS_COLORS = {\n", " \"worker\": (0,255,0),\n", " \"helmet\": (42,42,165),\n", " \"head\": (0,0,255)\n", "}\n", "# 자동 할당용 딕셔너리 (whhb가 아닌 모델 및 고정되지 않은 클래스)\n", "ASSIGNED_CLASS_COLORS = {}\n", "\n", "# 7) 앙상블 로직 파라미터\n", "WHHB_FALL_WORKER_IOU_THR = 0.3\n", "WHHB_FALL_WORKER_SCORE_WEIGHT_WHHB = 0.8\n", "WHHB_FALL_WORKER_SCORE_WEIGHT_FALL = 0.3\n", "HELMET_HEAD_IOU_THR = 0.2\n", "HELMET_SCORE_WEIGHT = 0.8\n", "HEAD_SCORE_WEIGHT = 0.3\n", "\n", "# 8) 다각형 구역 알람 관련\n", "ENABLE_REGION_ALARM = True \n", "TARGET_CLASS_FOR_ALARM = \"worker\" \n", "POLYGON_COORDS = [(400,400),(1200,400),(1500,1000),(700,1000)]\n", "ALARM_TEXT = \"WARNING: Worker in restricted area!\" \n", "ALARM_TEXT_COLOR = (0,0,255)\n", "ALARM_TEXT_SCALE = 2\n", "ALARM_TEXT_THICKNESS = 2\n", "ALARM_POSITION = (50,50)\n", "\n", "########################################\n", "# (!!!) 추가 전역: 영상 처리 on/off & 리사이즈 배율\n", "VIDEO_INFERENCE_ON = True\n", "VIDEO_SAVE_ON = True\n", "RESIZE_FACTOR = 1.0\n", "\n", "# [추가] 낮은 신뢰도 검출 로깅 on/off 및 임계값\n", "ENABLE_LOW_CONF_LOGGING = True\n", "LOW_CONF_THRESHOLD = 0.2\n", "\n", "# 로그 저장용 폴더 (영상별로 개별 파일 생성)\n", "LOW_CONF_LOG_DIR = os.path.join(SAVE_RESULT, \"low_conf_logs\")\n", "os.makedirs(LOW_CONF_LOG_DIR, exist_ok=True)\n", "\n", "# 헬멧 착용 경고 관련 전역 변수\n", "HELMET_WINDOW_DURATION = 3.0 # seconds\n", "HELMET_IOU_LOWER = 0.1 # 3초 윈도우 내 평균 IOU threshold\n", "HELMET_WARNING_ENABLED = False\n", "HELMET_WARNING_TEXT = \"Helmet not detected!\"\n", "HELMET_WARNING_POSITION = (100,100) # 기본 위치, 필요시 조정\n", "HELMET_WARNING_COLOR = (0,0,255)\n", "HELMET_WARNING_FONT_SCALE = 1\n", "HELMET_WARNING_FONT_THICKNESS = 1\n", "\n", "# 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", "########################################\n", "# [추가] 시각화 옵션 설정 (on/off 및 텍스트 위치 지정)\n", "VISUALIZATION_OPTIONS = {\n", " \"draw_bbox\": True,\n", " \"draw_class\": True,\n", " \"draw_confidence\": True,\n", " \"draw_track_id\": False,\n", " \"text_positions\": {\n", " \"class\": \"top\",\n", " \"confidence\": \"top\",\n", " \"track_id\": \"top\"\n", " },\n", " \"text_padding\": 15,\n", "}\n", "\n", "# [추가] 클래스별 텍스트 위치 옵션 (클래스마다 개별 지정 가능)\n", "CLASS_TEXT_POSITION = {\n", " \"worker\": {\"class\": \"top\", \"confidence\": \"top\", \"track_id\": \"top\"},\n", " \"helmet\": {\"class\": \"bottom\", \"confidence\": \"bottom\", \"track_id\": \"bottom\"},\n", " \"head\": {\"class\": \"right\", \"confidence\": \"right\", \"track_id\": \"right\"},\n", "}\n", "\n", "########################################\n", "# [추가] 영역(Zone) 시각화 옵션 설정\n", "ZONE_VIS_OPTIONS = {\n", " \"danger\": {\n", " \"enabled\": True,\n", " \"color\": (0, 0, 255),\n", " \"line_thickness\": 2,\n", " \"text\": \"Danger Zone\",\n", " \"text_enabled\": True,\n", " \"text_color\": (0, 0, 255),\n", " \"font_scale\": 1.5,\n", " \"font_thickness\": 2,\n", " \"font_face\": cv2.FONT_HERSHEY_SIMPLEX,\n", " \"text_position_offset\": (10, 30)\n", " },\n", " \"detect\": {\n", " \"enabled\": True,\n", " \"color\": (0, 255, 0),\n", " \"line_thickness\": 2,\n", " \"text\": \"Detect Zone\",\n", " \"text_enabled\": True,\n", " \"text_color\": (0, 255, 0),\n", " \"font_scale\": 1.5,\n", " \"font_thickness\": 2,\n", " \"font_face\": cv2.FONT_HERSHEY_SIMPLEX,\n", " \"text_position_offset\": (-10, 30)\n", " }\n", "}\n", "\n", "print(\"===== [Cell 1] 전역 설정 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 2] 첫 번째 영상 디버깅 + 그리드 표시 + 위험구역 폴리곤\n", "DEBUG_FRAME_IDX = 0 \n", "GRID_INTERVAL = 200 \n", "\n", "def debug_first_video_and_polygon():\n", " \"\"\"\n", " VIDEO_DIR 내 첫 번째 영상을 열어, 영상 정보와 DEBUG_FRAME_IDX 프레임의 그리드 및 위험구역(Polygon)을 시각화\n", " \"\"\"\n", " exts = {\".mp4\", \".avi\", \".mov\", \".mkv\"}\n", " files = sorted(os.listdir(VIDEO_DIR))\n", " first_video = None\n", " for f in files:\n", " base, ex = os.path.splitext(f)\n", " if ex.lower() in exts:\n", " first_video = os.path.join(VIDEO_DIR, f)\n", " break\n", " if first_video is None:\n", " print(\"[경고] VIDEO_DIR 내 처리할 영상 없음\")\n", " return\n", "\n", " cap = cv2.VideoCapture(first_video)\n", " if not cap.isOpened():\n", " print(\"[오류] 첫번째 영상 열기 실패 =>\", first_video)\n", " return\n", "\n", " frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n", " frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n", " fps = cap.get(cv2.CAP_PROP_FPS)\n", " frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n", " duration = frame_count / fps if fps > 0 else 0.0\n", "\n", " print(\"📌 첫번째 영상 정보:\")\n", " print(f\" - 경로: {first_video}\")\n", " print(f\" - 해상도: {frame_width} x {frame_height}\")\n", " print(f\" - FPS: {fps}\")\n", " print(f\" - 총 프레임: {frame_count}\")\n", " print(f\" - 영상 길이: {duration:.2f} 초\")\n", "\n", " cap.set(cv2.CAP_PROP_POS_FRAMES, DEBUG_FRAME_IDX)\n", " ret, debug_frame = cap.read()\n", " cap.release()\n", " if not ret:\n", " print(f\"[오류] DEBUG_FRAME_IDX({DEBUG_FRAME_IDX}) 프레임 읽기 실패\")\n", " return\n", "\n", " debug_frame_rgb = cv2.cvtColor(debug_frame, cv2.COLOR_BGR2RGB)\n", " plt.figure(figsize=(8,6))\n", " plt.imshow(debug_frame_rgb)\n", " plt.title(f\"디버그 프레임 idx={DEBUG_FRAME_IDX}\")\n", " plt.xticks(np.arange(0, debug_frame_rgb.shape[1], GRID_INTERVAL))\n", " plt.yticks(np.arange(0, debug_frame_rgb.shape[0], GRID_INTERVAL))\n", " plt.grid(True, which='both', linestyle='--', linewidth=0.5)\n", " plt.axis(\"on\")\n", "\n", " poly_x = [pt[0] for pt in POLYGON_COORDS] + [POLYGON_COORDS[0][0]]\n", " poly_y = [pt[1] for pt in POLYGON_COORDS] + [POLYGON_COORDS[0][1]]\n", " plt.plot(poly_x, poly_y, color=\"yellow\", linestyle=\"--\", marker=\"o\", markersize=5)\n", "\n", " plt.show()\n", "\n", " debug_frame_path = os.path.join(SAVE_RESULT, \"frame_debug_original.jpg\")\n", " cv2.imwrite(debug_frame_path, debug_frame)\n", " print(f\"✅ 디버그 프레임 저장: {debug_frame_path}\")\n", "\n", "debug_first_video_and_polygon()\n", "print(\"===== [Cell 2] 디버깅(첫 프레임+그리드+폴리곤) 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 3] 다중 ROI & 여러 알람대상 클래스 예시 (주석)\n", "\"\"\"\n", "아래는 다중 ROI & 여러 알람 대상 클래스 예시 주석입니다.\n", "\n", "# 예) 다중 ROI\n", "POLYGON_COORDS_LIST = [\n", " [(100,100),(200,100),(200,200),(100,200)],\n", " [(300,300),(400,300),(400,400),(300,400)]\n", "]\n", "\n", "# 예) 여러 알람 대상 클래스\n", "TARGET_CLASSES_FOR_ALARM = [\"worker\",\"helmet\"]\n", "\n", "실제로 적용하려면, process_video() 내에서\n", "- polygon check를 for poly in POLYGON_COORDS_LIST 형태로 반복\n", "- if dd[\"label\"] in TARGET_CLASSES_FOR_ALARM: ... 로 검사하세요.\n", "\"\"\"\n", "print(\"===== [Cell 3] 다중 ROI, 여러 클래스 알람 예시 주석 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 4] 작업구역 폴리곤 시각화 및 디버깅 (수정된 코드)\n", "RESTRICT_VIZ_TO_POLYGON = True\n", "RESTRICT_VIZ_CLASSES = [\"worker\", \"helmet\"]\n", "WORKZONE_POLYGON = [(0, 0), (1400, 0), (1400, 1080), (0, 1080)]\n", "\n", "def is_point_in_any_polygon(point, polygons):\n", " for poly in polygons:\n", " cnt = np.array([poly], dtype=np.int32)\n", " val = cv2.pointPolygonTest(cnt, point, False)\n", " if val >= 0:\n", " return True\n", " return False\n", "\n", "def debug_workzone_polygon(frame, polygon_coords):\n", " debug_img = frame.copy()\n", " pts = np.array([polygon_coords], dtype=np.int32)\n", " cv2.polylines(debug_img, pts, True, (0, 255, 255), 2)\n", " save_path_cv = os.path.join(SAVE_RESULT, \"debug_workzone_polygon.jpg\")\n", " cv2.imwrite(save_path_cv, debug_img)\n", " print(f\"[디버그] 작업구역 폴리곤 이미지 저장 (OpenCV): {save_path_cv}\")\n", " \n", " debug_img_rgb = cv2.cvtColor(debug_img, cv2.COLOR_BGR2RGB)\n", " plt.figure(figsize=(8, 6))\n", " plt.imshow(debug_img_rgb)\n", " plt.title(\"디버그: 작업구역 폴리곤 시각화\")\n", " GRID_INTERVAL = 200\n", " height, width = debug_img.shape[:2]\n", " plt.xticks(np.arange(0, width, GRID_INTERVAL))\n", " plt.yticks(np.arange(0, height, GRID_INTERVAL))\n", " plt.grid(True, which='both', linestyle='--', linewidth=0.5)\n", " poly_x = [pt[0] for pt in polygon_coords] + [polygon_coords[0][0]]\n", " poly_y = [pt[1] for pt in polygon_coords] + [polygon_coords[0][1]]\n", " plt.plot(poly_x, poly_y, color=\"yellow\", linestyle=\"--\", marker=\"o\", markersize=5)\n", " \n", " save_path_plt = os.path.join(SAVE_RESULT, \"debug_workzone_polygon_matplotlib.jpg\")\n", " plt.savefig(save_path_plt)\n", " print(f\"[디버그] 작업구역 폴리곤 matplotlib 이미지 저장: {save_path_plt}\")\n", " \n", " plt.show()\n", "\n", "def debug_workzone_polygon_from_video():\n", " exts = {\".mp4\", \".avi\", \".mov\", \".mkv\"}\n", " files = sorted(os.listdir(VIDEO_DIR))\n", " first_video = None\n", " for f in files:\n", " base, ex = os.path.splitext(f)\n", " if ex.lower() in exts:\n", " first_video = os.path.join(VIDEO_DIR, f)\n", " break\n", " if first_video is None:\n", " print(\"[경고] VIDEO_DIR 내 처리할 영상 없음\")\n", " return\n", "\n", " cap = cv2.VideoCapture(first_video)\n", " if not cap.isOpened():\n", " print(f\"[오류] 첫번째 영상 열기 실패 => {first_video}\")\n", " return\n", "\n", " frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n", " frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n", " fps = cap.get(cv2.CAP_PROP_FPS)\n", " frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n", " duration = frame_count / fps if fps > 0 else 0.0\n", " print(\"📌 첫번째 영상 정보 (WORKZONE 디버그용):\")\n", " print(f\" - 경로: {first_video}\")\n", " print(f\" - 해상도: {frame_width} x {frame_height}\")\n", " print(f\" - FPS: {fps}\")\n", " print(f\" - 총 프레임: {frame_count}\")\n", " print(f\" - 영상 길이: {duration:.2f} 초\")\n", "\n", " cap.set(cv2.CAP_PROP_POS_FRAMES, DEBUG_FRAME_IDX)\n", " ret, frame = cap.read()\n", " cap.release()\n", " if not ret:\n", " print(f\"[오류] DEBUG_FRAME_IDX({DEBUG_FRAME_IDX}) 프레임 읽기 실패\")\n", " return\n", "\n", " debug_workzone_polygon(frame, WORKZONE_POLYGON)\n", "\n", "debug_workzone_polygon_from_video()\n", "print(\"===== [Cell 4] 작업구역 폴리곤 시각화 및 디버깅 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 5] 헬퍼 함수 (IoU, distance, color 등)\n", "def iou_calc(a, b):\n", " inter_x1 = max(a[0], b[0])\n", " inter_y1 = max(a[1], b[1])\n", " inter_x2 = min(a[2], b[2])\n", " inter_y2 = min(a[3], 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 = (a[2]-a[0])*(a[3]-a[1])\n", " area_b = (b[2]-b[0])*(b[3]-b[1])\n", " union = area_a + area_b - inter_area\n", " if union <= 0:\n", " return 0.0\n", " return inter_area / union\n", "\n", "def distance(a, b):\n", " return np.hypot(a[0]-b[0], a[1]-b[1])\n", "\n", "# [수정] get_color_for_class: whhb 모델은 고정 색상 사용, 그 외는 자동 할당\n", "def get_color_for_class(cname, model_key=None):\n", " if model_key == \"whhb\" and cname in FIXED_CLASS_COLORS:\n", " return FIXED_CLASS_COLORS[cname]\n", " if cname in ASSIGNED_CLASS_COLORS:\n", " return ASSIGNED_CLASS_COLORS[cname]\n", " for color in COLOR_PALETTE:\n", " if color not in ASSIGNED_CLASS_COLORS.values():\n", " ASSIGNED_CLASS_COLORS[cname] = color\n", " return color\n", " assigned_color = COLOR_PALETTE[len(ASSIGNED_CLASS_COLORS) % len(COLOR_PALETTE)]\n", " ASSIGNED_CLASS_COLORS[cname] = assigned_color\n", " return assigned_color\n", "\n", "def is_point_in_polygon(point, polygon):\n", " cnt = np.array(polygon, dtype=np.int32)\n", " val = cv2.pointPolygonTest(cnt, point, False)\n", " return (val >= 0)\n", "\n", "print(\"===== [Cell 5] 헬퍼 함수 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 6] 모델 로드\n", "LOADED_MODELS = {}\n", "def load_all_models():\n", " loaded = {}\n", " for mk, info in MODELS_CONFIG.items():\n", " if not info[\"do_inference\"]:\n", " print(f\"[모델 스킵] {mk}\")\n", " continue\n", " print(f\"[모델 로드] {mk}\")\n", " model_obj = init_detector(info[\"config\"], info[\"checkpoint\"],\n", " device=DEVICE if torch.cuda.is_available() else \"cpu\")\n", " model_obj.dataset_meta = {\"CLASSES\": info[\"class_names\"]}\n", " loaded[mk] = model_obj\n", " return loaded\n", "\n", "LOADED_MODELS = load_all_models()\n", "print(\"===== [Cell 6] 모델 로드 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 7] 디텍션 & 시각화 함수 (텍스트 시각화 개선)\n", "def get_class_param(model_key, class_label, param_key):\n", " if model_key not in MODELS_CONFIG:\n", " return None\n", " cparams = MODELS_CONFIG[model_key].get(\"class_params\", {})\n", " if class_label in cparams:\n", " return cparams[class_label].get(param_key, cparams.get(\"default\", {}).get(param_key, None))\n", " else:\n", " return cparams.get(\"default\", {}).get(param_key, None)\n", "\n", "def detect_objects_single_model(frame, model_key, model_obj):\n", " accepted = []\n", " low_conf_logs = []\n", " cnames = MODELS_CONFIG[model_key][\"class_names\"]\n", " result = inference_detector(model_obj, frame)\n", " if not hasattr(result, \"pred_instances\"):\n", " return accepted, low_conf_logs\n", " bboxes = result.pred_instances.bboxes\n", " scores = result.pred_instances.scores\n", " labels = result.pred_instances.labels\n", " for i in range(len(bboxes)):\n", " cid = labels[i].item()\n", " if cid >= len(cnames):\n", " continue\n", " cname = cnames[cid]\n", " sc = scores[i].item()\n", " box = bboxes[i].tolist()\n", " thr = get_class_param(model_key, cname, \"detect_threshold\")\n", " if thr is None:\n", " thr = 0.0\n", " if sc >= thr:\n", " accepted.append({\n", " \"bbox\": box,\n", " \"score\": sc,\n", " \"label\": cname,\n", " \"model_key\": model_key\n", " })\n", " if ENABLE_LOW_CONF_LOGGING and sc < LOW_CONF_THRESHOLD:\n", " low_conf_logs.append({\n", " \"bbox\": box,\n", " \"score\": sc,\n", " \"label\": cname,\n", " \"model_key\": model_key\n", " })\n", " return accepted, low_conf_logs\n", "\n", "def draw_text_with_options(image, text_lines, anchor_point,\n", " text_position=\"top\",\n", " font_scale=None, font_thickness=None,\n", " color=(0,255,0),\n", " auto_adjust=True):\n", " if font_scale is None:\n", " font_scale = FONT_SCALE_DEFAULT\n", " if font_thickness is None:\n", " font_thickness = FONT_THICKNESS_DEFAULT\n", " font_face = CUSTOM_FONT_FACE if USE_CUSTOM_FONT_FACE else cv2.FONT_HERSHEY_SIMPLEX\n", " xA, yA = anchor_point\n", " for i, line in enumerate(text_lines):\n", " sz, _ = cv2.getTextSize(line, font_face, font_scale, font_thickness)\n", " tw, th = sz\n", " offset = i * (th + 5)\n", " if text_position == \"top\":\n", " org = (xA, yA - 5 - offset)\n", " elif text_position == \"bottom\":\n", " org = (xA, yA + th + 5 + offset)\n", " elif text_position == \"left\":\n", " org = (xA - tw - 5, yA + th + offset)\n", " else:\n", " org = (xA + 5, yA + th + offset)\n", " if auto_adjust:\n", " if org[0] < 0:\n", " org = (5, org[1])\n", " elif org[0] + tw > image.shape[1]:\n", " org = (image.shape[1] - tw - 5, org[1])\n", " if org[1] - th < 0:\n", " org = (org[0], th + 5)\n", " elif org[1] > image.shape[0]:\n", " org = (org[0], image.shape[0] - 5)\n", " cv2.putText(image, line, org, font_face, font_scale, color, font_thickness)\n", "\n", "def draw_detection_texts(image, detection, vis_opts, class_text_opts):\n", " x1, y1, x2, y2 = map(int, detection[\"bbox\"])\n", " label = detection.get(\"label\", \"unknown\")\n", " model_key = detection.get(\"model_key\")\n", " color = get_color_for_class(label, model_key)\n", " texts = {}\n", " if vis_opts.get(\"draw_class\", True):\n", " texts[\"class\"] = f'{model_key}:{label}'\n", " if vis_opts.get(\"draw_confidence\", True):\n", " texts[\"confidence\"] = f'{detection.get(\"score\", 0.0):.2f}'\n", " if vis_opts.get(\"draw_track_id\", True) and detection.get(\"track_id\") is not None:\n", " texts[\"track_id\"] = f'ID:{detection.get(\"track_id\")}'\n", " pos_opts = class_text_opts.get(label, vis_opts.get(\"text_positions\", {}))\n", " offsets = {\"top\": 0, \"bottom\": 0, \"left\": 0, \"right\": 0}\n", " font_face = CUSTOM_FONT_FACE if USE_CUSTOM_FONT_FACE else cv2.FONT_HERSHEY_SIMPLEX\n", " font_scale = FONT_SCALE_DEFAULT\n", " thickness = FONT_THICKNESS_DEFAULT\n", " padding = vis_opts.get(\"text_padding\", 5)\n", " for text_type, text_str in texts.items():\n", " side = pos_opts.get(text_type, \"top\")\n", " if side == \"top\":\n", " base_x, base_y = x1, y1\n", " elif side == \"bottom\":\n", " base_x, base_y = x1, y2\n", " elif side == \"left\":\n", " base_x, base_y = x1, y1\n", " elif side == \"right\":\n", " base_x, base_y = x2, y1\n", " else:\n", " base_x, base_y = x1, y1\n", " (tw, th), _ = cv2.getTextSize(text_str, font_face, font_scale, thickness)\n", " if side == \"top\":\n", " anchor = (base_x, base_y - offsets[\"top\"] * (th + padding) - 5)\n", " offsets[\"top\"] += 1\n", " elif side == \"bottom\":\n", " anchor = (base_x, base_y + offsets[\"bottom\"] * (th + padding) + th + 5)\n", " offsets[\"bottom\"] += 1\n", " elif side == \"left\":\n", " anchor = (base_x - offsets[\"left\"] * (tw + padding) - 5, base_y + th)\n", " offsets[\"left\"] += 1\n", " elif side == \"right\":\n", " anchor = (base_x + offsets[\"right\"] * (tw + padding) + 5, base_y + th)\n", " offsets[\"right\"] += 1\n", " else:\n", " anchor = (base_x, base_y)\n", " cv2.putText(image, text_str, anchor, font_face, font_scale, color, thickness)\n", "\n", "def draw_detections(frame, detections):\n", " out = frame.copy()\n", " for det in detections:\n", " mk = det.get(\"model_key\", \"?\")\n", " # 1. 모델별 클래스 시각화 on/off 필터링: 만약 해당 모델의 class_vis에서 False면 스킵\n", " model_conf = MODELS_CONFIG.get(mk, {})\n", " class_vis = model_conf.get(\"class_vis\", {})\n", " lbl = det.get(\"label\", \"unknown\")\n", " if not class_vis.get(lbl, True):\n", " continue\n", " box = det[\"bbox\"]\n", " x1, y1, x2, y2 = map(int, box)\n", " if VISUALIZATION_OPTIONS.get(\"draw_bbox\", True):\n", " cv2.rectangle(out, (x1, y1), (x2, y2), get_color_for_class(lbl, mk), BBOX_LINE_THICKNESS)\n", " draw_detection_texts(out, det, VISUALIZATION_OPTIONS, CLASS_TEXT_POSITION)\n", " return out\n", "\n", "print(\"===== [Cell 7] 디텍션 & 시각화 함수 (텍스트 시각화 개선) 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 8] 다각형 구역(ROI) 디버깅 함수\n", "import matplotlib\n", "matplotlib.use('Agg')\n", "def check_polygon_region_and_debug(frame, polygon_coords, debug_save_path):\n", " dbg = frame.copy()\n", " pts = np.array([polygon_coords], dtype=np.int32)\n", " cv2.polylines(dbg, pts, True, (0,255,255), 2)\n", " cv2.imwrite(debug_save_path, dbg)\n", " print(\"[DEBUG] polygon =>\", debug_save_path)\n", "\n", "print(\"===== [Cell 8] Polygon ROI 디버깅 함수 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 9] ByteTracker + 슬라이딩 윈도우\n", "def create_tracker():\n", " return BYTETracker(SimpleNamespace(\n", " track_thresh=TRACK_THRESH,\n", " track_buffer=TRACK_BUFFER,\n", " match_thresh=MATCH_THRESH\n", " ))\n", "\n", "def update_track_states(tracks, track_dict, fps):\n", " for t in tracks:\n", " tid = t.track_id\n", " x1, y1, w, h = t.tlwh\n", " x2 = x1 + w\n", " y2 = y1 + h\n", " cnow = ((x1+x2)/2, (y1+y2)/2)\n", " if tid not in track_dict:\n", " track_dict[tid] = {\"last_center\": None, \"is_filtered\": False,\n", " \"score_history\": [], \"score_passed\": False,\n", " \"class_label\": None, \"model_key\": None}\n", " if track_dict[tid][\"last_center\"] is None:\n", " track_dict[tid][\"last_center\"] = cnow\n", " else:\n", " dist_val = distance(cnow, track_dict[tid][\"last_center\"])\n", " track_dict[tid][\"last_center\"] = cnow\n", " if track_dict[tid][\"is_filtered\"]:\n", " continue\n", " sc = t.score\n", " track_dict[tid][\"score_history\"].append(sc)\n", " mk = track_dict[tid][\"model_key\"]\n", " clbl = track_dict[tid][\"class_label\"]\n", " if mk is None or clbl is None:\n", " stable_tw = 1.0; req_ratio = 0.3; thr = 0.0\n", " else:\n", " stable_tw = get_class_param(mk, clbl, \"stable_time_window\")\n", " req_ratio = get_class_param(mk, clbl, \"required_pass_ratio\")\n", " thr = get_class_param(mk, clbl, \"detect_threshold\")\n", " if stable_tw is None: stable_tw = 1.0\n", " if req_ratio is None: req_ratio = 0.3\n", " if thr is None: thr = 0.0\n", " window_size = int(fps * stable_tw)\n", " if window_size < 1: window_size = 1\n", " hist = track_dict[tid][\"score_history\"]\n", " if len(hist) > window_size:\n", " track_dict[tid][\"score_history\"] = hist[-window_size:]\n", " pass_count = sum([1 for s in hist if s >= thr])\n", " ratio = pass_count / len(hist) if len(hist) > 0 else 0\n", " if ratio >= req_ratio:\n", " track_dict[tid][\"score_passed\"] = True\n", "\n", "print(\"===== [Cell 9] ByteTracker + 슬라이딩 윈도우 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 10] 앙상블 로직 (whhb_worker+fall_worker, helmet+head)\n", "def custom_ensemble_detections(dets):\n", " \"\"\"\n", " 앙상블 로직:\n", " - worker: whhb와 fall의 worker 검출이 IoU>=WHHB_FALL_WORKER_IOU_THR이면 가중평균 병합\n", " - helmet vs head: IoU>=HELMET_HEAD_IOU_THR이면 가중치가 높은 쪽 선택\n", " \"\"\"\n", " used = [False] * len(dets)\n", " final = []\n", " for i in range(len(dets)):\n", " if used[i]:\n", " continue\n", " d1 = dets[i]\n", " sc1 = d1[\"score\"]\n", " lbl1 = d1[\"label\"]\n", " mk1 = d1[\"model_key\"]\n", " merged = None\n", " for j in range(i+1, len(dets)):\n", " if used[j]:\n", " continue\n", " d2 = dets[j]\n", " sc2 = d2[\"score\"]\n", " lbl2 = d2[\"label\"]\n", " mk2 = d2[\"model_key\"]\n", " iouv = iou_calc(d1[\"bbox\"], d2[\"bbox\"])\n", " if lbl1==\"worker\" and lbl2==\"worker\" and mk1==\"whhb\" and mk2==\"fall\" and iouv>=WHHB_FALL_WORKER_IOU_THR:\n", " new_sc = sc1*WHHB_FALL_WORKER_SCORE_WEIGHT_WHHB + sc2*WHHB_FALL_WORKER_SCORE_WEIGHT_FALL\n", " merged = {\"bbox\": d1[\"bbox\"], \"score\": new_sc, \"label\": \"worker\", \"model_key\": \"whhb_fall\"}\n", " used[j] = True\n", " break\n", " if ((lbl1==\"helmet\" and lbl2==\"head\") or (lbl1==\"head\" and lbl2==\"helmet\")) and iouv>=HELMET_HEAD_IOU_THR:\n", " sc1w = sc1 * (HELMET_SCORE_WEIGHT if lbl1==\"helmet\" else HEAD_SCORE_WEIGHT)\n", " sc2w = sc2 * (HELMET_SCORE_WEIGHT if lbl2==\"helmet\" else HEAD_SCORE_WEIGHT)\n", " if sc1w >= sc2w:\n", " used[j] = True\n", " else:\n", " used[i] = True\n", " break\n", " if merged:\n", " final.append(merged)\n", " used[i] = True\n", " else:\n", " if not used[i]:\n", " final.append(d1)\n", " return final\n", "\n", "print(\"===== [Cell 10] 앙상블 로직 완료 =====\")\n", "print(\" [설명] worker는 whhb와 fall의 worker가 IoU 조건 만족 시 병합, helmet과 head는 높은 가중치 검출 선택\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 11] 메인 처리 함수\n", "# 새로운 기능: bbox 간 IOU>=0.6인 경우 자동 크기 조정\n", "def adjust_overlapping_boxes(detections, iou_thresh=0.6, scale_factor=0.9):\n", " for i in range(len(detections)):\n", " for j in range(i+1, len(detections)):\n", " iou_val = iou_calc(detections[i][\"bbox\"], detections[j][\"bbox\"])\n", " if iou_val >= iou_thresh:\n", " # 두 번째 박스 크기 조절\n", " x1, y1, x2, y2 = detections[j][\"bbox\"]\n", " cx = (x1+x2)/2\n", " cy = (y1+y2)/2\n", " w = (x2-x1) * scale_factor\n", " h = (y2-y1) * scale_factor\n", " detections[j][\"bbox\"] = [int(cx-w/2), int(cy-h/2), int(cx+w/2), int(cy+h/2)]\n", " return detections\n", "\n", "# 헬멧 착용 여부 판단을 위한 windowing (전역 변수 helmet_window는 process_video 내에서 초기화)\n", "# 각 worker (whhb 모델)의 track id별로 최근 헬멧 IOU 값을 저장하여 평균이 낮으면 경고 표시\n", "def analyze_low_conf_detection(detection):\n", " frame_shape = detection.get(\"frame_shape\")\n", " frame_area = frame_shape[0]*frame_shape[1] if frame_shape is not None else 1\n", " x1, y1, x2, y2 = map(int, detection[\"bbox\"])\n", " width = x2 - x1; height = y2 - y1; area = width * height\n", " ratio = area / frame_area\n", " if ratio < 0.0005:\n", " reason = \"객체 크기가 작음 (작은 영역)\"\n", " elif ratio > 0.5:\n", " reason = \"객체 크기가 큼 (너무 큰 영역)\"\n", " else:\n", " aspect = width / height if height != 0 else 0\n", " if aspect < 0.5 or aspect > 2.0:\n", " reason = \"비정상적인 형태 (aspect ratio 이상)\"\n", " else:\n", " reason = \"오탐 가능성 높음 (모호한 특징)\"\n", " return reason\n", "\n", "def generate_recommendations(low_conf_log, video_name):\n", " recommendations = f\"Video: {video_name}\\n\"\n", " if not low_conf_log:\n", " recommendations += \"모든 객체의 신뢰도가 양호합니다.\\n\"\n", " return recommendations\n", " class_counts = {}\n", " for item in low_conf_log:\n", " label = item[\"label\"]\n", " class_counts[label] = class_counts.get(label, 0) + 1\n", " for label, count in class_counts.items():\n", " recommendations += f\"클래스 '{label}': 낮은 신뢰도 검출 {count}건. 검출 객체의 크기, 형태, 색상 등을 확인하고 threshold 조정 또는 데이터 보강 고려.\\n\"\n", " return recommendations\n", "\n", "def process_video(video_path, do_inference=True, do_save=True, resize_factor=1.0):\n", " if not do_inference:\n", " print(\"[스킵] =>\", video_path)\n", " return\n", " base_name = os.path.basename(video_path)\n", " name, ex = os.path.splitext(base_name)\n", " output_video_path = os.path.join(SAVE_RESULT, f\"detected_{name}.mp4\")\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: 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", " frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n", " writer = None\n", " if do_save:\n", " fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n", " ww = int(w_org * resize_factor)\n", " hh = int(h_org * resize_factor)\n", " writer = cv2.VideoWriter(output_video_path, fourcc, fps, (ww, hh))\n", " if not writer.isOpened():\n", " print(\"[경고] VideoWriter 생성 실패 =>\", output_video_path)\n", " writer = None\n", "\n", " tracker = create_tracker()\n", " track_dict = {}\n", " frame_times = []\n", " low_conf_logs_video = []\n", " # 헬멧 착용 여부 판단 window: track_id별로 최근 IOU 값 리스트\n", " helmet_window = {}\n", " alarm_on = ENABLE_REGION_ALARM\n", " st_total = time.time()\n", " print(\"[INFO] 영상 처리 =>\", video_path)\n", " frame_idx = 0\n", " with tqdm(total=frame_count, desc=f\"Processing {base_name}\") as pbar:\n", " while True:\n", " st = time.time()\n", " ret, frame = cap.read()\n", " if not ret:\n", " break\n", " all_dets = []\n", " low_conf_dets_all = []\n", " for mk, mobj in LOADED_MODELS.items():\n", " accepted, low_conf = detect_objects_single_model(frame, mk, mobj)\n", " all_dets.extend(accepted)\n", " for item in low_conf:\n", " item[\"frame_shape\"] = frame.shape\n", " item[\"frame_idx\"] = frame_idx\n", " item[\"time_sec\"] = frame_idx / fps\n", " low_conf_dets_all.extend(low_conf)\n", " low_conf_logs_video.extend(low_conf_dets_all)\n", " final_dets = custom_ensemble_detections(all_dets)\n", " xyxys = []\n", " for d in final_dets:\n", " x1, y1, x2, y2 = d[\"bbox\"]\n", " s = d[\"score\"]\n", " xyxys.append([x1, y1, x2, y2, s])\n", " xyxys_arr = np.array(xyxys, dtype=np.float32) if len(xyxys) > 0 else np.zeros((0,5), dtype=np.float32)\n", " online_tracks = tracker.update(xyxys_arr, [h_org, w_org])\n", " for t in online_tracks:\n", " tid = t.track_id\n", " x1, y1, w, h = t.tlwh\n", " x2 = x1 + w; y2 = y1 + h\n", " if tid not in track_dict:\n", " track_dict[tid] = {\"last_center\": None, \"is_filtered\": False,\n", " \"score_history\": [], \"score_passed\": False,\n", " \"class_label\": None, \"model_key\": None}\n", " if track_dict[tid][\"class_label\"] is not None:\n", " continue\n", " best_iou = 0; best_det = None\n", " for dd in final_dets:\n", " iouv = iou_calc([x1, y1, x2, y2], dd[\"bbox\"])\n", " if iouv > best_iou:\n", " best_iou = iouv; best_det = dd\n", " if best_det and best_iou >= 0.5:\n", " track_dict[tid][\"model_key\"] = best_det[\"model_key\"]\n", " track_dict[tid][\"class_label\"] = best_det[\"label\"]\n", " update_track_states(online_tracks, track_dict, fps)\n", " show_dets = []\n", " for t in online_tracks:\n", " if not track_dict[t.track_id][\"score_passed\"]:\n", " continue\n", " x1, y1, w, h = t.tlwh\n", " x2 = x1 + w; y2 = y1 + h\n", " mk = track_dict[t.track_id][\"model_key\"]\n", " clbl = track_dict[t.track_id][\"class_label\"]\n", " det_dict = {\"bbox\": [x1, y1, x2, y2], \"score\": t.score, \"label\": clbl,\n", " \"track_id\": t.track_id, \"model_key\": mk}\n", " show_dets.append(det_dict)\n", " vis_frame = frame.copy()\n", " # 모델·클래스 on/off 필터링\n", " final_vis = []\n", " for dd in show_dets:\n", " c_label = dd[\"label\"]\n", " mk = dd[\"model_key\"]\n", " model_conf = MODELS_CONFIG.get(mk, {})\n", " if not model_conf.get(\"class_vis\", {}).get(c_label, True):\n", " continue\n", " # 작업구역 필터링\n", " if RESTRICT_VIZ_TO_POLYGON and (c_label in RESTRICT_VIZ_CLASSES):\n", " bx1, by1, bx2, by2 = dd[\"bbox\"]\n", " cX = (bx1+bx2)/2; cY = (by1+by2)/2\n", " if not is_point_in_polygon((cX, cY), WORKZONE_POLYGON):\n", " continue\n", " final_vis.append(dd)\n", " # bbox 조정: 겹치는 박스 IOU>=0.6 조절\n", " final_vis = adjust_overlapping_boxes(final_vis, iou_thresh=0.6, scale_factor=0.9)\n", " # 헬멧 착용 여부 판단 (whhb 모델 worker의 경우)\n", " if HELMET_WARNING_ENABLED:\n", " for det in final_vis:\n", " if det[\"label\"] == \"worker\" and det.get(\"model_key\") == \"whhb\":\n", " track_id = det.get(\"track_id\")\n", " worker_box = det[\"bbox\"]\n", " max_helmet_iou = 0.0\n", " for other in final_vis:\n", " if other[\"label\"] == \"helmet\" and other.get(\"model_key\") == \"whhb\":\n", " iou_val = iou_calc(worker_box, other[\"bbox\"])\n", " if iou_val > max_helmet_iou:\n", " max_helmet_iou = iou_val\n", " if track_id not in helmet_window:\n", " helmet_window[track_id] = []\n", " helmet_window[track_id].append(max_helmet_iou)\n", " window_length = int(fps * HELMET_WINDOW_DURATION)\n", " if len(helmet_window[track_id]) > window_length:\n", " helmet_window[track_id] = helmet_window[track_id][-window_length:]\n", " avg_iou = sum(helmet_window[track_id]) / len(helmet_window[track_id])\n", " if avg_iou < HELMET_IOU_LOWER:\n", " x1, y1, x2, y2 = map(int, worker_box)\n", " warning_pos = (x1, y1 - 10)\n", " cv2.putText(vis_frame, HELMET_WARNING_TEXT, warning_pos,\n", " cv2.FONT_HERSHEY_SIMPLEX, HELMET_WARNING_FONT_SCALE,\n", " HELMET_WARNING_COLOR, HELMET_WARNING_FONT_THICKNESS)\n", " vis_frame = draw_detections(vis_frame, final_vis)\n", " if resize_factor != 1.0:\n", " oh, ow = vis_frame.shape[:2]\n", " vis_frame = cv2.resize(vis_frame, (int(ow*resize_factor), int(oh*resize_factor)))\n", " if writer:\n", " writer.write(vis_frame)\n", " dt = time.time() - st\n", " frame_times.append(dt)\n", " frame_idx += 1\n", " pbar.update(1)\n", " cap.release()\n", " if writer:\n", " writer.release()\n", " if len(frame_times) > 0:\n", " avg_fps = 1 / np.mean(frame_times)\n", " print(f\"[INFO] 평균 FPS => {avg_fps:.2f}\")\n", " print(f\"[INFO] 결과 => {output_video_path}\")\n", " print(f\"[INFO] 총 처리 시간 => {time.time()-st_total:.1f} s\")\n", " if ENABLE_LOW_CONF_LOGGING and low_conf_logs_video:\n", " log_file_path = os.path.join(LOW_CONF_LOG_DIR, f\"{name}_low_conf_log.txt\")\n", " with open(log_file_path, \"w\", encoding=\"utf-8\") as f:\n", " f.write(f\"Video: {video_path}\\n\")\n", " f.write(\"Frame_idx\\tTime(sec)\\tLabel\\tScore\\tBBox\\tAnalysis\\n\")\n", " for log_item in low_conf_logs_video:\n", " reason = analyze_low_conf_detection(log_item)\n", " f.write(f\"{log_item['frame_idx']}\\t{log_item['time_sec']:.2f}\\t{log_item['label']}\\t{log_item['score']:.2f}\\t{log_item['bbox']}\\t{reason}\\n\")\n", " print(f\"[INFO] 낮은 신뢰도 로그 저장: {log_file_path}\")\n", " rec_file_path = os.path.join(LOW_CONF_LOG_DIR, f\"{name}_recommendations.txt\")\n", " rec_text = generate_recommendations(low_conf_logs_video, name)\n", " with open(rec_file_path, \"w\", encoding=\"utf-8\") as f:\n", " f.write(rec_text)\n", " print(f\"[INFO] 추천 메시지 저장: {rec_file_path}\")\n", "\n", "print(\"===== [Cell 11] 메인 처리 함수 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 12] 실행부(main)\n", "def main():\n", " exts = {\".mp4\", \".avi\", \".mov\", \".mkv\"}\n", " all_files = sorted(os.listdir(VIDEO_DIR))\n", " for f in all_files:\n", " base, ex = os.path.splitext(f)\n", " if ex.lower() in exts:\n", " vp = os.path.join(VIDEO_DIR, f)\n", " process_video(vp, do_inference=True, do_save=True, resize_factor=1.0)\n", " else:\n", " print(f\"[스킵] 영상 파일 아님 => {f}\")\n", "\n", "if __name__==\"__main__\":\n", " print(\"===== [Cell 12] 실행부 start =====\")\n", " main()\n", " print(\"===== [Cell 12] 전체 영상 처리 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 13] 임의 프레임/이미지 디버그\n", "def debug_arbitrary_image(img_path=None):\n", " if not img_path or (not os.path.exists(img_path)):\n", " print(\"[오류] img_path 없음 =>\", img_path)\n", " return\n", " img = cv2.imread(img_path)\n", " if img is None:\n", " print(\"[오류] imread 실패 =>\", img_path)\n", " return\n", " if \"whhb\" not in LOADED_MODELS:\n", " print(\"[스킵] whhb 모델 없음\")\n", " return\n", " accepted, _ = detect_objects_single_model(img, \"whhb\", LOADED_MODELS[\"whhb\"])\n", " debug_out = draw_detections(img, accepted)\n", " debug_path = os.path.join(SAVE_RESULT, \"debug_arbitrary_image.jpg\")\n", " cv2.imwrite(debug_path, debug_out)\n", " print(\"[디버그] 임의 이미지 디텍션 =>\", debug_path)\n", "\n", "print(\"===== [Cell 13] 임의 프레임/이미지 디버그 함수 준비 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 14] 추가 주석/가이드/유의점\n", "\"\"\"\n", "1. 각 모델별 'class_vis' 옵션을 통해, 특정 클래스의 최종 시각화를 on/off 할 수 있습니다.\n", "2. whhb 모델의 경우 고정 색상(FIXED_CLASS_COLORS)을 사용하고, 그 외 모델은 자동 할당됩니다.\n", "3. 최종 시각화 단계에서, bbox 간 IOU가 0.6 이상인 경우 자동으로 박스 크기를 조절하여 겹침을 완화합니다.\n", "4. 헬멧 착용 여부 판단은 whhb 모델의 worker 검출에 대해 슬라이딩 윈도우(기본 3초)를 사용하여, helmet 검출 IOU가 낮으면 경고 문구를 출력합니다.\n", "5. 낮은 신뢰도 검출은 별도로 로그 파일 및 추천 메시지로 저장됩니다.\n", "6. 전체 코드는 셀 단위로 분할되어 주피터 노트북에서 단계별 실행 및 디버깅이 용이하도록 구성되어 있습니다.\n", "\"\"\"\n", "print(\"===== [Cell 14] 추가 가이드 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 15] 최종 확인\n", "def final_check():\n", " print(\"✅ 모든 셀 로직이 연결되었습니다. 필요 시 main() 실행하면 됩니다.\")\n", " print(\"✅ Cell 2 & 3에서 위험구역, 작업구역 관련 변수를 수정 및 디버그 이미지를 확인하세요.\")\n", " print(\"✅ VISUALIZATION_OPTIONS, CLASS_TEXT_POSITION, ZONE_VIS_OPTIONS, 각 모델의 'class_vis' 옵션, bbox 조정 및 헬멧 착용 경고 기능을 확인하세요.\")\n", "\n", "final_check()\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "핵심 요약 및 장기 기억 가이드\n", "핵심 요약:\n", "\n", "모델·클래스별 on/off 옵션: 각 모델 설정에 추가한 \"class_vis\"를 통해 특정 클래스의 최종 시각화를 제어합니다.\n", "\n", "고정 색상 제한: whhb 모델의 경우에만 고정 색상(FIXED_CLASS_COLORS)을 사용하고, 다른 모델은 자동 할당(ASSIGNED_CLASS_COLORS)하여 색상 겹침을 방지합니다.\n", "\n", "겹치는 bbox 조정: final_vis 내에서 IOU가 0.6 이상인 박스에 대해 자동으로 크기를 조절하여 가시성을 개선합니다.\n", "\n", "헬멧 착용 경고: whhb 모델의 worker 검출에 대해 슬라이딩 윈도우(기본 3초)로 helmet 검출 IOU를 분석하고, 평균 IOU가 낮으면 경고 문구를 출력합니다.\n", "\n", "낮은 신뢰도 로그: LOW_CONF_THRESHOLD 이하의 검출 정보를 프레임별로 기록하여 로그 파일과 추천 메시지를 생성합니다.\n", "\n", "장기 기억 가이드:\n", "\n", "코드 구조 및 디버깅:\n", "\n", "코드는 셀 단위로 분할되어 있으며, 각 셀 상단에 해당 기능에 대한 주석이 포함되어 있습니다.\n", "\n", "시각화 옵션 제어:\n", "\n", "전역 변수(VISUALIZATION_OPTIONS, CLASS_TEXT_POSITION, ZONE_VIS_OPTIONS)와 각 모델의 \"class_vis\" 설정을 통해 시각화 세부 옵션을 제어할 수 있습니다.\n", "\n", "색상 할당 정책:\n", "\n", "whhb 모델의 worker, helmet, head는 고정 색상을 사용하고, 그 외 모델은 COLOR_PALETTE를 기반으로 자동 할당됩니다.\n", "\n", "겹치는 박스 조정:\n", "\n", "final_vis에서 겹치는 bbox를 자동으로 조절하여 시각화의 명확성을 높입니다.\n", "\n", "헬멧 착용 판단:\n", "\n", "슬라이딩 윈도우를 통해 worker와 helmet의 IOU를 분석하고, 조건에 따라 경고 메시지를 출력합니다.\n", "\n", "낮은 신뢰도 검출 및 추천:\n", "\n", "낮은 신뢰도 검출 로그를 기록하여, 추후 threshold 조정이나 데이터 보강에 참고할 수 있도록 별도 파일로 저장합니다.\n", "\n", "이 가이드와 요약을 참고하여 향후 코드 개선 및 기능 추가 시 핵심 사항을 관리하시기 바랍니다." ] } ], "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 }