{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "아래에는 귀하께서 요청하신 “슬라이딩 윈도우 기반 점수 필터링(객체 트래킹) + 클래스별 파라미터 제어 + 시각화 옵션(on/off, 위치 조정) + 기타 코드 구조 개선” 등을 모두 반영한 예시 코드 전체를 제시합니다.\n",
    "\n",
    "목차\n",
    "사용 및 예시 안내\n",
    "\n",
    "전체 Jupyter Notebook 예시 코드\n",
    "\n",
    "[Cell 1] 전역 설정\n",
    "\n",
    "[Cell 2] 헬퍼 함수 (IoU, distance 등)\n",
    "\n",
    "[Cell 3] 디텍션 모델 로드\n",
    "\n",
    "[Cell 4] 디텍션 및 시각화 함수 정의\n",
    "\n",
    "[Cell 5] 여백 추가 함수\n",
    "\n",
    "[Cell 6] ByteTracker + 슬라이딩 윈도우 기반 필터링 로직\n",
    "\n",
    "[Cell 7] 메인 영상 처리 함수\n",
    "\n",
    "[Cell 8] 실행부\n",
    "\n",
    "코드 내에 [수정] 주석을 달아, 기존 코드 대비 어떤 부분이 보강되었는지 표시했습니다. 또한, 중간중간 주석과 print를 통해 파라미터, 경로, 디버깅 정보를 추가로 확인할 수 있도록 했습니다.\n",
    "\n",
    "1. 사용 및 예시 안내\n",
    "아래는 주요 전역 변수 및 하이퍼파라미터 사용 예시를 정리한 것입니다.\n",
    "\n",
    "CLASS_PARAMS: 각 클래스별 슬라이딩 윈도우 파라미터(윈도우 크기, pass_ratio 등) 및 Detection Threshold를 딕셔너리로 관리합니다.\n",
    "\n",
    "VISUALIZATION_OPTIONS: 시각화(뷰) 옵션을 전역 또는 클래스별로 관리합니다. (예: BBox 표시 여부, Label 표시 여부, 점수 표시 여부, TrackID 표시 여부, 텍스트를 BBox 위/아래/왼쪽/오른쪽 어디에 그릴지 등)\n",
    "\n",
    "특정 클래스만 다르게 하고 싶으면, VISUALIZATION_OPTIONS[\"클래스이름\"] = {...} 식으로 추가·변경합니다.\n",
    "\n",
    "미지정 시 VISUALIZATION_OPTIONS[\"default\"] 값을 사용합니다.\n",
    "\n",
    "AUTO_ADJUST_TEXT_POSITION: 텍스트가 프레임 바깥으로 나가는 경우 위치를 자동 조정할지 여부를 True/False로 설정합니다.\n",
    "\n",
    "파라미터/옵션 사용 예시\n",
    "python\n",
    "복사\n",
    "편집\n",
    "# 예) worker 클래스만 별도로 슬라이딩 윈도우 pass ratio를 0.5, 윈도우 크기도 2.0초로 설정\n",
    "CLASS_PARAMS[\"worker\"][\"stable_time_window\"] = 2.0\n",
    "CLASS_PARAMS[\"worker\"][\"required_pass_ratio\"] = 0.5\n",
    "\n",
    "# 예) helmet 클래스만 detect threshold를 0.3으로 조정\n",
    "CLASS_PARAMS[\"helmet\"][\"detect_threshold\"] = 0.3\n",
    "\n",
    "# 예) visualization off 예시 (helmet 클래스에 대해서만 label과 score를 숨긴다)\n",
    "VISUALIZATION_OPTIONS[\"helmet\"][\"show_label\"] = False\n",
    "VISUALIZATION_OPTIONS[\"helmet\"][\"show_score\"] = False\n",
    "\n",
    "# 예) 모든 클래스의 텍스트 표시 위치를 bbox 상단이 아니라 \"오른쪽\"으로 통일한다\n",
    "VISUALIZATION_OPTIONS[\"default\"][\"text_position\"] = \"right\"\n",
    "이처럼 전역 딕셔너리를 수정함으로써, 클래스별(또는 전역) 파라미터를 조정할 수 있습니다.\n",
    "\n",
    "코드 전체 개요\n",
    "**[Cell 1]**에서 필요한 라이브러리를 import하고, 경로 및 전역 파라미터를 설정합니다.\n",
    "\n",
    "**[Cell 2]**에서 IoU, distance 등을 계산하는 헬퍼 함수를 정의합니다.\n",
    "\n",
    "**[Cell 3]**에서 mmdetection 모델(config & weight) 로드.\n",
    "\n",
    "**[Cell 4]**에서 디텍션 후 시각화(박스 그리기, 텍스트 on/off, 위치 조절, 중복 방지 등) 함수를 정의합니다.\n",
    "\n",
    "**[Cell 5]**에서 필요 시 여백을 추가하는 함수를 정의합니다. (기존 코드와 동일)\n",
    "\n",
    "**[Cell 6]**에서 ByteTracker 및 슬라이딩 윈도우 필터링 로직을 정의합니다.\n",
    "\n",
    "**[Cell 7]**에서 메인 처리 함수(process_video)를 구현합니다.\n",
    "\n",
    "**[Cell 8]**에서 실제로 영상을 순회하며 최종 추론 및 저장을 수행합니다.\n",
    "\n",
    "디버깅 시 유의점\n",
    "GUI 환경이 아니므로, cv2.imshow() 대신 cv2.imwrite()로 저장하거나 matplotlib.pyplot으로 시각화합니다.\n",
    "\n",
    "AUTO_ADJUST_TEXT_POSITION = True이면, 텍스트가 프레임 밖으로 나가는 경우 자동으로 위치를 조정합니다.\n",
    "\n",
    "슬라이딩 윈도우 로직 상, track별로 \"현재까지 누적된 detection score\"를 관리하고 일정 횟수 이상 threshold를 넘으면 score_passed=True가 됩니다.\n",
    "\n",
    "각 클래스별 서로 다른 stable_time_window, required_pass_ratio도 설정 가능하지만, 코드가 다소 복잡해집니다(트랙 생성 시 어떤 클래스인지 추론해 기록해야 함). 이 예시 코드에서는 동일 트랙에서 여러 클래스가 섞이지 않는 정상 상황을 가정하여, 한 번 붙은 라벨은 계속 동일하다고 처리합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------------------------------------------------------------\n",
    "# [Cell 1] 라이브러리 임포트 및 전역 설정\n",
    "# --------------------------------------------------------------------------------\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 관련 임포트 (기존 코드 유지) ----\n",
    "from yolox.tracker.byte_tracker import BYTETracker\n",
    "from yolox.tracking_utils.timer import Timer\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')\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_현대_테스트할영상_중장비_con8_01'\n",
    "SAVE_SUBFOLDER = 'whhb_exp'  # [수정] exp 폴더명 등 변경 가능\n",
    "SAVE_RESULT = os.path.join(SAVE_RESULT_BASE, SAVE_SUBFOLDER)\n",
    "os.makedirs(SAVE_RESULT, exist_ok=True)\n",
    "\n",
    "# (2) 최종 클래스 목록 & 클래스별 Detection Threshold / 슬라이딩 윈도우 파라미터\n",
    "class_names_final = [\n",
    "    'doger',\n",
    "    'scissor_lift',\n",
    "    'evacavator',\n",
    "    'dump_truck',\n",
    "    'mixer_truck',\n",
    "    'crane_mobile',\n",
    "    'cargo_truck',\n",
    "    'forklift',\n",
    "]\n",
    "\n",
    "# ---- [수정/추가] 클래스별 파라미터 딕셔너리 ----\n",
    "#  - detect_threshold      : 디텍션 단계에서의 score 임계값\n",
    "#  - stable_time_window    : 슬라이딩 윈도우 크기(초)\n",
    "#  - required_pass_ratio   : 윈도우 내 score>=threshold가 차지하는 비율\n",
    "# ※ default 항목은 이 딕셔너리에 없는 클래스가 참조할 값\n",
    "CLASS_PARAMS = {\n",
    "    \"default\": {\n",
    "        \"detect_threshold\": 0.0,\n",
    "        \"stable_time_window\": 1.0,\n",
    "        \"required_pass_ratio\": 0.3\n",
    "    },\n",
    "    \"doger\": {\n",
    "        \"detect_threshold\": 0.0,\n",
    "        \"stable_time_window\": 1.0,\n",
    "        \"required_pass_ratio\": 0.3\n",
    "    },\n",
    "    \"scissor_lift\": {\n",
    "        \"detect_threshold\": 0.0,\n",
    "        \"stable_time_window\": 1.0,\n",
    "        \"required_pass_ratio\": 0.3\n",
    "    },\n",
    "    \"evacavator\": {\n",
    "        \"detect_threshold\": 0.0,\n",
    "        \"stable_time_window\": 1.0,\n",
    "        \"required_pass_ratio\": 0.3\n",
    "    },\n",
    "    \"dump_truck\": {\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",
    "    \"crane_mobile\": {\n",
    "        \"detect_threshold\": 0.0,\n",
    "        \"stable_time_window\": 1.0,\n",
    "        \"required_pass_ratio\": 0.3\n",
    "    },\n",
    "    \"cargo_truck\": {\n",
    "        \"detect_threshold\": 0.0,\n",
    "        \"stable_time_window\": 1.0,\n",
    "        \"required_pass_ratio\": 0.3\n",
    "    },\n",
    "    \"forklift\": {\n",
    "        \"detect_threshold\": 0.0,\n",
    "        \"stable_time_window\": 1.0,\n",
    "        \"required_pass_ratio\": 0.3\n",
    "    },\n",
    "}\n",
    "\n",
    "# (3) 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",
    "# (4) 시각화 옵션\n",
    "#  - show_bbox       : BBox 테두리 표시 여부\n",
    "#  - show_label      : 클래스 라벨 텍스트 표시 여부\n",
    "#  - show_score      : 점수 표시 여부\n",
    "#  - show_track_id   : 트랙 ID 표시 여부\n",
    "#  - text_position   : 텍스트 위치(top, bottom, left, right)\n",
    "# ※ 특정 클래스만 다르게 하고 싶다면, \"helmet\", \"worker\" 등으로 key를 넣고 별도 설정 가능\n",
    "VISUALIZATION_OPTIONS = {\n",
    "    \"default\": {\n",
    "        \"show_bbox\": True,\n",
    "        \"show_label\": True,\n",
    "        \"show_score\": True,\n",
    "        \"show_track_id\": False,\n",
    "        \"text_position\": \"top\"  # top | bottom | left | right\n",
    "    }\n",
    "    # 예) \"helmet\": {\"show_label\": False, \"show_score\": False, ...} 등으로 추가 가능\n",
    "}\n",
    "\n",
    "# (5) 텍스트가 화면 밖으로 나갈 경우 자동 조정 (True/False)\n",
    "AUTO_ADJUST_TEXT_POSITION = True\n",
    "\n",
    "print(\"===== [Cell 1] 전역 설정 완료 =====\")\n",
    "print(f\"✅ BASE_DIR: {BASE_DIR}\")\n",
    "print(f\"✅ VIDEO_DIR: {VIDEO_DIR}\")\n",
    "print(f\"✅ SAVE_RESULT: {SAVE_RESULT}\")\n",
    "print(f\"✅ DEVICE: {device}\")\n",
    "print(f\"✅ CLASS_PARAMS: {CLASS_PARAMS}\")\n",
    "print(f\"✅ VISUALIZATION_OPTIONS: {VISUALIZATION_OPTIONS}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------------------------------------------------------------\n",
    "# [Cell 2] 헬퍼 함수 정의 (IoU, distance 등)\n",
    "# --------------------------------------------------------------------------------\n",
    "\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\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] 헬퍼 함수 정의 완료 =====\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------------------------------------------------------------\n",
    "# [Cell 3] 디텍션 모델 로드\n",
    "# --------------------------------------------------------------------------------\n",
    "\n",
    "# *** [수정] 기존 코드 유지, config/weight 경로는 현재 6개 클래스용(통합 모델) ***\n",
    "config_file_final = '/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/con8/yolo8_large_960size_con_equip_8class.py'\n",
    "checkpoint_file_final = '/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_현대_영상용_weight모음/con8/best_coco_bbox_mAP_epoch_12.pth'\n",
    "\n",
    "det_model_final = init_detector(config_file_final, checkpoint_file_final, device=device)\n",
    "det_model_final.dataset_meta = {'CLASSES': class_names_final}\n",
    "\n",
    "print(\"===== [Cell 3] 모델 로드 완료 =====\")\n",
    "print(f\"✅ Using config:    {config_file_final}\")\n",
    "print(f\"✅ Using checkpoint:{checkpoint_file_final}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------------------------------------------------------------\n",
    "# [Cell 4] 디텍션 및 시각화 함수 정의\n",
    "# --------------------------------------------------------------------------------\n",
    "\n",
    "def get_class_param(label_name, param_key):\n",
    "    \"\"\"\n",
    "    label_name(문자열)과 param_key(\"detect_threshold\", \"stable_time_window\", \"required_pass_ratio\")로\n",
    "    CLASS_PARAMS 딕셔너리에서 해당 값을 가져오는 헬퍼 함수.\n",
    "    만약 label_name에 해당하는 키가 없으면 'default'를 사용.\n",
    "    \"\"\"\n",
    "    if label_name in CLASS_PARAMS:\n",
    "        return CLASS_PARAMS[label_name].get(param_key, CLASS_PARAMS[\"default\"][param_key])\n",
    "    else:\n",
    "        return CLASS_PARAMS[\"default\"][param_key]\n",
    "\n",
    "def detect_objects(frame, model, class_names):\n",
    "    \"\"\"\n",
    "    주어진 frame(이미지)과 모델, 클래스 목록을 이용해 디텍션 수행 후,\n",
    "    CLASS_PARAMS에 정의된 detect_threshold를 기준으로 필터링된 결과를 반환한다.\n",
    "    \"\"\"\n",
    "    result = inference_detector(model, 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",
    "    detections = []\n",
    "\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",
    "        # 클래스별 Detection Threshold 조회\n",
    "        thr = get_class_param(cname, \"detect_threshold\")\n",
    "        if sc < thr:\n",
    "            continue\n",
    "\n",
    "        detections.append({\n",
    "            'bbox': box,\n",
    "            'score': sc,\n",
    "            'label': cname\n",
    "        })\n",
    "    return detections\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",
    "    \"\"\"\n",
    "    [수정/추가] 텍스트를 BBox 주변에 그리는 함수.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    image : np.ndarray\n",
    "        실제 프레임 이미지 (BGR).\n",
    "    text_lines : list of str\n",
    "        표시할 텍스트 라인들 (예: [\"worker\", \"0.85\", \"ID: 3\"]).\n",
    "    anchor_point : (x, y)\n",
    "        BBox의 기준점(왼상단, 왼하단, 오른상단, 오른하단 등).\n",
    "    text_position : str\n",
    "        'top', 'bottom', 'left', 'right' 중 택일.\n",
    "    font_scale : float\n",
    "        글씨 크기.\n",
    "    font_thickness : int\n",
    "        글씨 두께.\n",
    "    color : (B, G, R)\n",
    "        텍스트 색상.\n",
    "    auto_adjust : bool\n",
    "        텍스트가 이미지 밖으로 나갈 경우 위치 자동 조정 여부.\n",
    "    \"\"\"\n",
    "    x_anchor, y_anchor = anchor_point\n",
    "    for i, line in enumerate(text_lines):\n",
    "        # 한 줄에 대한 사이즈 계산\n",
    "        text_size, _ = cv2.getTextSize(line, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)\n",
    "        text_w, text_h = text_size\n",
    "\n",
    "        offset = (i * (text_h + 5))  # 라인 간 간격\n",
    "\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",
    "        # 화면 밖으로 벗어날 경우 조정(옵션)\n",
    "        if auto_adjust:\n",
    "            # 가로 조정\n",
    "            if text_org[0] < 0:\n",
    "                # 왼쪽 밖으로 나간다면 오른쪽으로 이동\n",
    "                text_org = (5, text_org[1])\n",
    "            elif text_org[0] + text_w > image.shape[1]:\n",
    "                # 오른쪽 밖으로 나간다면 왼쪽으로 이동\n",
    "                text_org = (image.shape[1] - text_w - 5, text_org[1])\n",
    "\n",
    "            # 세로 조정\n",
    "            if text_org[1] - text_h < 0:\n",
    "                # 위로 나가면 아래로\n",
    "                text_org = (text_org[0], text_h + 5)\n",
    "            elif text_org[1] > image.shape[0]:\n",
    "                # 아래로 나가면 위로\n",
    "                text_org = (text_org[0], image.shape[0] - 5)\n",
    "\n",
    "        # 실제 텍스트 그리기\n",
    "        cv2.putText(\n",
    "            image, line, text_org,\n",
    "            cv2.FONT_HERSHEY_SIMPLEX, font_scale,\n",
    "            color, font_thickness\n",
    "        )\n",
    "\n",
    "\n",
    "def draw_detections(frame, detections, font_scale=0.5):\n",
    "    \"\"\"\n",
    "    detections: [{'bbox':[x1,y1,x2,y2], 'score':..., 'label':..., 'source':..., 'track_id':...}, ...]\n",
    "    \n",
    "    [수정/추가]\n",
    "    - 각 클래스별 혹은 디폴트 시각화 옵션(VISUALIZATION_OPTIONS)을 참조하여\n",
    "      BBox, Label, Score, TrackID 등을 그릴지 말지 결정.\n",
    "    - text_position도 클래스별로 다를 수 있도록 수정.\n",
    "    - 텍스트가 프레임 밖으로 나가지 않도록 AUTO_ADJUST_TEXT_POSITION 옵션 활용.\n",
    "    \"\"\"\n",
    "    output = frame.copy()\n",
    "\n",
    "    # 예시 color map (BGR)\n",
    "    color_map = {\n",
    "        'worker': (0, 0, 255),         # 빨강\n",
    "        'head': (255, 0, 255),          # 마젠타\n",
    "        'doger': (0, 255, 0),           # 초록\n",
    "        'scissor_lift': (0, 255, 255),  # 노랑\n",
    "        'evacavator': (255, 0, 0),      # 파랑\n",
    "        'dump_truck': (165, 0, 165),     # 보라\n",
    "        'mixer_truck': (0, 0, 128),     # 짙은 파랑\n",
    "        'crane_mobile': (128, 128, 128),# 회색\n",
    "        'cargo_truck': (255, 255, 0),    # 연노랑 (청색 계열)\n",
    "        'forklift': (0, 165, 255),       # 주황\n",
    "        'helmet': (0, 128, 255),          # 살짝 다르게 주황계\n",
    "        'default': (128, 128, 128)    # 회색\n",
    "    }\n",
    "\n",
    "    for det in detections:\n",
    "        x1, y1, x2, y2 = [int(x) for x in det['bbox']]\n",
    "        sc = det.get('score', 0.0)\n",
    "        lbl = det.get('label', 'unknown')\n",
    "        track_id = det.get('track_id', None)\n",
    "\n",
    "        # ---- [수정/추가] 클래스별 시각화 옵션 ----\n",
    "        if lbl in VISUALIZATION_OPTIONS:\n",
    "            vis_opt = VISUALIZATION_OPTIONS[lbl]\n",
    "        else:\n",
    "            vis_opt = VISUALIZATION_OPTIONS[\"default\"]\n",
    "\n",
    "        # BBox 색상\n",
    "        color = color_map.get(lbl, color_map['default'])\n",
    "\n",
    "        # 1) show_bbox\n",
    "        if vis_opt[\"show_bbox\"]:\n",
    "            cv2.rectangle(output, (x1, y1), (x2, y2), color, 2)\n",
    "\n",
    "        # 2) 텍스트 표시\n",
    "        text_lines = []\n",
    "        if vis_opt[\"show_label\"]:\n",
    "            text_lines.append(f\"{lbl}\")\n",
    "        if vis_opt[\"show_score\"]:\n",
    "            text_lines.append(f\"{sc:.2f}\")\n",
    "        if vis_opt[\"show_track_id\"] and (track_id is not None):\n",
    "            text_lines.append(f\"ID:{track_id}\")\n",
    "\n",
    "        # 텍스트가 전혀 없다면 draw_text 패스\n",
    "        if len(text_lines) > 0:\n",
    "            # text_position\n",
    "            tpos = vis_opt[\"text_position\"]\n",
    "            # BBox 기준점(anchor)은 top/left 모서리 쪽으로\n",
    "            # (사용자 지정에 따라 달리 할 수 있지만 여기서는 단순히 왼상단 혹은 오른상단 등)\n",
    "            if tpos in [\"top\", \"left\"]:\n",
    "                anchor_x, anchor_y = x1, y1\n",
    "            elif tpos in [\"bottom\"]:\n",
    "                anchor_x, anchor_y = x1, y2\n",
    "            else: # \"right\"\n",
    "                anchor_x, anchor_y = x2, y1\n",
    "\n",
    "            # draw_text_with_options\n",
    "            draw_text_with_options(\n",
    "                image=output,\n",
    "                text_lines=text_lines,\n",
    "                anchor_point=(anchor_x, anchor_y),\n",
    "                text_position=tpos,\n",
    "                color=color,\n",
    "                font_scale=font_scale,\n",
    "                auto_adjust=AUTO_ADJUST_TEXT_POSITION\n",
    "            )\n",
    "\n",
    "    return output\n",
    "\n",
    "print(\"===== [Cell 4] 디텍션 및 시각화 함수 정의 완료 =====\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------------------------------------------------------------\n",
    "# [Cell 5] 여백 추가 함수 (기존 코드 유지)\n",
    "# --------------------------------------------------------------------------------\n",
    "\n",
    "def add_margin_frame(frame, top_margin=0, side_margin=0):\n",
    "    \"\"\"\n",
    "    영상 상/하/좌/우 여백을 추가하는 함수.\n",
    "    \"\"\"\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] 여백 추가 함수 정의 완료 =====\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [Cell 6] ByteTracker + 슬라이딩 윈도우 기반 점수 필터링 로직\n",
    "def create_worker_helmet_head_tracker():\n",
    "    return BYTETracker(tracking_args)\n",
    "\n",
    "def update_track_states(tracks, track_dict, fps):\n",
    "    \"\"\"\n",
    "    1) 슬라이딩 윈도우 기반 score 업데이트\n",
    "    2) stable_time_window, required_pass_ratio는 track별 class_label에 따라 적용\n",
    "    \"\"\"\n",
    "    for t in tracks:\n",
    "        tid = t.track_id\n",
    "\n",
    "        # bbox\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",
    "        # 딕셔너리 미존재 시 생성 (혹은 누락 키를 추가) -------------------------\n",
    "        if tid not in track_dict:\n",
    "            track_dict[tid] = {\n",
    "                'last_center': None,  # 처음엔 None\n",
    "                'is_filtered': False,\n",
    "                'score_history': [],\n",
    "                'score_passed': False,\n",
    "                'class_label': None\n",
    "            }\n",
    "\n",
    "        # [수정] 'last_center'가 None이면 이번 프레임 center를 넣고 distance계산은 생략\n",
    "        if track_dict[tid]['last_center'] is None:\n",
    "            track_dict[tid]['last_center'] = center_now\n",
    "        else:\n",
    "            # 이동량 계산\n",
    "            last_center = track_dict[tid]['last_center']\n",
    "            dist_val = distance(center_now, last_center)\n",
    "            track_dict[tid]['last_center'] = center_now\n",
    "\n",
    "        if track_dict[tid]['is_filtered']:\n",
    "            continue\n",
    "\n",
    "        # (1) score 추가\n",
    "        current_score = t.score\n",
    "        track_dict[tid]['score_history'].append(current_score)\n",
    "\n",
    "        # (2) 클래스 라벨별 파라미터 조회\n",
    "        c_lbl = track_dict[tid]['class_label']\n",
    "        if c_lbl is None:\n",
    "            # 클래스가 아직 할당 안 된 상태면 default\n",
    "            stable_time_window = CLASS_PARAMS[\"default\"][\"stable_time_window\"]\n",
    "            required_pass_ratio = CLASS_PARAMS[\"default\"][\"required_pass_ratio\"]\n",
    "            sc_threshold = CLASS_PARAMS[\"default\"][\"detect_threshold\"]\n",
    "        else:\n",
    "            stable_time_window = get_class_param(c_lbl, \"stable_time_window\")\n",
    "            required_pass_ratio = get_class_param(c_lbl, \"required_pass_ratio\")\n",
    "            sc_threshold = get_class_param(c_lbl, \"detect_threshold\")\n",
    "\n",
    "        # (3) 슬라이딩 윈도우\n",
    "        stable_window_size = int(fps * stable_time_window)\n",
    "        if stable_window_size <= 0:\n",
    "            stable_window_size = 1\n",
    "\n",
    "        hist_scores = track_dict[tid]['score_history']\n",
    "        if len(hist_scores) > stable_window_size:\n",
    "            track_dict[tid]['score_history'] = hist_scores[-stable_window_size:]\n",
    "\n",
    "        pass_count = sum(1 for s in track_dict[tid]['score_history'] if s >= sc_threshold)\n",
    "        if len(track_dict[tid]['score_history']) == 0:\n",
    "            ratio = 0.0\n",
    "        else:\n",
    "            ratio = pass_count / len(track_dict[tid]['score_history'])\n",
    "\n",
    "        if ratio >= required_pass_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(\n",
    "    video_path,\n",
    "    model,\n",
    "    class_names,\n",
    "    do_inference=True,\n",
    "    do_save=True,\n",
    "    resize_factor=1.0\n",
    "):\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",
    "    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 fps <= 0:\n",
    "        fps = 30.0\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_worker_helmet_head_tracker()\n",
    "    track_dict = {}\n",
    "\n",
    "    frame_times = []\n",
    "    idx = 0\n",
    "\n",
    "    print(f\"[INFO] 영상 처리 시작 -> {video_path}\")\n",
    "    start_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",
    "            detections = detect_objects(frame, model, class_names)\n",
    "\n",
    "            xyxyscore = []\n",
    "            for d in detections:\n",
    "                x1, y1, x2, y2 = d['bbox']\n",
    "                s = d['score']\n",
    "                xyxyscore.append([x1, y1, x2, y2, s])\n",
    "            online_tracks = tracker.update(np.array(xyxyscore), [h_org, w_org])\n",
    "\n",
    "            # (3) 각 track 클래스 라벨 부여\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",
    "                # [수정] track_dict[tid]가 없으면 여기서도 초기화\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",
    "                    }\n",
    "\n",
    "                # 이미 라벨이 있으면 갱신 불필요\n",
    "                if track_dict[tid]['class_label'] is None:\n",
    "                    best_iou = 0\n",
    "                    best_label = None\n",
    "                    for d in detections:\n",
    "                        iou_val = iou_calc(track_box, d['bbox'])\n",
    "                        if iou_val > best_iou:\n",
    "                            best_iou = iou_val\n",
    "                            best_label = d['label']\n",
    "\n",
    "                    if best_label and best_iou > 0.5:\n",
    "                        track_dict[tid]['class_label'] = best_label\n",
    "\n",
    "            # (4) 슬라이딩 윈도우 기반 score 필터링\n",
    "            update_track_states(online_tracks, track_dict, fps)\n",
    "\n",
    "            # (5) 필터링 통과한 객체만 최종 디텍션\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",
    "                label_now = track_dict[t.track_id].get('class_label', 'unknown')\n",
    "                final_detections.append({\n",
    "                    'bbox': [x1, y1, x2, y2],\n",
    "                    'score': t.score,\n",
    "                    'label': label_now,\n",
    "                    'track_id': t.track_id\n",
    "                })\n",
    "\n",
    "            # (6) 시각화\n",
    "            vis_frame = draw_detections(frame, final_detections, font_scale=0.5)\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",
    "\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 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_total:.1f}초\")\n",
    "\n",
    "print(\"===== [Cell 7] 메인 영상 처리 함수 정의 완료 =====\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [Cell 8] 실행부\n",
    "video_inference_on = True\n",
    "video_save_on = True\n",
    "resize_factor_video = 1.0\n",
    "\n",
    "det_model = det_model_final\n",
    "cls_names = class_names_final\n",
    "\n",
    "exts = {'.mp4', '.avi', '.mov', '.mkv'}\n",
    "files = sorted(os.listdir(VIDEO_DIR))\n",
    "\n",
    "for f in files:\n",
    "    base, ex = os.path.splitext(f)\n",
    "    if ex.lower() in exts:\n",
    "        p = os.path.join(VIDEO_DIR, f)\n",
    "        process_video(\n",
    "            video_path=p,\n",
    "            model=det_model,\n",
    "            class_names=cls_names,\n",
    "            do_inference=video_inference_on,\n",
    "            do_save=video_save_on,\n",
    "            resize_factor=resize_factor_video\n",
    "        )\n",
    "    else:\n",
    "        print(f\"[스킵] 영상파일 아님 -> {f}\")\n",
    "\n",
    "print(\"===== [Cell 8] 전체 영상 처리 완료. =====\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "추가 설명 및 주의 사항\n",
    "클래스별 파라미터 적용 로직\n",
    "\n",
    "Detection 단계에서 각 bounding box가 해당 클래스의 detect_threshold를 통과해야 detections 목록에 포함됩니다.\n",
    "\n",
    "이후, ByteTracker로 넘어가면, Track별로 어떤 클래스인지(track_dict[tid]['class_label'])를 iou 매칭을 통해 찾고, 그 클래스에 맞는 stable_time_window, required_pass_ratio를 적용합니다.\n",
    "\n",
    "“윈도우 내 특정 비율 이상이 threshold를 넘으면 최종 필터 통과(score_passed=True)” 구조입니다.\n",
    "\n",
    "시각화 옵션\n",
    "\n",
    "VISUALIZATION_OPTIONS의 구조는 {\"default\": {...}, \"helmet\": {...}, \"worker\": {...}, ...} 식으로, 클래스명(key)별 세부 옵션을 둘 수 있습니다.\n",
    "\n",
    "각 항목(show_bbox, show_label, show_score, show_track_id, text_position)을 True/False로 끄고 켤 수 있으며, text_position을 \"top\", \"bottom\", \"left\", \"right\" 중 하나로 선택할 수 있습니다.\n",
    "\n",
    "만약 특정 클래스에 해당 옵션이 지정되지 않았다면, \"default\" 키의 옵션을 사용합니다.\n",
    "\n",
    "텍스트가 프레임 바깥으로 나갈 경우, AUTO_ADJUST_TEXT_POSITION = True이면 자동 보정(왼쪽/위쪽/오른쪽/아래쪽으로 살짝 옮김)을 수행합니다.\n",
    "\n",
    "텍스트 중첩 방지\n",
    "\n",
    "라인마다 약간씩 Y좌표(혹은 X좌표)를 오프셋(offset = i*(text_h+5))을 줘서, 클래스 이름 / 점수 / 트랙ID가 겹치지 않게 했습니다.\n",
    "\n",
    "만약 여러 객체가 서로 인접해 있어 텍스트가 서로 겹치는 상황은, 좀 더 고급 로직(충돌 감지 등)이 필요하겠으나, 현재 예시 코드에서는 지원하지 않습니다.\n",
    "\n",
    "코드 효율/최적화\n",
    "\n",
    "각 프레임마다 iou_calc를 전부 수행하므로, 인물이 많아지면 연산량이 증가합니다. 필요시 NMS 혹은 다른 최적화 기법을 적용할 수 있습니다.\n",
    "\n",
    "track_dict의 메모리가 커질 수 있으니, 종료된 트랙(사라진 트랙)은 일정 시간이 지난 뒤 track_dict에서 제거해 주면 메모리 사용량을 줄일 수 있습니다(본 예시에는 구현 X).\n",
    "\n",
    "오류 혹은 문제 상황\n",
    "\n",
    "FPS가 0으로 읽히거나, 영상을 열지 못하는 경우 등의 예외 처리를 추가했습니다(방어 코드).\n",
    "\n",
    "ByteTracker 내부 동작 관련 이슈는 yolox 라이브러리 자체 문제일 수 있으므로, 필요시 yolox 버전과 호환 여부를 확인해주세요.\n",
    "\n",
    "디버깅/테스트\n",
    "\n",
    "Jupyter Notebook에서 실행하면, [Cell 8]이 진행 상황을 tqdm으로 표시하며, 최종적으로 mp4 파일을 SAVE_RESULT 경로에 저장합니다.\n",
    "\n",
    "저장된 mp4 파일을 로컬 환경에서 플레이어로 열어 확인해 주세요 (GUI 환경이 아니면 ffmpeg나 opencv-python을 통해 다시 확인 가능).\n",
    "\n",
    "위와 같이 전체 코드를 셀 단위로 분리하였으며, 원본 코드를 최대한 유지하면서 클래스별 파라미터, 시각화 옵션 on/off, 텍스트 위치/중첩 방지 등을 모두 반영했습니다. 필요에 따라 더 세밀한 커스터마이징(예: 특정 클래스만 별도 색상, 더 복잡한 자동 위치 조정 등)을 적용하실 수 있습니다."
   ]
  }
 ],
 "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
}