{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# %% [셀 1]: 라이브러리 임포트 및 전역 설정\n",
    "# --------------------------------------------------\n",
    "# 1) 필요한 라이브러리 임포트\n",
    "# 2) 경로, 클래스 이름, 스코어 임계값 등 전역 변수 설정\n",
    "# 3) mmdet, yolox, mmpose 관련 셋업\n",
    "# 4) 슬라이딩 윈도우 기반 score 안정성 판정을 위한 전역 변수 설정\n",
    "# --------------------------------------------------\n",
    "\n",
    "import os\n",
    "import os.path as osp\n",
    "import time\n",
    "import cv2\n",
    "import numpy as np\n",
    "from loguru import logger\n",
    "from types import SimpleNamespace\n",
    "import glob\n",
    "import sys\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",
    "# from mmpose.evaluation.functional import nms\n",
    "\n",
    "from yolox.tracker.byte_tracker import BYTETracker\n",
    "from yolox.tracking_utils.timer import Timer\n",
    "\n",
    "# from mmpose.apis import inference_topdown, init_model as init_pose_estimator\n",
    "# from mmpose.registry import VISUALIZERS\n",
    "# from mmpose.structures import merge_data_samples\n",
    "\n",
    "# mmyolo 및 mmpose 경로 추가\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",
    "# mmpose_dir = \"/DATA2/ltb/mmpose\"\n",
    "# if mmpose_dir not in sys.path:\n",
    "#     sys.path.append(mmpose_dir)\n",
    "\n",
    "register_all_modules()\n",
    "\n",
    "# 전역 경로 및 설정\n",
    "BASE_DIR = '/DATA2/ltb'\n",
    "# 기존 영상 경로 사용 (필요에 따라 변경 가능)\n",
    "VIDEO_DIR = ('/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/크기비교용_샘플')\n",
    "# 저장 결과 폴더 (11개 클래스용으로 별도 지정)\n",
    "SAVE_RESULT = ('/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/output_val_samples/크기비교용_샘플_0.3')\n",
    "os.makedirs(SAVE_RESULT, exist_ok=True)\n",
    "\n",
    "# --------------------------------------------------\n",
    "# *** [변경/추가] 새로운 11개 클래스 설정 및 각 클래스별 스코어 임계값 ***\n",
    "# 최종 클래스 목록 (순서가 config와 weight에 맞춰져 있어야 함)\n",
    "class_names_final = [\n",
    "    'worker', \n",
    "    'helmet', \n",
    "    'head'\n",
    "]\n",
    "# 모든 클래스에 대해 임계값 0.0 (필요시 값 조정)\n",
    "class_score_thresholds_final = {\n",
    "    'worker': 0.0,\n",
    "    'helmet': 0.0,\n",
    "    'head': 0.0\n",
    "}\n",
    "\n",
    "# 디바이스 및 ByteTracker 파라미터 (기존과 동일)\n",
    "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\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",
    "# 슬라이딩 윈도우 기반 Score 판정 관련 변수\n",
    "required_pass_ratio = 0.3    # 윈도우 내 통과비율\n",
    "stable_time_window = 1.0     # 슬라이딩 윈도우 크기(초)\n",
    "score_threshold = 0.1        # 디텍션 score 임계값\n",
    "\n",
    "print(\"[셀 1] 라이브러리 임포트 및 전역 설정 완료.\")\n",
    "print(f\"[INFO] BASE_DIR: {BASE_DIR}\")\n",
    "print(f\"[INFO] Input Video Directory: {VIDEO_DIR}\")\n",
    "print(f\"[INFO] Output Directory: {SAVE_RESULT}\")\n",
    "print(f\"[INFO] 디바이스: {device}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# %% [셀 2]: 헬퍼 함수 (IoU, distance 등)\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(\"[셀 2] 헬퍼 함수 정의 완료.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# %% [셀 3]: 디텍션 모델 로드\n",
    "import sys\n",
    "sys.path.append(\"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmdetection\")\n",
    "\n",
    "# *** [변경/추가] 새로운 weight 및 config 경로 (11개 클래스 통합) ***\n",
    "config_file_final = '/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/0.3x_1819_cutrota_01/0.3x_1819_cutrota.py'\n",
    "checkpoint_file_final = '/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/0.3x_1819_cutrota_01/best_coco_bbox_mAP_epoch_5.pth'\n",
    "\n",
    "os.chdir(mmyolo_dir)\n",
    "# 새로운 모델 초기화 (11개 클래스)\n",
    "det_model_final = init_detector(config_file_final, checkpoint_file_final, device=device)\n",
    "# 모델 메타 정보에 최종 클래스 목록 전달\n",
    "det_model_final.dataset_meta = {'CLASSES': class_names_final}\n",
    "\n",
    "print(\"[모델 로드 경로 안내]\")\n",
    "print(f\" [INFO] final config: {config_file_final}\")\n",
    "print(f\" [INFO] final weight: {checkpoint_file_final}\")\n",
    "print(\"[셀 3] 디텍션 모델 로드 완료.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# %% [셀 4]: 디텍션 및 시각화 함수 정의 (공용)\n",
    "def detect_objects(frame, model, class_score_thresholds_init, class_names):\n",
    "    \"\"\"\n",
    "    주어진 frame(이미지)과 모델, 클래스 스코어 임계값 등을 이용해\n",
    "    디텍션을 수행 후, 필터링된 결과를 반환한다.\n",
    "    \"\"\"\n",
    "    result = inference_detector(model, frame)\n",
    "    if not hasattr(result, 'pred_instances'):\n",
    "        return []\n",
    "    bboxes = result.pred_instances.bboxes\n",
    "    scores = result.pred_instances.scores\n",
    "    labels = result.pred_instances.labels\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 = class_score_thresholds_init.get(cname, 0.3)\n",
    "        if sc < thr:\n",
    "            continue\n",
    "        detections.append({\n",
    "            'bbox': box,\n",
    "            'score': sc,\n",
    "            'label': cname\n",
    "        })\n",
    "    return detections\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",
    "    output = frame.copy()\n",
    "    \n",
    "    # *** [추가됨] 11개 클래스에 대해 색상 매핑 dictionary (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",
    "    }\n",
    "    \n",
    "    for det in detections:\n",
    "        x1, y1, x2, y2 = [int(x) for x in det['bbox']]\n",
    "        sc = det['score']\n",
    "        lbl = det['label']\n",
    "        color = color_map.get(lbl, (128, 128, 128))  # 기본 회색\n",
    "        \n",
    "        # 트랙 id 출력 (존재하면)\n",
    "        if 'track_id' in det:\n",
    "            label_text = f\"{lbl} {det['track_id']} {sc:.2f}\"\n",
    "        else:\n",
    "            label_text = f\"{lbl} {sc:.2f}\"\n",
    "        cv2.rectangle(output, (x1, y1), (x2, y2), color, 2)\n",
    "        cv2.putText(output, label_text, (x1, max(y1-10, 0)),\n",
    "                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, 2)\n",
    "    return output\n",
    "\n",
    "print(\"[셀 4] 디텍션 및 시각화 함수 정의 완료.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# %% [셀 5]: 여백 추가 함수\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(\"[셀 5] 여백 추가 함수 정의 완료.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# %% [셀 7]: ByteTracker 기반 필터링 로직 준비 (슬라이딩 윈도우 방식)\n",
    "from yolox.tracker.byte_tracker import BYTETracker\n",
    "\n",
    "def create_worker_helmet_head_tracker():\n",
    "    \"\"\"\n",
    "    (기존 기능 그대로) Worker, helmet, head 등 및 기타 객체(필요시) 추적을 위한 ByteTracker 생성 함수.\n",
    "    \"\"\"\n",
    "    return BYTETracker(tracking_args)\n",
    "\n",
    "def update_track_states(\n",
    "    tracks, \n",
    "    track_dict, \n",
    "    fps,\n",
    "    stable_time_window, \n",
    "    score_threshold\n",
    "):\n",
    "    \"\"\"\n",
    "    [주요 기능]\n",
    "      1) 슬라이딩 윈도우 기반 score 업데이트 및 score_passed 여부 갱신\n",
    "\n",
    "    Parameters:\n",
    "      tracks: ByteTracker 업데이트 결과 (각 트랙 정보)\n",
    "      track_dict: 트랙 정보를 저장하는 dict\n",
    "      fps: 영상 FPS\n",
    "      stable_time_window: 슬라이딩 윈도우 크기(초)\n",
    "      score_threshold: detection score 임계값\n",
    "    \"\"\"\n",
    "    # 슬라이딩 윈도우 크기 (프레임 단위)\n",
    "    stable_window_size = int(fps * stable_time_window)\n",
    "    # 슬라이딩 윈도우 내 몇 프레임 이상 score가 통과해야 score_passed=True\n",
    "    required_pass_count = int(stable_window_size * 0.67)\n",
    "\n",
    "    for t in tracks:\n",
    "        tid = t.track_id\n",
    "        # track_dict 에 트랙아이디가 없으면 초기화\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",
    "            }\n",
    "\n",
    "        # 이미 필터된 트랙이면 무시\n",
    "        if track_dict[tid]['is_filtered']:\n",
    "            continue\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",
    "        last_center = track_dict[tid]['last_center']\n",
    "        if last_center is None:\n",
    "            track_dict[tid]['last_center'] = center_now\n",
    "        else:\n",
    "            # 이동 거리 계산 (필요시 추가 로직 가능)\n",
    "            dist_val = distance(center_now, last_center)\n",
    "            track_dict[tid]['last_center'] = center_now\n",
    "\n",
    "        # 슬라이딩 윈도우 기반 Score 업데이트\n",
    "        current_score = t.score\n",
    "        score_history = track_dict[tid]['score_history']\n",
    "        score_history.append(current_score)\n",
    "        if len(score_history) > stable_window_size:\n",
    "            score_history.pop(0)\n",
    "\n",
    "        pass_count = sum(1 for s in score_history if s >= score_threshold)\n",
    "        if pass_count >= required_pass_count:\n",
    "            track_dict[tid]['score_passed'] = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# %% [셀 8]: 쓰러짐 감지 함수 및 경고 그리기\n",
    "# (현재 pose 사용 안 하므로 주석. 필요한 경우 참고 가능.)\n",
    "\n",
    "# def fall_detection_single(data_sample):\n",
    "#     \"\"\"\n",
    "#     단일 인물 Pose 결과에 대한 쓰러짐 감지 함수.\n",
    "#     data_sample: MMPose의 DataSample (예: 하나의 사람에 대한 포즈 추정 결과)\n",
    "#\n",
    "#     return:\n",
    "#       - fall_detected: 쓰러짐 감지 여부 (bool)\n",
    "#       - bbox: 감지된 사람의 바운딩 박스\n",
    "#     \"\"\"\n",
    "#     if not hasattr(data_sample, 'pred_instances'):\n",
    "#         return False, None\n",
    "#     if len(data_sample.pred_instances) == 0:\n",
    "#         return False, None\n",
    "#\n",
    "#     keypoints = data_sample.pred_instances.keypoints[0]\n",
    "#     keypoint_scores = data_sample.pred_instances.keypoint_scores[0]\n",
    "#     bbox = data_sample.pred_instances.bboxes[0]\n",
    "#\n",
    "#     fall_detected = False\n",
    "#     return fall_detected, bbox\n",
    "#\n",
    "# def draw_fall_alert(image, bbox):\n",
    "#     x_min, y_min, x_max, y_max = bbox.astype(int)\n",
    "#     cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 0, 255), 2)\n",
    "#     cv2.putText(image, 'Fall Detected', (x_min, max(y_min - 10, 0)),\n",
    "#                 cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)\n",
    "#\n",
    "# print(\"[셀 8] 쓰러짐 감지 함수 및 경고 시각화 함수 정의 완료.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# %% [셀 9]: 전체 영상 처리 (앙상블) + 트래킹 및 필터링\n",
    "def process_video(\n",
    "    video_path, \n",
    "    models, \n",
    "    score_thresh_inits, \n",
    "    class_names_lists,\n",
    "    stable_time_window=1.0,\n",
    "    do_inference=True, \n",
    "    do_save=True, \n",
    "    resize_factor=1.0\n",
    "):\n",
    "    \"\"\"\n",
    "    영상에 대해 다음 작업을 수행:\n",
    "      1) 단일 디텍션 모델로 디텍션 수행 (11개 클래스)\n",
    "      2) ByteTracker를 통한 트래킹 및 슬라이딩 윈도우 기반 score 필터링\n",
    "      3) 최종 디텍션 결과 시각화 후 영상 저장\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    video_path : str\n",
    "        처리할 영상 파일 경로.\n",
    "    models : list\n",
    "        디텍션 모델 리스트 (여기서는 하나의 모델만 사용).\n",
    "    score_thresh_inits : list\n",
    "        클래스별 score 임계값 dict 리스트.\n",
    "    class_names_lists : list\n",
    "        클래스 이름 리스트.\n",
    "    stable_time_window : float\n",
    "        슬라이딩 윈도우 크기(초).\n",
    "    do_inference : bool\n",
    "        디텍션 수행 여부.\n",
    "    do_save : bool\n",
    "        결과 영상 저장 여부.\n",
    "    resize_factor : float\n",
    "        결과 영상 리사이즈 배율.\n",
    "    \"\"\"\n",
    "    if not do_inference:\n",
    "        print(f\"[스킵] 디텍션 off -> {video_path}\")\n",
    "        return\n",
    "    base_name = os.path.basename(video_path)\n",
    "    name, ext = os.path.splitext(base_name)\n",
    "    output_video_path = os.path.join(SAVE_RESULT, f\"detected_{name}.mp4\")\n",
    "\n",
    "    cap = cv2.VideoCapture(video_path)\n",
    "    if not cap.isOpened():\n",
    "        print(f\"[오류] 영상을 열 수 없습니다: {video_path}\")\n",
    "        return\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 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",
    "    frame_times = []\n",
    "    idx = 0\n",
    "\n",
    "    # *** [변경됨] 단일 모델 사용 (11개 클래스 통합) ***\n",
    "    final_model = models[0]\n",
    "    final_score = score_thresh_inits[0]\n",
    "    final_cls = class_names_lists[0]\n",
    "\n",
    "    # 트래커 생성 (필요시 사용; 기존 기능 유지)\n",
    "    tracker = create_worker_helmet_head_tracker()\n",
    "    track_dict = {}\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",
    "            # (1) 단일 모델 디텍션 수행 (11개 클래스)\n",
    "            detections = detect_objects(frame, final_model, final_score, final_cls)\n",
    "            # source 태그 추가 (원하는 경우 활용)\n",
    "            for d in detections:\n",
    "                d['source'] = 'final'\n",
    "\n",
    "            # (2) ByteTracker 업데이트 (디텍션 결과 기반)\n",
    "            # 디텍션 결과를 [x1, y1, x2, y2, score] 형식으로 변환\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) 트래킹 상태 업데이트 (슬라이딩 윈도우 기반 score 필터링)\n",
    "            update_track_states(online_tracks, track_dict, fps, stable_time_window, score_threshold)\n",
    "\n",
    "            # (4) 최종 디텍션 결과 생성 (score_passed True 인 경우만)\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",
    "                box_now = [x1, y1, x2, y2]\n",
    "                found_source = None\n",
    "                found_label = None  # *** [추가됨] 실제 클래스 라벨 저장 변수\n",
    "                for d in detections:\n",
    "                    if iou_calc(box_now, d['bbox']) > 0.5:\n",
    "                        found_source = d.get('source', None)\n",
    "                        found_label = d['label']  # 디텍션 결과의 실제 라벨을 추출\n",
    "                        break\n",
    "                # 만약 매칭되는 디텍션이 없으면 'unknown'으로 처리\n",
    "                if found_label is None:\n",
    "                    found_label = \"unknown\"\n",
    "                final_detections.append({\n",
    "                    'bbox': box_now,\n",
    "                    'score': t.score,\n",
    "                    'label': found_label,  # 'object'가 아닌 실제 라벨 사용\n",
    "                    'source': found_source,\n",
    "                    'track_id': t.track_id\n",
    "                })\n",
    "\n",
    "            # (5) 최종 시각화: 디텍션 bbox 및 트랙 정보 표시\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# %% [셀 10]: 실행부 (영상 폴더 내 모든 파일 처리)\n",
    "video_inference_on = True\n",
    "video_save_on = True\n",
    "resize_factor_video = 1.0\n",
    "\n",
    "# *** [변경됨] 새로운 모델, 클래스, score 임계값 사용 (11개 클래스) ***\n",
    "det_models = [ det_model_final ]\n",
    "scores_inits = [ class_score_thresholds_final ]\n",
    "names_lists = [ 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",
    "            models=det_models,\n",
    "            score_thresh_inits=scores_inits,\n",
    "            class_names_lists=names_lists,\n",
    "            stable_time_window=stable_time_window,\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(\"[셀 10] 전체 영상 처리 완료.\")"
   ]
  }
 ],
 "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
}
