{ "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": "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_현대_테스트할영상_편집본_모음_위험구역_전체테스트_06\"\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", "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\": {\n", " \"detect_threshold\": 0.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " },\n", " \"worker\": {\n", " \"detect_threshold\": 0.4, # [수정] 임계값 조정\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " },\n", " \"helmet\": {\n", " \"detect_threshold\": 0.2, # [수정] 임계값 조정\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " }\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\": {\n", " \"detect_threshold\": 0.05,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.2\n", " }\n", " }\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\": {\n", " \"detect_threshold\": 0.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.1\n", " },\n", " \"worker\": {\n", " \"detect_threshold\": 0.3,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " }\n", " }\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\": [\n", " 'worker',\n", " 'signalman',\n", " 'helmet',\n", " 'harness',\n", " 'mixer_truck',\n", " 'excavator',\n", " ],\n", " \"do_inference\": True,\n", " \"vis_on\": True,\n", " \"class_params\": {\n", " \"default\": {\n", " \"detect_threshold\": 0.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.1\n", " },\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", "# [수정] 클래스별 색상 지정 (FIXED_CLASS_COLORS는 고정, 나머지는 순차할당)\n", "COLOR_PALETTE = [\n", " (0,255,0),\n", " (42,42,165),\n", " (0,0,255),\n", " (211,0,148),\n", " (255,0,0),\n", " (255,255,0),\n", " (0,255,255),\n", " (255,0,255),\n", " (100,100,100),\n", " (128,128,0),\n", " (128,0,128),\n", " (0,128,128),\n", " (128,128,255),\n", " (128,255,128),\n", " (255,128,128),\n", " (192,192,192),\n", " (0,0,128),\n", " (128,0,0),\n", " (128,255,255),\n", " (255,128,255)\n", "]\n", "# 고정할 클래스 색상 지정 (worker, helmet, head)\n", "FIXED_CLASS_COLORS = {\n", " \"worker\": (0,255,0),\n", " \"helmet\": (42,42,165),\n", " \"head\": (0,0,255)\n", "}\n", "# 자동 할당용 딕셔너리 (고정 색상이 아닌 클래스에 대해)\n", "ASSIGNED_CLASS_COLORS = {}\n", "\n", "# 7) 앙상블 로직 파라미터\n", "WHHB_FALL_WORKER_IOU_THR = 0.2\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.4 # 예: 0.4 이하이면 로깅 대상\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", "# 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", "# 위험구역(POLYGON_COORDS)과 작업구역(WORKZONE_POLYGON)에 대해 각각 다른 색상 및 텍스트를 지정\n", "ZONE_VIS_OPTIONS = {\n", " \"danger\": { # 위험구역: POLYGON_COORDS\n", " \"enabled\": True,\n", " \"color\": (0, 0, 255), # 빨강 (BGR)\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\": { # 작업구역: WORKZONE_POLYGON (아래 Cell 4에서 정의됨)\n", " \"enabled\": True,\n", " \"color\": (0, 255, 0), # 초록 (BGR)\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) # 우측 상단 기준 오프셋 (x는 음수)\n", " }\n", "}\n", "\n", "print(\"===== [Cell 1] 전역 설정 완료 =====\")\n", "\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] 디버깅(첫 프레임+그리드+폴리곤) 완료 =====\")\n" ] }, { "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)], # 사각형1\n", " [(300,300),(400,300),(400,400),(300,400)] # 사각형2\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", "등으로 수정하면 됩니다.\n", "\"\"\"\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] 작업구역 폴리곤 시각화 및 디버깅 완료 =====\")" ] }, { "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: 고정 클래스(worker, helmet, head)는 FIXED_CLASS_COLORS 사용, 나머지는 COLOR_PALETTE에서 순차 할당\n", "def get_color_for_class(cname):\n", " if 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", " # COLOR_PALETTE에서 이미 고정된 색상을 제외하고 순차 할당\n", " for color in COLOR_PALETTE:\n", " if color not in FIXED_CLASS_COLORS.values() and color not in ASSIGNED_CLASS_COLORS.values():\n", " ASSIGNED_CLASS_COLORS[cname] = color\n", " return color\n", " # 만약 모두 사용되었으면 순환\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", "# [수정] detect_objects_single_model: accepted detections와 낮은 신뢰도(low_conf) 검출 모두 기록 (후자는 로깅용)\n", "def detect_objects_single_model(frame, model_key, model_obj):\n", " accepted = []\n", " low_conf_logs = [] # 낮은 신뢰도 검출 기록 (global threshold 미만)\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", " # accepted 검출: 클래스별 threshold를 만족하면\n", " if sc >= thr:\n", " accepted.append({\n", " \"bbox\": box,\n", " \"score\": sc,\n", " \"label\": cname,\n", " \"model_key\": model_key\n", " })\n", " # 낮은 신뢰도 검출: 전역 LOW_CONF_THRESHOLD 미만이면 로깅 (on/off 옵션)\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", "# 텍스트 시각화 함수 (기존 함수 유지)\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", "\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", "\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", "\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", "\n", " cv2.putText(image, line, org, font_face, font_scale, color, font_thickness)\n", "\n", "# [추가] 새로운 텍스트 시각화 함수: 각 항목별 on/off 및 위치 지정, 겹침 방지\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", " color = get_color_for_class(label)\n", " texts = {}\n", " if vis_opts.get(\"draw_class\", True):\n", " texts[\"class\"] = f'{detection.get(\"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", " \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", " \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", " \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", " if mk in MODELS_CONFIG and not MODELS_CONFIG[mk].get(\"vis_on\", False):\n", " continue\n", " box = det[\"bbox\"]\n", " x1, y1, x2, y2 = map(int, box)\n", " lbl = det.get(\"label\", \"unknown\")\n", " if VISUALIZATION_OPTIONS.get(\"draw_bbox\", True):\n", " cv2.rectangle(out, (x1, y1), (x2, y2), get_color_for_class(lbl), 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", "\n", " if tid not in track_dict:\n", " track_dict[tid] = {\n", " \"last_center\": None,\n", " \"is_filtered\": False,\n", " \"score_history\": [],\n", " \"score_passed\": False,\n", " \"class_label\": None,\n", " \"model_key\": None\n", " }\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", "\n", " if track_dict[tid][\"is_filtered\"]:\n", " continue\n", "\n", " sc = t.score\n", " track_dict[tid][\"score_history\"].append(sc)\n", "\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\n", " req_ratio = 0.3\n", " 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", "\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", "\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", "\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를 만족하면, 두 모델의 score를 가중 평균하여 하나의 worker 객체로 병합합니다.\n", " - helmet vs head 앙상블: 한 쪽이 helmet, 다른 쪽이 head인 경우 IoU>=HELMET_HEAD_IOU_THR를 만족하면, 각 모델의 가중치(HELMET_SCORE_WEIGHT, HEAD_SCORE_WEIGHT)를 곱한 후 높은 쪽만 남깁니다.\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", "\n", " # worker 앙상블: whhb와 fall에서 worker 검출\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 = {\n", " \"bbox\": d1[\"bbox\"],\n", " \"score\": new_sc,\n", " \"label\": \"worker\",\n", " \"model_key\": \"whhb_fall\"\n", " }\n", " used[j] = True\n", " break\n", "\n", " # helmet vs head 앙상블: 두 객체 간 IoU 만족 시 가중치가 높은 쪽만 선택\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는 IoU 조건 만족 시 가중치가 높은 검출만 남깁니다.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 11] 메인 처리 함수\n", "def analyze_low_conf_detection(detection):\n", " \"\"\"\n", " 낮은 신뢰도 검출에 대해 간단한 히스토리 기반 분석을 수행.\n", " - 객체 영역의 크기를 계산하여, 너무 작으면 '작음', 너무 크면 '큼'으로 판별\n", " - 비정상적인 aspect ratio 등도 체크\n", " 검출 항목에 저장된 frame_shape 정보를 사용합니다.\n", " \"\"\"\n", " frame_shape = detection.get(\"frame_shape\")\n", " if frame_shape is None:\n", " frame_area = 1\n", " else:\n", " frame_area = frame_shape[0] * frame_shape[1]\n", " \n", " x1, y1, x2, y2 = map(int, detection[\"bbox\"])\n", " width = x2 - x1\n", " height = y2 - y1\n", " 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", "# 영상별 낮은 신뢰도 로그에 기반한 추천 메시지 생성 함수\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", " recommendations += \"검출 객체의 크기, 형태, 색상 등을 확인하고, 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", "\n", " cap = cv2.VideoCapture(video_path)\n", " if not cap.isOpened():\n", " print(\"[오류] 열기 실패 =>\", video_path)\n", " return\n", "\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", "\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", "\n", " alarm_on = ENABLE_REGION_ALARM\n", " st_total = time.time()\n", "\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", "\n", " all_dets = []\n", " low_conf_dets_all = []\n", " # 각 모델별 검출 수행 (accepted와 낮은 신뢰도 검출 모두 수집)\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", " # 낮은 신뢰도 검출 항목에 현재 프레임 shape, frame_idx, 시간 정보 추가\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", "\n", " low_conf_logs_video.extend(low_conf_dets_all)\n", "\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 len(xyxys) > 0 else np.zeros((0,5), dtype=np.float32)\n", " online_tracks = tracker.update(xyxys_arr, [h_org, w_org])\n", "\n", " for t in online_tracks:\n", " tid = t.track_id\n", " x1, y1, w, h = t.tlwh\n", " x2 = x1 + w\n", " y2 = y1 + h\n", " if tid not in track_dict:\n", " track_dict[tid] = {\n", " \"last_center\": None,\n", " \"is_filtered\": False,\n", " \"score_history\": [],\n", " \"score_passed\": False,\n", " \"class_label\": None,\n", " \"model_key\": None\n", " }\n", " if track_dict[tid][\"class_label\"] is not None:\n", " continue\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\n", " 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", "\n", " update_track_states(online_tracks, track_dict, fps)\n", "\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\n", " y2 = y1 + h\n", " mk = track_dict[t.track_id][\"model_key\"]\n", " clbl = track_dict[t.track_id][\"class_label\"]\n", " show_dets.append({\n", " \"bbox\": [x1, y1, x2, y2],\n", " \"score\": t.score,\n", " \"label\": clbl,\n", " \"track_id\": t.track_id,\n", " \"model_key\": mk\n", " })\n", "\n", " vis_frame = frame.copy()\n", "\n", " # [수정] 기존 알람 처리: final_vis(실제 시각화 대상)에서 TARGET_CLASS_FOR_ALARM을 검사\n", " final_vis = []\n", " for dd in show_dets:\n", " c_label = dd[\"label\"]\n", " # 작업구역 필터링 (Cell 4 on/off)\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\n", " cY = (by1+by2)/2\n", " if not is_point_in_polygon((cX, cY), WORKZONE_POLYGON):\n", " continue\n", " final_vis.append(dd)\n", "\n", " # 위험구역 알람: final_vis에서 TARGET_CLASS_FOR_ALARM 검출 여부 확인\n", " if alarm_on:\n", " cnt = np.array([POLYGON_COORDS], dtype=np.int32)\n", " cv2.polylines(vis_frame, cnt, True, (255,255,0), 2)\n", " alarm_trigger = False\n", " for dd in final_vis:\n", " if dd[\"label\"] == TARGET_CLASS_FOR_ALARM:\n", " bx1, by1, bx2, by2 = dd[\"bbox\"]\n", " cX = (bx1+bx2)/2\n", " cY = (by1+by2)/2\n", " if is_point_in_polygon((cX, cY), POLYGON_COORDS):\n", " alarm_trigger = True\n", " break\n", " if alarm_trigger:\n", " draw_text_with_options(\n", " image=vis_frame,\n", " text_lines=[ALARM_TEXT],\n", " anchor_point=ALARM_POSITION,\n", " text_position=\"bottom\",\n", " font_scale=ALARM_TEXT_SCALE,\n", " font_thickness=ALARM_TEXT_THICKNESS,\n", " color=ALARM_TEXT_COLOR\n", " )\n", "\n", " vis_frame = draw_detections(vis_frame, final_vis)\n", "\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", "\n", " dt = time.time() - st\n", " frame_times.append(dt)\n", " frame_idx += 1\n", " pbar.update(1)\n", "\n", " cap.release()\n", " if writer:\n", " writer.release()\n", "\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", "\n", " # [추가] 낮은 신뢰도 검출 로그 저장 및 추천 파일 생성 (영상별)\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", "\n", " # 추천 메시지 파일 생성\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] 메인 처리 함수 완료 =====\")" ] }, { "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", " \"\"\"\n", " 임의 이미지 경로 지정하여 draw_detections 등을 테스트\n", " \"\"\"\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", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 14] 추가 주석/가이드/유의점\n", "\"\"\"\n", "(1) Polygon이 여러 개라면, POLYGON_COORDS_LIST = [ [...], [...], ... ] 식으로 두고,\n", " is_point_in_polygon() 대신 is_point_in_any_polygon() 호출하면 됩니다.\n", "\n", "(2) 알람 대상이 여러 클래스라면, TARGET_CLASSES_FOR_ALARM = [\"worker\",\"helmet\",\"head\"] 식으로 리스트를 써서,\n", " if dd[\"label\"] in TARGET_CLASSES_FOR_ALARM: ... 로 검사하세요.\n", "\n", "(3) 특정 작업구역(Polygon)에서 시각화:\n", " - RESTRICT_VIZ_TO_POLYGON = True\n", " - RESTRICT_VIZ_CLASSES = [\"worker\",\"helmet\", ...]\n", " - WORKZONE_POLYGON = [...]\n", " 이 로직은 [Cell 11] process_video() 후반부에서 적용.\n", "\n", "(4) GUI 없는 서버 환경이라, \"마우스 클릭으로 ROI 지정\"은 직접 구현이 어렵습니다.\n", " => 수동 좌표 수정 또는 offline GUI 툴에서 polygon 좌표 추출 후 저장.\n", "\n", "(5) harness 모델, do_inference=False로 스킵되지만, 코드에는 남아있어도 문제 없습니다.\n", "\n", "(6) [추가 개선] 시각화 옵션 \n", " - VISUALIZATION_OPTIONS와 CLASS_TEXT_POSITION을 통해, bbox, 클래스, 신뢰도, track id 등의 정보를 개별적으로 on/off 및 위치 조절할 수 있습니다.\n", " - 추가로 ZONE_VIS_OPTIONS를 통해 위험구역과 작업구역의 경계 및 텍스트를 원하는 스타일로 표시할 수 있습니다.\n", "\n", "(7) 낮은 신뢰도 검출(예: score < 0.5)에 대해 영상 이름, 타임라인, 객체명, 분석 결과를 로그 파일로 저장하며,\n", " 해당 결과를 바탕으로 threshold 조정, 데이터 보강 등의 추천 메시지를 함께 출력합니다.\n", "\n", "(8) 출력 파일은 SAVE_RESULT 및 LOW_CONF_LOG_DIR 경로에 저장되며, 파일명에 정보가 포함되도록 하였습니다.\n", "\n", "(9) 전체 코드는 셀 단위로 분할되어 있어, 주피터 노트북에서 단계별 실행 및 디버깅이 용이합니다.\n", "\"\"\"\n", "\n", "print(\"===== [Cell 14] 추가 가이드 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "# %% [Cell 15] 최종 확인\n", "def final_check():\n", " print(\"✅ 모든 셀 로직이 연결되었습니다. 필요 시 main() 실행하면 됩니다.\")\n", " print(\"✅ Cell 2 & 3에서 위험구역, 작업구역 관련 변수를 쉽게 수정하고, 디버그 이미지를 확인하세요.\")\n", " print(\"✅ VISUALIZATION_OPTIONS, CLASS_TEXT_POSITION, ZONE_VIS_OPTIONS 및 낮은 신뢰도 로깅 기능을 통해 세밀한 검출 결과 및 개선점을 확인할 수 있습니다.\")\n", "\n", "final_check()" ] } ], "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 }