{ "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": "markdown", "metadata": {}, "source": [ "아래는 지금까지의 개선사항과 추가 기능(각 모델·클래스 시각화 on/off, 고정 색상 적용 범위 제한, 겹치는 bbox 자동 조절, 헬멧 미착용 경고 개선 및 상세 분석)을 반영한 최종 수정 코드를 셀 단위로 나눠서 작성한 예시입니다.\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", "이 가이드와 요약을 참고하여 향후 코드 개선 및 기능 추가 시 핵심 사항을 관리하시기 바랍니다." ] }, { "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_현대_테스트할영상_편집본_모음_위험구역_전체테스트_13_하네스클래스몇개끔_윈도우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) 모델 목록 \n", "# 각 모델 설정에 'class_vis' 항목을 추가하여, 해당 모델의 각 클래스별 최종 시각화 여부를 on/off 합니다.\n", "# (예: whhb 모델의 경우 worker, helmet, head는 시각화, background는 off)\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.0, \"stable_time_window\": 3.0, \"required_pass_ratio\": 0.5},\n", " \"helmet\": {\"detect_threshold\": 0.0, \"stable_time_window\": 2.0, \"required_pass_ratio\": 0.3}\n", " },\n", " \"class_vis\": { # whhb는 고정 색상을 사용할 예정 (아래 Cell 5 참고)\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.00, \"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.0, \"stable_time_window\": 3.0, \"required_pass_ratio\": 0.5}\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\": 3.0, \"required_pass_ratio\": 0.4}\n", " },\n", " \"class_vis\": { # 예시로 mixer_truck, excavator는 시각화를 끄도록 설정\n", " \"worker\": False,\n", " \"helmet\": False,\n", " \"mixer_truck\": False,\n", " \"excavator\": False\n", " }\n", " },\n", "}\n", "\n", "# 3) 앙상블 모드: \"union\" 또는 \"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", " (255,0,0), (211,0,148),\n", " (255,255,0), (0,255,255), (255,0,255), (100,100,100),\n", " (128,128,0), (128,0,128), (0,128,128), (128,128,255),\n", " (128,255,128), (255,128,128), (192,192,192), (0,0,128),\n", " (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 이외 모델의 클래스 및 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 # 초 단위 슬라이딩 윈도우 길이\n", "HELMET_IOU_LOWER = 0.1 # 슬라이딩 윈도우 내 평균 IOU 임계값\n", "HELMET_WARNING_ENABLED = True\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", "# [추가] 시각화 옵션 설정\n", "VISUALIZATION_OPTIONS = {\n", " \"draw_bbox\": True,\n", " \"draw_class\": True,\n", " \"draw_confidence\": True,\n", " \"draw_track_id\": False,\n", " \"text_positions\": {\"class\": \"top\", \"confidence\": \"top\", \"track_id\": \"top\"},\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 내 첫 번째 영상을 열어, 영상 정보와 \n", " 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] 디버깅 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 3] 다중 ROI 및 알람 대상 클래스 예시 (주석)\n", "\"\"\"\n", "예시:\n", "# 다중 ROI\n", "POLYGON_COORDS_LIST = [\n", " [(100,100), (200,100), (200,200), (100,200)], # ROI 1\n", " [(300,300), (400,300), (400,400), (300,400)] # ROI 2\n", "]\n", "\n", "# 여러 알람 대상 클래스\n", "TARGET_CLASSES_FOR_ALARM = [\"worker\", \"helmet\"]\n", "\n", "실제 적용 시 process_video() 내에서 위와 같이 반복문으로 처리하면 됩니다.\n", "\"\"\"\n", "print(\"===== [Cell 3] 예시 주석 완료 =====\")\n" ] }, { "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), (1600, 0), (1600, 1080), (0, 1080)]\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\"[디버그] 작업구역 폴리곤 이미지 저장: {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", " plt.xticks([])\n", " plt.yticks([])\n", " plt.grid(True, which='both', linestyle='--', linewidth=0.5)\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", " 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", " cap = cv2.VideoCapture(first_video)\n", " if not cap.isOpened():\n", " print(f\"[오류] 영상 열기 실패 => {first_video}\")\n", " return\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", " 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, 거리, 색상 할당 등)\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", "# 또한, 앙상블 결과(예: \"whhb_fall\")의 경우 별도 색상 (예, 시안)을 사용.\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", " # 앙상블 결과 전용 색상\n", " if model_key == \"whhb_fall\":\n", " return (0, 255, 255) # 시안\n", " # 예를 들어, 추적 보정된 객체는 별도 색상 + 기호 처리 (아래 draw_detections에서 처리)\n", " if \"is_track_corrected\" in cname: # (예시로 cname에 플래그를 넣는 경우)\n", " return (255, 255, 255) # 흰색\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", " return cv2.pointPolygonTest(cnt, point, False) >= 0\n", "\n", "print(\"===== [Cell 5] 헬퍼 함수 완료 =====\")\n" ] }, { "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] 모델 로드 완료 =====\")\n" ] }, { "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), 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", " # 추가: 만약 detection에 \"is_track_corrected\"가 True면, 별도 표시 (예: 작은 원)\n", " for det in detections:\n", " mk = det.get(\"model_key\", \"?\")\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", " # Draw bbox with color\n", " cv2.rectangle(out, (x1, y1), (x2, y2), get_color_for_class(lbl, mk), BBOX_LINE_THICKNESS)\n", " # 만약 detection이 tracking에 의해 보정되었다면 (예: score_history 길이가 1보다 클 경우)\n", " if det.get(\"is_track_corrected\", False):\n", " # Draw a small circle at top-left of bbox\n", " cv2.circle(out, (x1, y1), 5, (255, 255, 255), -1)\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 디버깅 함수 완료 =====\")\n" ] }, { "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; 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, req_ratio, thr = 1.0, 0.3, 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] 앙상블 로직 함수\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(\" [설명] whhb:worker와 fall:worker는 병합, helmet vs head는 높은 가중치 선택\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 11] 메인 처리 함수\n", "# 추가 기능:\n", "# 1. adjust_overlapping_boxes(): final_vis 내 겹치는 bbox(IOU>=0.6) 크기 조절\n", "# 2. 헬멧 착용 경고: whhb 모델 worker에 대해, 슬라이딩 윈도우(3초) 내 helmet IOU 평균이 낮으면 경고.\n", "# - 슬라이딩 윈도우는 buffer가 full 되어야 경고가 발생하도록 수정.\n", "# 3. Zone 시각화: 위험구역 및 작업구역을 출력 (on/off 옵션으로 제어).\n", "\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", " 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", "# 헬멧 경고 알고리즘 설명:\n", "# (1) whhb 모델의 \"worker\" 검출에 대해, 같은 프레임 내 \"helmet\" 검출과의 IOU를 계산합니다.\n", "# (2) 각 track_id별로 최근 HELMET_WINDOW_DURATION(3초) 동안의 최대 helmet IOU 값을 buffer(helmet_window)에 저장합니다.\n", "# (3) buffer 길이가 full이면(>= window_length) 평균 IOU를 계산하고, 평균이 HELMET_IOU_LOWER 이하이면 경고 문구를 출력합니다.\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 = \"비정상적 형태\"\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}건. 검토 필요.\\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, ext = 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", " helmet_window = {} # track_id별 헬멧 IOU sliding window\n", " # Zone 시각화: danger (위험구역)와 detect (작업구역) 출력\n", " # on/off은 ZONE_VIS_OPTIONS에서 제어됨.\n", " st_total = time.time()\n", " print(\"[INFO] 영상 처리 시작 =>\", video_path)\n", " frame_idx = 0\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", " # 추적을 위한 배열 생성\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 xyxys 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, y2 = x1+w, 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, \"last_detection_frame\": frame_idx}\n", " if track_dict[tid][\"class_label\"] is None:\n", " best_iou = 0\n", " 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", " track_dict[tid][\"last_detection_frame\"] = frame_idx\n", " # 만약 이전 frame와 detection이 업데이트되지 않았으면, mark as track corrected\n", " if frame_idx - track_dict[tid].get(\"last_detection_frame\", frame_idx) >= 1:\n", " # 추가 필드: is_track_corrected\n", " track_dict[tid][\"is_track_corrected\"] = True\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, y2 = x1+w, 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,\n", " \"label\": clbl, \"track_id\": t.track_id, \"model_key\": mk}\n", " if track_dict[t.track_id].get(\"is_track_corrected\"):\n", " det_dict[\"is_track_corrected\"] = True\n", " show_dets.append(det_dict)\n", " vis_frame = frame.copy()\n", " \n", " # Zone 시각화: danger zone 및 detect zone (on/off 제어)\n", " if ZONE_VIS_OPTIONS[\"danger\"][\"enabled\"]:\n", " cv2.polylines(vis_frame, np.array([POLYGON_COORDS], dtype=np.int32), True,\n", " ZONE_VIS_OPTIONS[\"danger\"][\"color\"], ZONE_VIS_OPTIONS[\"danger\"][\"line_thickness\"])\n", " if ZONE_VIS_OPTIONS[\"danger\"][\"text_enabled\"]:\n", " x, y, w, h = cv2.boundingRect(np.array([POLYGON_COORDS], dtype=np.int32))\n", " pos = (x + ZONE_VIS_OPTIONS[\"danger\"][\"text_position_offset\"][0], y + ZONE_VIS_OPTIONS[\"danger\"][\"text_position_offset\"][1])\n", " cv2.putText(vis_frame, ZONE_VIS_OPTIONS[\"danger\"][\"text\"], pos,\n", " ZONE_VIS_OPTIONS[\"danger\"][\"font_face\"],\n", " ZONE_VIS_OPTIONS[\"danger\"][\"font_scale\"],\n", " ZONE_VIS_OPTIONS[\"danger\"][\"text_color\"],\n", " ZONE_VIS_OPTIONS[\"danger\"][\"font_thickness\"])\n", " if ZONE_VIS_OPTIONS[\"detect\"][\"enabled\"]:\n", " cv2.polylines(vis_frame, np.array([WORKZONE_POLYGON], dtype=np.int32), True,\n", " ZONE_VIS_OPTIONS[\"detect\"][\"color\"], ZONE_VIS_OPTIONS[\"detect\"][\"line_thickness\"])\n", " if ZONE_VIS_OPTIONS[\"detect\"][\"text_enabled\"]:\n", " x, y, w, h = cv2.boundingRect(np.array([WORKZONE_POLYGON], dtype=np.int32))\n", " pos = (x + ZONE_VIS_OPTIONS[\"detect\"][\"text_position_offset\"][0], y + ZONE_VIS_OPTIONS[\"detect\"][\"text_position_offset\"][1])\n", " cv2.putText(vis_frame, ZONE_VIS_OPTIONS[\"detect\"][\"text\"], pos,\n", " ZONE_VIS_OPTIONS[\"detect\"][\"font_face\"],\n", " ZONE_VIS_OPTIONS[\"detect\"][\"font_scale\"],\n", " ZONE_VIS_OPTIONS[\"detect\"][\"text_color\"],\n", " ZONE_VIS_OPTIONS[\"detect\"][\"font_thickness\"])\n", " # 모델·클래스 on/off 필터링\n", " final_vis = []\n", " for dd in show_dets:\n", " mk = dd[\"model_key\"]\n", " lbl = dd[\"label\"]\n", " model_conf = MODELS_CONFIG.get(mk, {})\n", " if not model_conf.get(\"class_vis\", {}).get(lbl, True):\n", " continue\n", " if RESTRICT_VIZ_TO_POLYGON and (lbl 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 겹침 조정\n", " final_vis = adjust_overlapping_boxes(final_vis, iou_thresh=0.6, scale_factor=0.9)\n", " \n", " # 헬멧 착용 경고 기능 (whhb 모델의 worker 대상)\n", " if HELMET_WARNING_ENABLED:\n", " window_length = int(fps * HELMET_WINDOW_DURATION)\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", " if len(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", " # 추가: 트래킹 보정(detected by tracking, not fresh)인 경우 별도 표시는 draw_detections에서 처리됨.\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", " frame_times.append(time.time()-st)\n", " frame_idx += 1\n", " cap.release()\n", " if writer:\n", " writer.release()\n", " if frame_times:\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", " # 추가: 최종 final_vis 결과의 요약 (whhb:worker, whhb:helmet, harness:harness)\n", " summary = {\"whhb:worker\": 0, \"whhb:helmet\": 0, \"harness:harness\": 0}\n", " for det in final_vis:\n", " mk = det.get(\"model_key\")\n", " lbl = det.get(\"label\")\n", " if mk == \"whhb\" and lbl == \"worker\":\n", " summary[\"whhb:worker\"] += 1\n", " if mk == \"whhb\" and lbl == \"helmet\":\n", " summary[\"whhb:helmet\"] += 1\n", " if mk == \"harness\" and lbl == \"harness\":\n", " summary[\"harness:harness\"] += 1\n", " print(\"[INFO] 최종 시각화 검출 요약:\", summary)\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] 실행부 시작 =====\")\n", " main()\n", " print(\"===== [Cell 12] 전체 영상 처리 완료 =====\")\n" ] }, { "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] 임의 이미지 디버깅 함수 준비 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 14] 추가 주석/가이드 및 유의점\n", "\"\"\"\n", "1. 각 모델별 'class_vis' 옵션을 통해 특정 클래스의 최종 시각화를 on/off 할 수 있습니다.\n", "2. whhb 모델의 경우, worker, helmet, head는 고정 색상을 사용하고, 나머지 모델은 자동 할당되어 색상 겹침을 방지합니다.\n", "3. 최종 시각화 단계에서 IOU>=0.6인 bbox는 자동으로 크기가 조절되어 겹침을 완화합니다.\n", "4. 헬멧 착용 경고는 whhb 모델 worker의 슬라이딩 윈도우(기본 3초)를 기반으로 helmet의 평균 IOU가 HELMET_IOU_LOWER 이하이면 경고 문구를 출력합니다.\n", " - 단, 버퍼 길이가 충분해야 경고가 작동됩니다.\n", "5. 위험구역과 작업구역은 ZONE_VIS_OPTIONS에 따라 on/off되며, 영상 상에 해당 영역이 출력됩니다.\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(\"✅ Zone, 각 모델 'class_vis' 옵션, bbox 조정 및 헬멧 경고 기능을 확인 바랍니다.\")\n", " print(\"✅ 최종 검출 요약은 각 영상 처리 후 콘솔에 출력됩니다.\")\n", "\n", "final_check()\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# 핵심 요약:\n", "1. **모델·클래스별 on/off 옵션:** \n", " - 각 모델 설정에 추가한 `class_vis` 옵션을 통해 특정 클래스(예: whhb 모델의 worker, helmet, head)의 시각화 여부를 제어합니다.\n", "2. **고정 색상 제한:** \n", " - whhb 모델의 경우에 한해 고정 색상(`FIXED_CLASS_COLORS`)을 사용하며, 나머지 모델은 전역 `COLOR_PALETTE`를 기반으로 자동 할당되어 모든 클래스에 대해 서로 다른 색상이 적용됩니다.\n", " - 앙상블 결과(\"whhb_fall\")는 별도의 색상(예: 시안)으로 표시됩니다.\n", "3. **겹치는 bbox 조정:** \n", " - 최종 시각화(final_vis)에서 IOU가 0.6 이상인 bbox는 자동으로 크기를 조절(0.9배)하여 겹침을 완화합니다.\n", " - 또한, tracking으로 이어진 객체(새 검출이 없는 경우)는 추가 표시(예: 작은 원)를 통해 시각적으로 구분합니다.\n", "4. **헬멧 착용 경고:** \n", " - whhb 모델의 worker 검출에 대해, 슬라이딩 윈도우(기본 3초) 내 헬멧 검출과의 IOU를 분석합니다.\n", " - 만약 슬라이딩 윈도우 버퍼가 full 되고 평균 IOU가 설정 임계값(예: 0.1) 이하이면 헬멧 미착용 경고 문구를 출력합니다.\n", " - 경고 문구의 텍스트, 위치, 색상, 폰트 크기 등은 전역 변수로 조정 가능합니다.\n", "5. **낮은 신뢰도 검출 로그:** \n", " - LOW_CONF_THRESHOLD 이하의 검출 정보(영상 이름, 타임라인, 객체 정보 및 간단한 분석 결과)를 로그 파일과 추천 메시지 파일로 저장하여 추후 threshold 조정 및 데이터 보강에 참고할 수 있습니다.\n", "6. **Zone 시각화:** \n", " - 위험 구역(danger)과 작업 구역(detect)은 각각 다른 색상과 텍스트로 영상에 출력되며, on/off 옵션(`ZONE_VIS_OPTIONS`)을 통해 제어할 수 있습니다.\n", "\n", "# 장기 기억 가이드:\n", "- **코드 구조 및 디버깅:** \n", " - 코드는 각 기능별로 셀 단위로 분할되어 있으며, 각 셀 상단에 해당 기능에 대한 자세한 주석과 사용법이 포함되어 있습니다.\n", "- **시각화 옵션 제어:** \n", " - 전역 변수(`VISUALIZATION_OPTIONS`, `CLASS_TEXT_POSITION`, `ZONE_VIS_OPTIONS`, 각 모델의 `class_vis`)를 통해 시각화 세부 사항을 유연하게 제어할 수 있습니다.\n", "- **색상 할당 정책:** \n", " - whhb 모델의 고정 색상과 나머지 모델의 자동 색상 할당으로, 모든 클래스에 대해 서로 다른 색상을 보장하며, 앙상블 결과는 별도 색상으로 표기됩니다.\n", "- **겹치는 박스 조정:** \n", " - 최종 시각화 단계에서 겹치는 bbox는 자동으로 크기가 조절되어 가시성이 개선됩니다.\n", "- **헬멧 착용 판단:** \n", " - 슬라이딩 윈도우를 통해 worker와 helmet의 IOU를 분석, 충분한 버퍼가 채워진 경우에만 경고를 출력하여 과민 반응을 방지합니다.\n", "- **낮은 신뢰도 검출 로그:** \n", " - 낮은 신뢰도 검출은 로그 파일 및 추천 메시지로 저장되어, 후속 조치를 위한 중요한 정보를 제공합니다.\n", "- **요약 정보 출력:** \n", " - 영상 처리 후 최종 검출 결과(예: whhb:worker, whhb:helmet, harness:harness 등)에 대한 요약을 콘솔에 출력합니다.\n", " \n", "현재 내용은 위와 같이 관리하며, 추후 코드 수정 시 이 핵심 요약과 장기 기억 가이드를 참고하시기 바랍니다.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "헬멧 착용 경고 알고리즘 상세 분석\n", "작동 순서:\n", "\n", "whhb 모델로 worker를 검출한 후, 같은 프레임 내에서 해당 worker bbox와 helmet bbox들의 IOU를 계산합니다.\n", "\n", "각 worker의 track_id별로 최대 helmet IOU 값을 슬라이딩 윈도우(3초) 버퍼에 저장합니다.\n", "\n", "버퍼 길이가 full(프레임수 기준)되면 평균 IOU를 계산하며, 그 값이 HELMET_IOU_LOWER(예: 0.1) 미만이면 헬멧 미착용으로 간주하고, 지정된 위치에 경고 문구를 출력합니다.\n", "\n", "현재 문제점:\n", "\n", "헬멧 미착용 경고가 너무 빨리 켜지는 문제는 버퍼가 충분히 채워지기 전에 경고가 발생하거나, threshold 수치가 낮게 설정되어 있기 때문일 수 있습니다.\n", "\n", "개선 방법:\n", "\n", "버퍼 길이가 window_length보다 작으면 경고를 출력하지 않도록 조건을 추가했습니다.\n", "\n", "HELMET_WINDOW_DURATION과 HELMET_IOU_LOWER 수치를 조정하여 더 안정적인 판단을 하도록 구성할 수 있습니다.\n", "\n", "코드 내 설명:\n", "\n", "코드에서는 helmet_window 딕셔너리에 track_id별로 IOU 값을 저장하고, 길이가 window_length 이상이 될 때만 평균을 계산하여 경고 조건을 검사합니다." ] } ], "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 }