{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "아래에는 귀하께서 요청하신 “4가지 모델(whhb, con8, harness, fall)을 앙상블 + 클래스별·모델별 파라미터(Threshold, 윈도우 등) 분리 + 시각화 on/off 및 우선순위 설정” 등을 모두 반영한 통합 예시 코드를 제시합니다.\n", "\n", "개요\n", "주요 개선사항 정리\n", "다중 모델(최대 4개) 동시 사용\n", "\n", "각 모델별로 Inference on/off 가능.\n", "\n", "Inference를 수행한 모델별로 Detection 결과를 얻은 뒤, 하나의 리스트로 합쳐서 ByteTracker에 넘깁니다.\n", "\n", "클래스 이름이 같아도 모델별 파라미터/Threshold를 별도로 설정할 수 있음\n", "\n", "예: whhb 모델의 \"worker\"와 fall 모델의 \"worker\"가 다른 Threshold, 다른 stable_time_window를 가질 수 있음.\n", "\n", "시각화 on/off: 모델별·클래스별로 시각화를 끄고 켤 수 있음\n", "\n", "예: whhb 모델의 \"helmet\"은 시각화하고, harness 모델의 \"helmet\"은 시각화 끄는 식으로 설정 가능.\n", "\n", "동일 클래스에 대한 앙상블(중복 처리) 전략\n", "\n", "\"union\" 모드: 겹치는 박스라도 둘 다 표시(각기 다른 모델 결과를 그대로 모두 시각화)\n", "\n", "\"priority\" 모드: 만약 두 박스가 IoU≥0.5 & 클래스명이 동일하면, 더 높은 confidence 박스를 우선 시각화하고 나머지는 제거\n", "\n", "(원하시면 이 로직을 확장해서, IoU가 작은 건 서로 다른 박스로 표시, IoU가 큰 것은 하나로 합치는 “박스 병합(fusion)” 로직으로 발전시킬 수도 있습니다.)\n", "\n", "전역 변수 한번 수정으로 코드 전체 반영\n", "\n", "[Cell 1]에서 4개 모델의 config·checkpoint·class_names·CLASS_PARAMS·VIS_OPTIONS 등을 모두 등록하면, 아래 셀에서 자동 적용되도록 작성했습니다.\n", "\n", "코드 실행 흐름\n", "[Cell 1] 전역 설정\n", "\n", "4개 모델에 대한 경로, 체크포인트, 클래스 목록, Inference on/off, 파라미터(Threshold, stable_time_window, etc.), 시각화 옵션, 앙상블 모드(\"union\"|\"priority\") 등 모두 설정.\n", "\n", "[Cell 2] 헬퍼 함수 (IoU, distance 등)\n", "\n", "[Cell 3] 모델 로드: 4개 모델 각각 init_detector() 호출 (Inference on/off에 따라 안 불러올 수도 있음)\n", "\n", "[Cell 4] 디텍션 & 시각화 함수\n", "\n", "detect_objects_single_model(): 모델 1개와 frame을 입력받아, 해당 모델/클래스별 Threshold로 필터링한 bbox 목록을 반환.\n", "\n", "ensemble_detections(): 여러 모델의 detection 결과를 하나로 합침(“union” 또는 “priority”).\n", "\n", "시각화 시, 모델이름+클래스 단위의 시각화 on/off도 지원.\n", "\n", "[Cell 5] 여백 추가 함수 (기존 코드와 동일)\n", "\n", "[Cell 6] ByteTracker + 슬라이딩 윈도우 필터링\n", "\n", "update_track_states()에서 track_dict[tid]에 'model_name', 'class_label'를 같이 저장.\n", "\n", "파라미터 조회 시에도 CLASS_PARAMS[model_name][class_label] → 없으면 CLASS_PARAMS[model_name][\"default\"], 그래도 없으면 CLASS_PARAMS[\"global_default\"] 식으로 처리.\n", "\n", "[Cell 7] 메인 영상 처리 함수\n", "\n", "한 프레임마다 4개 모델 각각 detect_objects_single_model() → 합치기(ensemble) → ByteTracker 업데이트 → 슬라이딩 윈도우 → 시각화 후 저장\n", "\n", "[Cell 8] 실행부\n", "\n", "VIDEO_DIR 안의 영상 모두 처리.\n", "\n", "실제로 mp4 결과가 SAVE_RESULT 경로에 저장됨." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 1] 라이브러리 임포트 및 \"다중 모델\" 전역 파라미터 설정\n", "\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 imports\n", "from yolox.tracker.byte_tracker import BYTETracker\n", "from yolox.tracking_utils.timer import Timer\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_exp1' # 결과 저장 서브폴더 (매번 바꿔서 덮어쓰기 방지)\n", "SAVE_RESULT = os.path.join(SAVE_RESULT_BASE, SAVE_SUBFOLDER)\n", "os.makedirs(SAVE_RESULT, exist_ok=True)\n", "\n", "#### 2) 모델 목록 (최대 4개) ####\n", "# - \"model_key\"를 식별자로 사용\n", "# - 각 모델별로:\n", "# \"config\": mmyolo/mmdet config 파일 경로\n", "# \"checkpoint\": weight 파일 경로\n", "# \"class_names\": [클래스명들...]\n", "# \"do_inference\": True/False => 이 모델을 실제로 사용할지 여부\n", "# \"vis_on\": True/False => 이 모델의 detection 결과를 시각화할지 여부\n", "# \"class_params\": {\n", "# \"worker\": {...}, \"helmet\": {...}, ...\n", "# \"default\": {...}\n", "# } => 각 클래스(또는 default)에 대해 threshold, stable_time_window, required_pass_ratio\n", "# - 아래 예시로 4개 모델 등록\n", "\n", "MODELS_CONFIG = {\n", " \"whhb\": {\n", " \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/whh/it1_250326_1819_4class_rota_resize_0p5_1p0_2p0_5p0.py\",\n", " \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/whh/best_coco_bbox_mAP_epoch_7.pth\",\n", " \"class_names\": [\"worker\",\"helmet\",\"head\",\"background\"],\n", " \"do_inference\": True,\n", " \"vis_on\": True, # 이 모델의 결과를 시각화할지 여부\n", " # 모델별 CLASS PARAMS\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.0,\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", " \"background\": {\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.3\n", " },\n", " \"mixer_truck\": {\n", " \"detect_threshold\": 0.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\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", " \"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.0,\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", " \"harness\": {\n", " \"detect_threshold\": 0.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " },\n", " \"signalman\": {\n", " \"detect_threshold\": 0.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " },\n", " \"mixer_truck\": {\n", " \"detect_threshold\": 0.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " },\n", " \"excavator\": {\n", " \"detect_threshold\": 0.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\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, # 예: 일단 off\n", " \"vis_on\": True, # 시각화 off\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.0,\n", " \"stable_time_window\": 1.0,\n", " \"required_pass_ratio\": 0.3\n", " },\n", " \"background\": {\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", "# - \"union\": 여러 모델에서 동일 클래스 & IoU>=0.5인 박스가 나오더라도 다 표시\n", "# - \"priority\": 동일 클래스 & IoU>=0.5 겹칠 경우, confidence가 높은 박스만 살리고 나머지는 제거\n", "ENSEMBLE_STRATEGY = \"union\" # \"union\" or \"priority\"\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", "# - 이 예시에서는 \"모델key: { 클래스명: { show_bbox, show_label, ...}, 'default': {...} }\"\n", "# - 실제 사용시, 각 모델key 아래에 세부 클래스를 추가해 주세요.\n", "VISUALIZATION_OPTIONS = {\n", " \"whhb\": {\n", " \"default\": {\n", " \"show_bbox\": True,\n", " \"show_label\": True,\n", " \"show_score\": True,\n", " \"show_track_id\": False,\n", " \"text_position\": \"top\"\n", " },\n", " # \"helmet\": { ... },\n", " # ...\n", " },\n", " \"con8\": {\n", " \"default\": {\n", " \"show_bbox\": True,\n", " \"show_label\": True,\n", " \"show_score\": True,\n", " \"show_track_id\": False,\n", " \"text_position\": \"top\"\n", " },\n", " },\n", " \"harness\": {\n", " \"default\": {\n", " \"show_bbox\": True,\n", " \"show_label\": True,\n", " \"show_score\": True,\n", " \"show_track_id\": False,\n", " \"text_position\": \"top\"\n", " },\n", " },\n", " \"fall\": {\n", " \"default\": {\n", " \"show_bbox\": True,\n", " \"show_label\": True,\n", " \"show_score\": True,\n", " \"show_track_id\": False,\n", " \"text_position\": \"top\"\n", " },\n", " }\n", "}\n", "\n", "#### 6) 비디오 처리 on/off & 리사이즈 배율 ####\n", "VIDEO_INFERENCE_ON = True\n", "VIDEO_SAVE_ON = True\n", "RESIZE_FACTOR = 1.0\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", "\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", "print(\"===== [Cell 2] 헬퍼 함수 정의 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 3] 여러 모델 로드\n", "LOADED_MODELS = {}\n", "\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", "\n", "def get_class_param(model_key, class_label, param_key):\n", " \"\"\"\n", " model_key와 class_label, param_key(\"detect_threshold\", \"stable_time_window\", \"required_pass_ratio\")로\n", " MODELS_CONFIG[model_key][\"class_params\"]에서 해당 값을 가져오는 헬퍼.\n", " 없는 경우 \"default\" -> 그래도 없으면 None 반환\n", " \"\"\"\n", " if model_key not in MODELS_CONFIG:\n", " return None\n", " cparams = MODELS_CONFIG[model_key][\"class_params\"]\n", " if class_label in cparams:\n", " return cparams[class_label].get(param_key, cparams[\"default\"].get(param_key, None))\n", " else:\n", " # class_label이 없으면 default\n", " if \"default\" in cparams:\n", " return cparams[\"default\"].get(param_key, None)\n", " else:\n", " return None\n", "\n", "def detect_objects_single_model(frame, model_key, model_obj):\n", " \"\"\"\n", " 단일 모델(model_obj)에 대해 Detection 수행,\n", " 모델별 class_params 기반으로 detect_threshold 필터링\n", " Return: list of dict -> [{\n", " \"bbox\": [x1,y1,x2,y2],\n", " \"score\": float,\n", " \"label\": str,\n", " \"model_key\": str # 추가\n", " }]\n", " \"\"\"\n", " # 모델의 클래스 목록\n", " class_names = MODELS_CONFIG[model_key][\"class_names\"]\n", "\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", "\n", " # 모델별+클래스별 threshold 조회\n", " thr = get_class_param(model_key, cname, \"detect_threshold\")\n", " if thr is None:\n", " thr = 0.0 # fallback\n", "\n", " if sc < thr:\n", " continue\n", "\n", " detections.append({\n", " \"bbox\": box,\n", " \"score\": sc,\n", " \"label\": cname,\n", " \"model_key\": model_key\n", " })\n", " return detections\n", "\n", "def ensemble_detections(all_det_list, strategy=\"union\", iou_thr=0.5):\n", " \"\"\"\n", " 여러 모델에서 나온 detection들을 하나의 리스트로 앙상블.\n", " - strategy=\"union\": 중복되어도 모두 사용\n", " - strategy=\"priority\": 같은 class_label & model_key가 다른 detection이 IoU>=iou_thr이면,\n", " score가 더 낮은 쪽은 제거\n", " Return: 최종 detection list\n", " \"\"\"\n", " if strategy == \"union\":\n", " return all_det_list\n", "\n", " # strategy == \"priority\"\n", " # 동일 label & IoU>=0.5 -> score 높은 것만 살림\n", " final_dets = []\n", " used = [False]*len(all_det_list)\n", "\n", " for i in range(len(all_det_list)):\n", " if used[i]:\n", " continue\n", " d1 = all_det_list[i]\n", " box1 = d1[\"bbox\"]\n", " sc1 = d1[\"score\"]\n", " lbl1 = d1[\"label\"]\n", " mk1 = d1[\"model_key\"]\n", "\n", " # d1과 IoU>=0.5 & lbl 동일 -> score 비교\n", " for j in range(i+1, len(all_det_list)):\n", " if used[j]:\n", " continue\n", " d2 = all_det_list[j]\n", " lbl2 = d2[\"label\"]\n", " mk2 = d2[\"model_key\"]\n", " if lbl1 == lbl2: # 클래스 이름이 같을 때\n", " iouv = iou_calc(box1, d2[\"bbox\"])\n", " if iouv >= iou_thr:\n", " # 점수 높은쪽만 살리기\n", " if sc1 >= d2[\"score\"]:\n", " # d2 버림\n", " used[j] = True\n", " else:\n", " # d1 버림\n", " used[i] = True\n", " break\n", "\n", " if not used[i]:\n", " final_dets.append(d1)\n", " return final_dets\n", "\n", "\n", "def draw_text_with_options(\n", " image, text_lines, anchor_point, text_position=\"top\",\n", " font_scale=0.5, font_thickness=1, color=(0,255,0),\n", " auto_adjust=True\n", "):\n", " x_anchor, y_anchor = anchor_point\n", " for i, line in enumerate(text_lines):\n", " text_size, _ = cv2.getTextSize(line, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)\n", " text_w, text_h = text_size\n", " offset = i*(text_h+5)\n", "\n", " if text_position == \"top\":\n", " text_org = (x_anchor, y_anchor - 5 - offset)\n", " elif text_position == \"bottom\":\n", " text_org = (x_anchor, y_anchor + text_h + 5 + offset)\n", " elif text_position == \"left\":\n", " text_org = (x_anchor - text_w - 5, y_anchor + text_h + offset)\n", " else: # \"right\"\n", " text_org = (x_anchor + 5, y_anchor + text_h + offset)\n", "\n", " if auto_adjust:\n", " # 좌/우 경계 보정\n", " if text_org[0] < 0:\n", " text_org = (5, text_org[1])\n", " elif text_org[0] + text_w > image.shape[1]:\n", " text_org = (image.shape[1] - text_w - 5, text_org[1])\n", " # 상/하 경계 보정\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,\n", " cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, font_thickness)\n", "\n", "def draw_detections(frame, detections, font_scale=0.5):\n", " \"\"\"\n", " detections = [\n", " {\"bbox\":..., \"score\":..., \"label\":..., \"track_id\":..., \"model_key\":...},\n", " ...\n", " ]\n", " - 모델별 시각화 옵션: VISUALIZATION_OPTIONS[model_key][class_name or 'default']\n", " - 박스 색상도 model_key와 class_name을 혼합해서 구분 가능\n", " \"\"\"\n", "\n", " output = frame.copy()\n", "\n", " # 간단히 model_key별로 색상 구분 or 클래스별로도 추가\n", " color_map = {\n", " (\"whhb\", \"worker\"): (0, 0, 255),\n", " (\"whhb\", \"helmet\"): (0, 128, 255),\n", " # ... 필요하면 더\n", " # 기본값\n", " (\"default\", \"default\"): (128,128,128)\n", " }\n", "\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\", \"default\")\n", "\n", " # 1) 이 모델의 시각화 on/off 체크\n", " # MODELS_CONFIG[mk][\"vis_on\"] == True 인지?\n", " # 만약 vis_on=False면 표시 X\n", " if not MODELS_CONFIG[mk][\"vis_on\"]:\n", " # 이 모델은 시각화 안 함\n", " continue\n", "\n", " # 2) 이 모델의 시각화 옵션 가져오기\n", " if mk in VISUALIZATION_OPTIONS:\n", " if lbl in VISUALIZATION_OPTIONS[mk]:\n", " vis_opt = VISUALIZATION_OPTIONS[mk][lbl]\n", " else:\n", " vis_opt = VISUALIZATION_OPTIONS[mk][\"default\"]\n", " else:\n", " # 모델이름이 없으면 임시 default\n", " vis_opt = {\n", " \"show_bbox\": True,\n", " \"show_label\": True,\n", " \"show_score\": True,\n", " \"show_track_id\": False,\n", " \"text_position\": \"top\"\n", " }\n", "\n", " # BBox 색상 결정\n", " # 우선 color_map에 (mk,lbl) 키가 있으면 그 색, 없으면 fallback\n", " if (mk, lbl) in color_map:\n", " color = color_map[(mk, lbl)]\n", " else:\n", " color = color_map.get((\"default\",\"default\"), (128,128,128))\n", "\n", " # bbox on/off\n", " if vis_opt[\"show_bbox\"]:\n", " cv2.rectangle(output, (x1,y1), (x2,y2), color, 2)\n", "\n", " # text lines\n", " lines = []\n", " if vis_opt[\"show_label\"]:\n", " lines.append(f\"{mk}:{lbl}\") # 모델명+클래스\n", " if vis_opt[\"show_score\"]:\n", " lines.append(f\"{sc:.2f}\")\n", " if vis_opt[\"show_track_id\"] and (tid is not None):\n", " lines.append(f\"ID:{tid}\")\n", "\n", " if len(lines)>0:\n", " tpos = vis_opt[\"text_position\"]\n", " if tpos in [\"top\",\"left\"]:\n", " anchor=(x1,y1)\n", " elif tpos==\"bottom\":\n", " anchor=(x1,y2)\n", " else: # right\n", " anchor=(x2,y1)\n", "\n", " draw_text_with_options(\n", " image=output,\n", " text_lines=lines,\n", " anchor_point=anchor,\n", " text_position=tpos,\n", " color=color,\n", " font_scale=font_scale,\n", " auto_adjust=True\n", " )\n", "\n", " return output\n", "\n", "\n", "print(\"===== [Cell 4] 디텍션 & 시각화 함수 정의 완료 =====\")\n" ] }, { "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", " # 현재 score 추가\n", " current_score = t.score\n", " track_dict[tid][\"score_history\"].append(current_score)\n", "\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", " # 아직 결정 안 된 상태 => default\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", "\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", "\n", " if ratio >= req_ratio:\n", " track_dict[tid][\"score_passed\"] = True\n", "\n", "print(\"===== [Cell 6] ByteTracker + 슬라이딩 윈도우 필터링 로직 준비 완료 =====\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 7] 메인 처리 함수\n", "def process_video(video_path, do_inference=True, do_save=True, resize_factor=1.0):\n", " if not do_inference:\n", " print(f\"[스킵] 디텍션 off => {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", " if do_save:\n", " fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n", " writer_w = int(w_org * resize_factor)\n", " writer_h = int(h_org * resize_factor)\n", " writer = cv2.VideoWriter(output_video_path, fourcc, fps, (writer_w, writer_h))\n", " if not writer.isOpened():\n", " print(f\"[경고] VideoWriter 생성 실패 => {output_video_path}\")\n", " writer = None\n", " else:\n", " writer=None\n", "\n", " tracker = create_tracker()\n", " track_dict = {}\n", "\n", " frame_times=[]\n", " idx=0\n", "\n", " print(f\"[INFO] 영상 처리 시작 => {video_path}\")\n", " start_time_total = 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, model_obj in LOADED_MODELS.items():\n", " # 이 모델이 do_inference=True인지 이미 [Cell 3]에서 필터링됨\n", " # detect\n", " dets_single = detect_objects_single_model(frame, mk, model_obj)\n", " all_dets.extend(dets_single)\n", "\n", " # (2) 앙상블\n", " final_det_list = ensemble_detections(all_dets, strategy=ENSEMBLE_STRATEGY, iou_thr=0.5)\n", "\n", " # (3) ByteTracker update => [x1,y1,x2,y2,score]\n", " xyxyscore = []\n", " for d in final_det_list:\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) 트랙에 모델key, class_label 할당\n", " # IoU>0.5로 매칭 => track_dict[tid][\"model_key\"], [\"class_label\"]\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", " # 이미 model_key/class_label이 있으면 스킵\n", " if track_dict[tid][\"model_key\"] is not None and track_dict[tid][\"class_label\"] is not None:\n", " continue\n", "\n", " # 없으면 final_det_list중 IoU가 가장 큰 detection 찾아서 할당\n", " best_iou=0\n", " best_det=None\n", " for d in final_det_list:\n", " iouv = iou_calc(track_box, d[\"bbox\"])\n", " if iouv>best_iou:\n", " best_iou=iouv\n", " best_det=d\n", "\n", " # IoU>0.5이면 할당\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_detections=[]\n", " for t in online_tracks:\n", " if not track_dict[t.track_id].get(\"score_passed\",False):\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_detections.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", " # (7) draw\n", " vis_frame = draw_detections(frame, final_detections, font_scale=0.5)\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", " idx+=1\n", " pbar.update(1)\n", "\n", " cap.release()\n", " if writer is not None:\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()-start_time_total:.1f}s\")\n", "\n", "print(\"===== [Cell 7] 메인 영상 처리 함수 정의 완료 =====\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [Cell 8] 실행부\n", "def main():\n", " if not VIDEO_INFERENCE_ON:\n", " print(\"[스킵] VIDEO_INFERENCE_ON=False\")\n", " return\n", "\n", " exts = {'.mp4','.avi','.mov','.mkv'}\n", " files = sorted(os.listdir(VIDEO_DIR))\n", "\n", " for f in files:\n", " base, e = os.path.splitext(f)\n", " if e.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", "main()\n", "print(\"===== [Cell 8] 전체 영상 처리 완료 =====\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "사용 설명 / 주의점\n", "**[Cell 1]**에서 **4개 모델(whhb, con8, harness, fall)**의 설정을 한 번에 정의했습니다.\n", "\n", "do_inference=False인 모델은 로드하지 않으며, 당연히 Detection도 수행하지 않습니다.\n", "\n", "vis_on=False라면 Detection을 하더라도 시각화는 건너뜁니다.\n", "\n", "모델별 클래스 파라미터(class_params)에서 \"worker\", \"helmet\", … 마다 detect_threshold, stable_time_window, required_pass_ratio를 다르게 설정할 수 있습니다.\n", "\n", "예: whhb 모델의 \"worker\"와 harness 모델의 \"worker\" 임계값을 다르게 세팅 가능.\n", "\n", "시각화 on/off:\n", "\n", "각 모델별 vis_on=False이면 전체 detection을 해도 bounding box를 그리지 않습니다.\n", "\n", "또한, 모델 안에서도 특정 클래스 시각화를 세밀히 끄고 싶으면 VISUALIZATION_OPTIONS[\"whhb\"][\"helmet\"] = {\"show_bbox\":False, \"show_label\":False, ...} 처럼 세부 설정 추가 가능.\n", "\n", "클래스 이름이 동일한 경우(예: whhb: 'worker' vs harness: 'worker')\n", "\n", "각자 다른 threshold로 Detection, ByteTracker가 붙은 트랙에도 track_dict[tid][\"model_key\"]가 저장되어, window 파라미터도 서로 다르게 적용됩니다.\n", "\n", "앙상블(ENSEMBLE_STRATEGY)\n", "\n", "\"union\": 겹치는 박스라도 전부 유지 → ByteTracker로 넘어감\n", "\n", "\"priority\": 동일 class & IoU≥0.5이면, 높은 score 박스만 남김. → (주어진 코드에서 “합산” 로직은 없고, 단순 “필터링”만 적용)\n", "\n", "필요시 “합산해서 점수 올리는” 방식은 별도로 로직을 만들어야 합니다. (예: 2개 모델이 같은 박스를 찍으면 +0.1 점수, 등등)\n", "\n", "중간 디버깅\n", "\n", "CLI 환경이라 cv2.imshow() 대신, 필요 시 cv2.imwrite()나 plt.imshow() 등을 활용하여 특정 프레임을 저장·확인하세요.\n", "\n", "각종 확장\n", "\n", "트랙이 사라진 뒤(혹은 오래 lost된 뒤) track_dict에서 제거하는 기능, Box Fusion(NMS) 등은 구현하지 않았습니다. 필요시 추가하시면 됩니다.\n", "\n", "결론\n", "위와 같이 “여러 모델 앙상블 + 모델별/클래스별 파라미터 + 시각화 on/off + IoU기반 우선순위 전략 + 단일 ByteTracker 통합 관리” 구조를 완성했습니다.\n", "\n", "모델별 파라미터/Threshold/시각화를 쉽게 제어할 수 있고,\n", "\n", "동일 클래스라도 모델마다 개별 설정을 가질 수 있으며,\n", "\n", "합쳐진 Detection 결과를 단일 ByteTracker로 추적·슬라이딩 윈도우 필터링합니다.\n", "\n", "추가로 각종 디버깅 로그나 에러 처리, 메모리 관리(track_dict에서 사라진 트랙 제거) 등을 원하시면 동일한 방식으로 확장해 주시면 됩니다." ] } ], "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 }