{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "아래 예시 코드는 이전 앙상블 코드에 다음과 같은 추가 개선사항을 반영한 Jupyter Notebook 형식입니다:\n", "\n", "폰트 관련 개선\n", "\n", "폰트 크기를 크게 설정(기본 15)하고, 폰트 종류를 지정할 수 있는 on/off 옵션 추가.\n", "\n", "폰트 라인 굵기를 1로(기본 설정) 하고, 손쉽게 수정 가능하도록 전역 파라미터화.\n", "\n", "클래스별 색상 지정\n", "\n", "worker → 초록색\n", "\n", "helmet → 갈색\n", "\n", "head → 빨간색\n", "\n", "중장비 8종(con8 등) → 보라색\n", "\n", "그 외 20가지 미리 정의된 색상을 COLOR_PALETTE라는 리스트로 준비(확장 가능).\n", "\n", "원하는 클래스에 매핑하거나, 나중에 수정·확장하기 쉽게 정리.\n", "\n", "앙상블 로직 확장\n", "\n", "(3-A) whhb_worker와 fall_worker가 IoU≥어떤 값(조절 가능)으로 겹치면, 점수를 “가중치 곱한 합산”으로 통합하여 단일 detection으로 만듦.\n", "\n", "예: score = score_whhb * w_whhb + score_fall * w_fall (0~1사이 가중치 곱 가능)\n", "\n", "(3-B) helmet과 head가 IoU≥어떤 값(조절 가능)으로 겹치면, threshold(or score)가 더 높은 쪽만 살림. 또한 각각 0~1 사이 가중치를 곱해 비교 가능.\n", "\n", "겹치지 않는 경우는 기존대로 개별로 표시.\n", "\n", "코드 구조\n", "\n", "셀을 8개 이상으로 나누어, 각 기능별로 최대한 세분화했습니다.\n", "\n", "원본 코드를 주석과 [수정] 표시로 구분해 가며 개선사항을 반영했습니다.\n", "\n", "지침에 따라 출력 경로, 중간 디버깅 방법, 사용 예시 등을 주석으로 자세히 달았습니다.\n", "\n", "목차\n", "[Cell 1] 전역 설정\n", "\n", "[Cell 2] 헬퍼 함수\n", "\n", "[Cell 3] 모델 로드\n", "\n", "[Cell 4] 디텍션 + 시각화 함수 (폰트/색상 개선)\n", "\n", "[Cell 5] 여백 추가 함수\n", "\n", "[Cell 6] ByteTracker + 슬라이딩 윈도우\n", "\n", "[Cell 7] 확장된 앙상블 로직 (whhb_worker+fall_worker, helmet+head)\n", "\n", "[Cell 8] 메인 처리 함수 + 실행부\n", "\n", "아래 전체 코드를 복사해서 하나의 .ipynb로 사용하실 수 있습니다.\n", "(또는 .py 형태로 저장해도 되며, 주피터가 아니어도 문제 없이 동작합니다.)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# # (Ⅰ) 전역 설정 + 라이브러리 로드\n", "# ----------------------------------------------------------\n", "# 본 노트북에서는 4개 모델(whhb, con8, harness, fall)을 동시에\n", "# 앙상블/트래킹 + 클래스별 파라미터/시각화 on/off를 수행합니다.\n", "# 여기에 추가로 폰트 관련(크기, 종류) 및 색상 배치, 그리고\n", "# 특정 모델+클래스 조합 간 앙상블(점수 합산, 우선순위) 로직이 확장되었습니다.\n", "#\n", "# ## 주요 개선사항\n", "# 1) 폰트 크기, 라인 두께, 폰트 종류 on/off\n", "# 2) 클래스별/20가지 색상 팔레트\n", "# 3) whhb_worker & fall_worker => IoU>=? => 점수 합산\n", "# 4) helmet & head => IoU>=? => 더 높은 점수만 표시 (가중치 적용 가능)\n", "#\n", "# 이 노트북 셀을 순차적으로 실행하면 됩니다.\n", "\n", "# %% [Cell 1] 전역 설정\n", "import os\n", "import sys\n", "import time\n", "import glob\n", "import cv2\n", "import numpy as np\n", "from loguru import logger\n", "from types import SimpleNamespace\n", "import mmcv\n", "import gc\n", "\n", "import torch\n", "import torch.nn as nn\n", "import matplotlib.pyplot as plt\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\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_현대_테스트할영상_편집본_모음_final_ensemble'\n", "SAVE_SUBFOLDER = 'ensemble_exp3' # [수정] exp 폴더명 증가\n", "SAVE_RESULT = os.path.join(SAVE_RESULT_BASE, SAVE_SUBFOLDER)\n", "os.makedirs(SAVE_RESULT, exist_ok=True)\n", "\n", "# 2) 모델 목록\n", "# - harness 모델을 예시로 do_inference=False로 설정(원하면 True로)\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.35,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " },\n", " \"helmet\": {\n", " \"detect_threshold\": 0.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " },\n", " \"head\": {\n", " \"detect_threshold\": 0.0,\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\",\"mixer_truck\",\"crane_mobile\",\"cargo_truck\",\"forklift\"],\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.05\n", " }\n", " # 필요시 각 클래스별 항목 추가\n", " },\n", " },\n", " \"harness\": {\n", " \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/dl_1_worker_helmet_harness_signalman_mixertruck_excavator_01/dl_1_worker_helmet_harness_signalman_mixertruck_excavator.py\",\n", " \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/dl_1_worker_helmet_harness_signalman_mixertruck_excavator_01/best_coco_bbox_mAP_epoch_15.pth\",\n", " \"class_names\": [\"worker\",\"helmet\",\"harness\",\"signalman\",\"mixer_truck\",\"excavator\"],\n", " \"do_inference\": False, # 예시로 끔\n", " \"vis_on\": False,\n", " \"class_params\": {}\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.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " }\n", " }\n", " },\n", "}\n", "\n", "# 3) 앙상블 모드\n", "ENSEMBLE_STRATEGY = \"union\" # priority/union는 기본, 새로이 helmet~head, whhb_worker~fall_worker 로직 추가\n", "\n", "# 4) ByteTracker 파라미터\n", "TRACK_THRESH = 0.05\n", "TRACK_BUFFER = 300\n", "MATCH_THRESH = 0.95\n", "tracking_args = SimpleNamespace(\n", " track_thresh=TRACK_THRESH,\n", " track_buffer=TRACK_BUFFER,\n", " match_thresh=MATCH_THRESH\n", ")\n", "\n", "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n", "\n", "# 5) 시각화 관련 - 폰트/색상\n", "# [수정] 폰트 크기, 라인 굵기, 폰트 종류 on/off\n", "FONT_SCALE_DEFAULT = 1 # (!!!) 폰트 크기\n", "FONT_THICKNESS_DEFAULT = 1 # 라인 굵기\n", "USE_CUSTOM_FONT_FACE = True # True면 아래 font_face를 적용, False면 cv2.FONT_HERSHEY_SIMPLEX\n", "CUSTOM_FONT_FACE = cv2.FONT_HERSHEY_DUPLEX # 원하는 cv2 폰트\n", "\n", "# [수정] 클래스별 색상 지정\n", "# - worker = 초록\n", "# - helmet = 갈색\n", "# - head = 빨강\n", "# - 중장비 8종 = 보라\n", "# - 그 외 20가지 color palette 준비\n", "COLOR_PALETTE = [\n", " (0,255,0), # 초록\n", " (42,42,165), # 갈색(bgr로 표현 시 대략..)\n", " (0,0,255), # 빨강\n", " (211,0,148), # 보라\n", " # 이하 확장 예시 (임의로 20개)\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", "\n", "# 6) 비디오 처리 on/off & 리사이즈 배율\n", "VIDEO_INFERENCE_ON = True\n", "VIDEO_SAVE_ON = True\n", "RESIZE_FACTOR = 1.0\n", "\n", "# 7) whhb_worker & fall_worker 앙상블 가중치 + IoU 임계값\n", "WHHB_FALL_WORKER_IOU_THR = 0.3\n", "WHHB_FALL_WORKER_SCORE_WEIGHT_WHHB = 0.7\n", "WHHB_FALL_WORKER_SCORE_WEIGHT_FALL = 0.6\n", "\n", "# 8) helmet & head 앙상블 가중치 + IoU 임계값\n", "HELMET_HEAD_IOU_THR = 0.4\n", "HELMET_SCORE_WEIGHT = 0.8\n", "HEAD_SCORE_WEIGHT = 0.5\n", "\n", "print(\"===== [Cell 1] 전역 설정 완료 =====\")\n", "print(f\"✅ VIDEO_DIR: {VIDEO_DIR}\")\n", "print(f\"✅ SAVE_RESULT: {SAVE_RESULT}\")\n", "print(f\"✅ ENSEMBLE_STRATEGY: {ENSEMBLE_STRATEGY}\")\n", "print(f\"✅ DEVICE: {device}\")\n", "print(f\"✅ FONT_SCALE_DEFAULT: {FONT_SCALE_DEFAULT}\")\n", "print(f\"✅ USE_CUSTOM_FONT_FACE: {USE_CUSTOM_FONT_FACE}\")\n", "print(f\"✅ CUSTOM_FONT_FACE: {CUSTOM_FONT_FACE}\")\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", "\n", "register_all_modules()\n", "init_default_scope('mmdet')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 2] 헬퍼 함수 정의\n", "def iou_calc(a, b):\n", " \"\"\"\n", " 두 bbox (x1,y1,x2,y2) 간 IoU 계산\n", " \"\"\"\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", " \"\"\"\n", " 두 점 (x,y) 사이의 유클리드 거리 계산\n", " \"\"\"\n", " return np.hypot(a[0] - b[0], a[1] - b[1])\n", "\n", "def get_color_for_class(cname):\n", " \"\"\"\n", " [수정] 클래스명에 따라 특정 색상(BGR)을 반환.\n", " - worker=초록, helmet=갈색, head=빨강, 중장비(8종)=보라\n", " - 그 외 COLOR_PALETTE에서 순서대로 반환\n", " \"\"\"\n", " # worker -> 초록\n", " if cname==\"worker\":\n", " return (0,255,0)\n", " # helmet -> 갈색 (대충)\n", " if cname==\"helmet\":\n", " return (42,42,165) # BGR 갈색 계열\n", " # head -> 빨강\n", " if cname==\"head\":\n", " return (0,0,255)\n", " # 중장비 8종 -> 보라\n", " if cname in [\"doger\",\"scissor_lift\",\"evacavator\",\"dump_truck\",\"mixer_truck\",\"crane_mobile\",\"cargo_truck\",\"forklift\",\"excavator\"]:\n", " return (211,0,148) # 보라(bgr)\n", "\n", " # 그 외 -> color palette 순서 (예시)\n", " # 첫 번째 색부터 반환(실 사용시엔 map으로 할당 가능)\n", " return COLOR_PALETTE[0]\n", "\n", "print(\"===== [Cell 2] 헬퍼 함수 정의 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 3] 여러 모델 로드\n", "LOADED_MODELS = {}\n", "for mk, mkinfo in MODELS_CONFIG.items():\n", " if not mkinfo[\"do_inference\"]:\n", " print(f\"[모델 스킵] {mk} -> do_inference=False\")\n", " continue\n", " print(f\"[모델 로드 시작] {mk} ...\")\n", " config_path = mkinfo[\"config\"]\n", " ckpt_path = mkinfo[\"checkpoint\"]\n", " print(f\" - CONFIG: {config_path}\")\n", " print(f\" - CHECKPOINT: {ckpt_path}\")\n", "\n", " model_obj = init_detector(config_path, ckpt_path, device=device)\n", " # 클래스 리스트 부착\n", " model_obj.dataset_meta = {\"CLASSES\": mkinfo[\"class_names\"]}\n", " LOADED_MODELS[mk] = model_obj\n", " print(f\"[모델 로드 완료] {mk}\")\n", "\n", "print(\"===== [Cell 3] 모델 로드 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 4] 디텍션 & 시각화 관련 함수 정의 (폰트/색상 개선)\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", " \"\"\"\n", " 모델별 class_params 기반으로 detect_threshold 적용\n", " \"\"\"\n", " class_names = MODELS_CONFIG[model_key][\"class_names\"]\n", " result = inference_detector(model_obj, frame)\n", " if not hasattr(result, \"pred_instances\"):\n", " return []\n", "\n", " bboxes = result.pred_instances.bboxes\n", " scores = result.pred_instances.scores\n", " labels = result.pred_instances.labels\n", "\n", " detections = []\n", " for i in range(len(bboxes)):\n", " cid = labels[i].item()\n", " if cid >= len(class_names):\n", " continue\n", " cname = class_names[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: thr=0.0\n", " if scimage.shape[1]:\n", " text_org = (image.shape[1]-text_w-5, text_org[1])\n", " if text_org[1]-text_h<0:\n", " text_org = (text_org[0], text_h+5)\n", " elif text_org[1]>image.shape[0]:\n", " text_org = (text_org[0], image.shape[0]-5)\n", "\n", " cv2.putText(image, line, text_org, font_face, font_scale, color, font_thickness)\n", "\n", "\n", "def draw_detections(frame, detections):\n", " \"\"\"\n", " detections = [{\"bbox\":..., \"score\":..., \"label\":..., \"track_id\":..., \"model_key\":...}, ...]\n", " \"\"\"\n", " output = frame.copy()\n", " for det in detections:\n", " x1,y1,x2,y2 = [int(v) for v in det[\"bbox\"]]\n", " sc = det.get(\"score\",0.0)\n", " lbl = det.get(\"label\",\"unknown\")\n", " tid = det.get(\"track_id\",None)\n", " mk = det.get(\"model_key\",\"none\")\n", "\n", " # vis_on 확인\n", " if mk in MODELS_CONFIG:\n", " if not MODELS_CONFIG[mk].get(\"vis_on\",False):\n", " continue\n", "\n", " # 시각화용 색상\n", " color = get_color_for_class(lbl)\n", "\n", " # Bbox\n", " cv2.rectangle(output,(x1,y1),(x2,y2), color, 2)\n", "\n", " # 텍스트\n", " txt_lines=[]\n", " txt_lines.append(f\"{mk}:{lbl}\")\n", " txt_lines.append(f\"{sc:.2f}\")\n", " if tid is not None:\n", " txt_lines.append(f\"ID:{tid}\")\n", "\n", " # anchor\n", " anchor_x, anchor_y = x1, y1\n", " draw_text_with_options(\n", " image=output,\n", " text_lines=txt_lines,\n", " anchor_point=(anchor_x, anchor_y),\n", " text_position=\"top\",\n", " color=color\n", " )\n", "\n", " return output\n", "\n", "\n", "print(\"===== [Cell 4] 디텍션 & 시각화 함수 정의 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 5] 여백 추가 함수\n", "def add_margin_frame(frame, top_margin=0, side_margin=0):\n", " if top_margin<=0 and side_margin<=0:\n", " return frame\n", " h,w = frame.shape[:2]\n", " new_w = w+2*side_margin\n", " new_h = h+top_margin\n", " canvas = np.full((new_h,new_w,3),255,dtype=frame.dtype)\n", " canvas[top_margin:top_margin+h, side_margin:side_margin+w] = frame\n", " return canvas\n", "\n", "print(\"===== [Cell 5] 여백 추가 함수 정의 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 6] ByteTracker + 슬라이딩 윈도우\n", "def create_tracker():\n", " return BYTETracker(tracking_args)\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", " center_now = ((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", "\n", " if track_dict[tid][\"last_center\"] is None:\n", " track_dict[tid][\"last_center\"] = center_now\n", " else:\n", " dist_val = distance(center_now, track_dict[tid][\"last_center\"])\n", " track_dict[tid][\"last_center\"] = center_now\n", "\n", " if track_dict[tid][\"is_filtered\"]:\n", " continue\n", "\n", " current_score = t.score\n", " track_dict[tid][\"score_history\"].append(current_score)\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", " stable_window_size=int(fps*stable_tw)\n", " if stable_window_size<1: stable_window_size=1\n", " hist = track_dict[tid][\"score_history\"]\n", " if len(hist)>stable_window_size:\n", " track_dict[tid][\"score_history\"]=hist[-stable_window_size:]\n", "\n", " pass_count=sum([1 for s in track_dict[tid][\"score_history\"] if s>=thr])\n", " ratio = pass_count/len(track_dict[tid][\"score_history\"]) if len(track_dict[tid][\"score_history\"])>0 else 0\n", " if ratio>=req_ratio:\n", " track_dict[tid][\"score_passed\"] = True\n", "\n", "\n", "print(\"===== [Cell 6] ByteTracker + 슬라이딩 윈도우 로직 준비 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 7] 앙상블 로직 확장 (whhb_worker+fall_worker, helmet+head 등)\n", "def custom_ensemble_detections(dets, iou_thr=0.5):\n", " \"\"\"\n", " [중요 확장 로직]\n", " 1) whhb_worker vs fall_worker -> IoU>=WHHB_FALL_WORKER_IOU_THR 시, 점수 = sum(score*g)로 합쳐 단일 det로\n", " 2) helmet vs head -> IoU>=HELMET_HEAD_IOU_THR 시, 더 높은 (score*g) 만 살림\n", " 3) 그 외에는 ENSEMBLE_STRATEGY=\"union\" 처럼 전부 유지 (혹은 priority?)\n", " \"\"\"\n", "\n", " used = [False]*len(dets)\n", " final_dets = []\n", " for i in range(len(dets)):\n", " if used[i]:\n", " continue\n", " d1 = dets[i]\n", " box1 = d1[\"bbox\"]\n", " sc1 = d1[\"score\"]\n", " lbl1 = d1[\"label\"]\n", " mk1 = d1[\"model_key\"]\n", "\n", " # whhb_worker vs fall_worker 합산\n", " # helmet vs head 중 하나만\n", " # 또는 union\n", " merged_det = None\n", "\n", " for j in range(i+1, len(dets)):\n", " if used[j]:\n", " continue\n", " d2 = dets[j]\n", " box2 = d2[\"bbox\"]\n", " sc2 = d2[\"score\"]\n", " lbl2 = d2[\"label\"]\n", " mk2 = d2[\"model_key\"]\n", "\n", " iouv = iou_calc(box1, box2)\n", "\n", " # (1) whhb_worker & fall_worker\n", " if (lbl1==\"worker\" and mk1==\"whhb\") and (lbl2==\"worker\" and mk2==\"fall\") and iouv>=WHHB_FALL_WORKER_IOU_THR:\n", " # score = s1*g1 + s2*g2\n", " new_score = sc1*WHHB_FALL_WORKER_SCORE_WEIGHT_WHHB + sc2*WHHB_FALL_WORKER_SCORE_WEIGHT_FALL\n", " new_box = box1 # 임의로 box1 채택\n", " # track_id 같은건 아직 없음 => 앙상블중\n", " merged_det = {\n", " \"bbox\": new_box,\n", " \"score\": new_score,\n", " \"label\": \"worker\", # 동일\n", " \"model_key\": \"whhb_fall\" # 임의 식별\n", " }\n", " used[j] = True\n", " break\n", "\n", " # (2) helmet vs head => 더 높은 score*g 만 살림\n", " # if iouv>=HELMET_HEAD_IOU_THR\n", " # (lbl1==\"helmet\" & lbl2==\"head\") or 반대\n", " if ( (lbl1==\"helmet\" and lbl2==\"head\") or (lbl1==\"head\" and lbl2==\"helmet\") ) and iouv>=HELMET_HEAD_IOU_THR:\n", " # weighted score\n", " sc1_w = sc1*HELMET_SCORE_WEIGHT if lbl1==\"helmet\" else sc1*HEAD_SCORE_WEIGHT\n", " sc2_w = sc2*HELMET_SCORE_WEIGHT if lbl2==\"helmet\" else sc2*HEAD_SCORE_WEIGHT\n", " if sc1_w>=sc2_w:\n", " # d1 살리고 d2 버림\n", " used[j] = True\n", " else:\n", " # d2 살리고 d1 버림\n", " used[i] = True\n", " # 이미 처리했으니 break\n", " break\n", "\n", " if merged_det is not None:\n", " # i번 detection이 merged로 변경\n", " final_dets.append(merged_det)\n", " used[i] = True\n", " else:\n", " if not used[i]:\n", " final_dets.append(d1)\n", "\n", " return final_dets\n", "\n", "\n", "print(\"===== [Cell 7] 확장 앙상블 로직 정의 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 8] 메인 처리 함수 + 실행부\n", "def process_video(video_path, do_inference=True, do_save=True, resize_factor=1.0):\n", " if not do_inference:\n", " print(f\"[스킵] => {video_path}\")\n", " return\n", "\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", "\n", " cap = cv2.VideoCapture(video_path)\n", " if not cap.isOpened():\n", " print(f\"[오류] 영상 열기 실패 => {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(f\"[경고] VideoWriter 생성 실패 => {output_video_path}\")\n", " writer=None\n", "\n", " tracker = create_tracker()\n", " track_dict = {}\n", " frame_times=[]\n", "\n", " print(f\"[INFO] 처리 시작 => {video_path}\")\n", " start_t = time.time()\n", "\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", " # (1) 여러 모델 detect\n", " all_dets=[]\n", " for mk,mobj in LOADED_MODELS.items():\n", " det_single = detect_objects_single_model(frame, mk, mobj)\n", " all_dets.extend(det_single)\n", "\n", " # (2) 기존 ensemble (union or priority) -> 생략, 대신 custom_ensemble\n", " final_dets = custom_ensemble_detections(all_dets)\n", "\n", " # (3) ByteTracker update\n", " xyxyscore=[]\n", " for d in final_dets:\n", " x1,y1,x2,y2 = d[\"bbox\"]\n", " s = d[\"score\"]\n", " xyxyscore.append([x1,y1,x2,y2,s])\n", " xyxyscore_arr = np.array(xyxyscore) if len(xyxyscore)>0 else np.zeros((0,5),dtype=np.float32)\n", " online_tracks = tracker.update(xyxyscore_arr, [h_org,w_org])\n", "\n", " # (4) 트랙 dict에 model_key/class_label 할당 (iou>0.5등)\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", " track_box=[x1,y1,x2,y2]\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", "\n", " if track_dict[tid][\"model_key\"] is not None and track_dict[tid][\"class_label\"] is not None:\n", " continue\n", "\n", " best_iou=0\n", " best_det=None\n", " for d in final_dets:\n", " iouv = iou_calc(track_box,d[\"bbox\"])\n", " if iouv>best_iou:\n", " best_iou=iouv\n", " best_det=d\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", " # (5) 슬라이딩 윈도우\n", " update_track_states(online_tracks, track_dict, fps)\n", "\n", " # (6) 필터 통과한 객체 시각화\n", " final_vis=[]\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", " final_vis.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 = draw_detections(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", "\n", " if writer is not None:\n", " writer.write(vis_frame)\n", "\n", " dt=time.time()-st\n", " frame_times.append(dt)\n", " pbar.update(1)\n", "\n", " cap.release()\n", " if writer: writer.release()\n", "\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()-start_t:.1f} s\")\n", "\n", "\n", "def main():\n", " if not VIDEO_INFERENCE_ON:\n", " print(\"[스킵] VIDEO_INFERENCE_ON=False\")\n", " return\n", " exts={\".mp4\",\".avi\",\".mov\",\".mkv\"}\n", " fs=sorted(os.listdir(VIDEO_DIR))\n", " for f in fs:\n", " base,ex = os.path.splitext(f)\n", " if ex.lower() in exts:\n", " vp = os.path.join(VIDEO_DIR,f)\n", " process_video(\n", " video_path=vp,\n", " do_inference=VIDEO_INFERENCE_ON,\n", " do_save=VIDEO_SAVE_ON,\n", " resize_factor=RESIZE_FACTOR\n", " )\n", " else:\n", " print(f\"[스킵] 영상 아님 => {f}\")\n", "\n", "if __name__==\"__main__\":\n", " main()\n", " print(\"===== [Cell 8] 전체 영상 처리 완료 =====\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "추가 설명\n", "폰트 관련\n", "\n", "전역 변수 FONT_SCALE_DEFAULT = 15, FONT_THICKNESS_DEFAULT = 1를 사용합니다.\n", "\n", "USE_CUSTOM_FONT_FACE = True라면 CUSTOM_FONT_FACE(예: cv2.FONT_HERSHEY_DUPLEX)를 사용합니다.\n", "\n", "USE_CUSTOM_FONT_FACE = False로 바꾸면 기본 폰트(cv2.FONT_HERSHEY_SIMPLEX)가 적용됩니다.\n", "\n", "draw_text_with_options() 내부에서 이 전역 설정을 참조하므로, 한 번만 수정하면 모든 시각화에 반영됩니다.\n", "\n", "클래스별 색상\n", "\n", "함수 get_color_for_class(cname) 내부에서,\n", "\n", "worker → 초록, helmet → 갈색, head → 빨강,\n", "\n", "8종(중장비) → 보라,\n", "\n", "그 외는 COLOR_PALETTE[0]을 사용(필요시 인덱스를 달리하거나, dict 매핑으로 확장).\n", "\n", "COLOR_PALETTE에 20개 정도 RGB를 (B,G,R) 형식으로 작성해 두었습니다. 필요 시 원하는 컬러로 변경/확장해주세요.\n", "\n", "앙상블 로직 (custom_ensemble_detections)\n", "\n", "입력: dets(모든 모델에서 모은 detection)\n", "\n", "내부에서 특수 규칙:\n", "\n", "whhb, label=worker 와 fall, label=worker가 IoU≥ WHHB_FALL_WORKER_IOU_THR →\n", "점수를 score = score1*g1 + score2*g2로 합산 후 단일 detection으로 병합.\n", "\n", "helmet vs head가 IoU≥ HELMET_HEAD_IOU_THR → 가중치를 곱해 비교 후 더 큰 쪽만 유지\n", "\n", "그 외 → union(겹치면 그냥 둘 다 표시)\n", "\n", "이 로직은 예시이며, 실제 활용 시 필요에 맞게 수정 가능합니다. 예: box 병합 시 bbox 위치를 중간점으로 정한다든가, score를 normalize한다든가.\n", "\n", "주의/확장\n", "\n", "“트랙 사라짐” 후 track_dict에서 제거, “합산 점수에 대한 안정성”, “박스 실제 위치” 등은 아직 단순 버전입니다.\n", "\n", "COLOR_PALETTE도 임의로 20개를 넣었습니다. BGR 순서라 “갈색, 보라”가 약간 어긋날 수 있으니, 실제 색감은 테스트하며 조정하세요.\n", "\n", "helmet vs head 로직과 whhb_worker vs fall_worker 로직이 동시에 적용되더라도, 호출 순서에 따라 중복 처리가 있을 수 있으니, 필요하면 더욱 정교한 로직을 도입해야 합니다(현재 코드에서는 i