{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "# # [1] 초기 설정 & 라이브러리 임포트\n",
    "\n",
    "# %%\n",
    "import os\n",
    "import sys\n",
    "import cv2\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "\n",
    "# mmengine & mmdet (또는 mmyolo) 관련\n",
    "from mmengine.registry import init_default_scope\n",
    "from mmdet.apis import init_detector, inference_detector\n",
    "from mmdet.utils import register_all_modules\n",
    "\n",
    "# 필요시 mmyolo 디렉토리를 파이썬 path에 추가\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",
    "# mmdet/mmyolo 레지스트리 등록\n",
    "register_all_modules()\n",
    "init_default_scope('mmdet')\n",
    "\n",
    "# --------------------------------------------------------\n",
    "# (!!!) 자주 바꿀 수 있는 파라미터들\n",
    "# --------------------------------------------------------\n",
    "VIDEO_PATH = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/포스코_광주_고양_합침/2024-12-24 10_50_25 T_C #3.mp4\"\n",
    "\n",
    "# 최종 결과를 저장할 메인 폴더\n",
    "MAIN_OUTPUT_DIR = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/모델비교_아웃풋들\"\n",
    "\n",
    "# 디바이스 설정\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"🚀 디바이스: {device}\")\n",
    "\n",
    "# --------------------------------------------------------\n",
    "# (!!!) 4개 모델 정보\n",
    "# --------------------------------------------------------\n",
    "MODEL_INFOS = [\n",
    "    {\n",
    "        \"model_name\": \"model_1\",\n",
    "        \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/ff_base_1819+cut+rotate+cutrota_1,1,3_01/yolo8_large_960size_ff_base_1819+cut+rotate+cutrota_1,1,3.py\",\n",
    "        \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/ff_base_1819+cut+rotate+cutrota_1,1,3_01/best_coco_bbox_mAP_epoch_6.pth\",\n",
    "        \"class_names\": [\"worker\", \"helmet\", \"head\"],\n",
    "        \"score_thresholds\": {\"worker\":0.1, \"helmet\":0.1, \"head\":0.1}\n",
    "    },\n",
    "    {\n",
    "        \"model_name\": \"model_2\",\n",
    "        \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/ff_base_resize_1,2,0.5_1819+cut+rotate+cutrota_1,1,3_01/yolo8_large_960size_ff_base_resize_1,2,0.5_1819+cut+rotate+cutrota_1,1,3.py\",\n",
    "        \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/ff_base_resize_1,2,0.5_1819+cut+rotate+cutrota_1,1,3_01/best_coco_bbox_mAP_epoch_8.pth\",\n",
    "        \"class_names\": [\"worker\", \"helmet\", \"head\"],\n",
    "        \"score_thresholds\": {\"worker\":0.1, \"helmet\":0.1, \"head\":0.1}\n",
    "    },\n",
    "    {\n",
    "        \"model_name\": \"model_3\",\n",
    "        \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_load_epoch16_1819_add_head_2,1,10_01/yolov8_l_960_custom_add_head_ff_result_2,1,10.py\",\n",
    "        \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_load_epoch16_1819_add_head_2,1,10_01/best_coco_bbox_mAP_epoch_19.pth\",\n",
    "        \"class_names\": [\"worker\", \"helmet\", \"head\"],\n",
    "        \"score_thresholds\": {\"worker\":0.1, \"helmet\":0.1, \"head\":0.1}\n",
    "    },\n",
    "    {\n",
    "        \"model_name\": \"model_4\",\n",
    "        \"config\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_roboflow_fall_only_rotate_960_yolo_01/yolov8_l_960_custom_fall.py\",\n",
    "        \"checkpoint\": \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmyolo/work_dirs/it1_roboflow_fall_only_rotate_960_yolo_01/best_coco_bbox_mAP_epoch_15.pth\",\n",
    "        \"class_names\": [\"worker\"],  # fall 모델에서의 클래스 명칭(실은 1개만)\n",
    "        \"score_thresholds\": {\"worker\":0.1}\n",
    "    },\n",
    "]\n",
    "\n",
    "# 시각화 파라미터\n",
    "BOX_COLOR = (0, 255, 0)\n",
    "THICKNESS = 2\n",
    "FONT_SCALE = 0.5\n",
    "FONT_THICKNESS = 1\n",
    "\n",
    "print(\"✅ [1] 초기 설정 완료.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "# # [2] 결과 폴더 생성 함수\n",
    "# \n",
    "# - exp, exp2... 식으로 자동 생성해, 덮어씌우지 않도록\n",
    "\n",
    "# %%\n",
    "def get_new_output_folder(base_dir, prefix=\"exp\"):\n",
    "    i = 1\n",
    "    while True:\n",
    "        folder_name = prefix if i == 1 else f\"{prefix}{i}\"\n",
    "        new_dir = os.path.join(base_dir, folder_name)\n",
    "        if not os.path.exists(new_dir):\n",
    "            os.makedirs(new_dir, exist_ok=True)\n",
    "            return new_dir\n",
    "        i += 1\n",
    "\n",
    "output_folder = get_new_output_folder(MAIN_OUTPUT_DIR, prefix=\"exp\")\n",
    "print(f\"✅ 결과 저장 폴더 생성 완료: {output_folder}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "# # [3] 모델 로드 & 디텍션 함수\n",
    "# \n",
    "# - init_detector()로 모델 객체 로딩\n",
    "# - detect_objects(), draw_detections() 함수 정의\n",
    "\n",
    "# %%\n",
    "def load_model(config_path, checkpoint_path, device=\"cpu\"):\n",
    "    model = init_detector(config_path, checkpoint_path, device=device)\n",
    "    return model\n",
    "\n",
    "def detect_objects(frame_bgr, model, class_names, class_score_thresholds):\n",
    "    \"\"\"\n",
    "    frame_bgr: BGR 이미지\n",
    "    model: mmdet/mmyolo 기반 디텍터\n",
    "    class_names: ex) [\"worker\",\"helmet\",\"head\"]\n",
    "    class_score_thresholds: ex) {\"worker\":0.1, \"helmet\":0.1, ...}\n",
    "    \"\"\"\n",
    "    result = inference_detector(model, frame_bgr)\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",
    "        box = bboxes[i].tolist()  # [x1,y1,x2,y2]\n",
    "        sc = scores[i].item()\n",
    "        cid = labels[i].item()\n",
    "        if cid < len(class_names):\n",
    "            cname = class_names[cid]\n",
    "        else:\n",
    "            cname = f\"cls_{cid}\"\n",
    "        \n",
    "        thr = class_score_thresholds.get(cname, 0.3)\n",
    "        if sc < thr:\n",
    "            continue\n",
    "        \n",
    "        x1, y1, x2, y2 = box\n",
    "        w = x2 - x1\n",
    "        h = y2 - y1\n",
    "        area = w*h\n",
    "        \n",
    "        detections.append({\n",
    "            \"bbox\": [x1,y1,x2,y2],\n",
    "            \"score\": sc,\n",
    "            \"class_id\": cid,\n",
    "            \"class_name\": cname,\n",
    "            \"area\": area\n",
    "        })\n",
    "    return detections\n",
    "\n",
    "def draw_detections(frame_bgr, detections,\n",
    "                    box_color=(0,255,0),\n",
    "                    thickness=2,\n",
    "                    font_scale=0.5,\n",
    "                    font_thickness=1):\n",
    "    output = frame_bgr.copy()\n",
    "    for det in detections:\n",
    "        x1, y1, x2, y2 = map(int, det[\"bbox\"])\n",
    "        sc = det[\"score\"]\n",
    "        cname = det[\"class_name\"]\n",
    "        \n",
    "        cv2.rectangle(output, (x1, y1), (x2, y2), box_color, thickness)\n",
    "        text = f\"{cname} {sc:.2f}\"\n",
    "        cv2.putText(output, text, (x1, y1-5),\n",
    "                    cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                    font_scale, box_color, font_thickness, cv2.LINE_AA)\n",
    "    return output\n",
    "\n",
    "print(\"✅ [3] 모델 로드 & 디텍션 함수 정의 완료.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "# # [4] 부분 프레임 추론 & 결과 수집\n",
    "# \n",
    "# - 4개 모델 로드\n",
    "# - 지정된 초(time_in_seconds) 주위 ±2프레임(기본값)만 추론\n",
    "# - 각 모델별로 디텍션 → 시각화 이미지 저장 & 노트북 출력\n",
    "# - 동시에 (model_name, frame_idx, class_name, score, area, bbox...)를 리스트에 저장, 최종 리턴\n",
    "\n",
    "# %%\n",
    "loaded_models = []\n",
    "for info in MODEL_INFOS:\n",
    "    print(f\"🔄 로딩 중: {info['model_name']}\")\n",
    "    mdl = load_model(info[\"config\"], info[\"checkpoint\"], device=device)\n",
    "    loaded_models.append({\n",
    "        \"model_name\": info[\"model_name\"],\n",
    "        \"model_obj\": mdl,\n",
    "        \"class_names\": info[\"class_names\"],\n",
    "        \"score_thresholds\": info[\"score_thresholds\"]\n",
    "    })\n",
    "print(f\"✅ 총 {len(loaded_models)}개 모델 로드 완료.\")\n",
    "\n",
    "def run_detection_for_time(video_path,\n",
    "                           time_in_seconds,\n",
    "                           loaded_models,\n",
    "                           frame_range=2,\n",
    "                           save_dir=None):\n",
    "    \"\"\"\n",
    "    - time_in_seconds → frame_index (fps 사용)\n",
    "    - frame_index ± frame_range만 추론\n",
    "    - 모델별로 detection → 시각화 & 저장\n",
    "    - 리턴: detection_results (list of dict)\n",
    "       예: [{\n",
    "              \"model_name\": \"model_1\",\n",
    "              \"frame_idx\": 123,\n",
    "              \"class_name\": \"helmet\",\n",
    "              \"score\": 0.88,\n",
    "              \"area\": 12000,\n",
    "              \"x1\":..., \"y1\":..., \"x2\":..., \"y2\":...\n",
    "           }, ...]\n",
    "    \"\"\"\n",
    "    det_records = []  # 최종 리턴용 (모든 모델+프레임 결과)\n",
    "    \n",
    "    cap = cv2.VideoCapture(video_path)\n",
    "    if not cap.isOpened():\n",
    "        print(f\"[에러] 영상 로드 실패: {video_path}\")\n",
    "        return det_records  # 빈 리스트\n",
    "    \n",
    "    fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "    \n",
    "    target_frame = int(time_in_seconds * fps)\n",
    "    start_f = max(0, target_frame - frame_range)\n",
    "    end_f = min(total_frames - 1, target_frame + frame_range)\n",
    "    \n",
    "    print(f\"🎯 중심 초={time_in_seconds} → 중심프레임={target_frame}\")\n",
    "    print(f\"   처리 구간: [{start_f} ~ {end_f}]\")\n",
    "    \n",
    "    for fidx in range(start_f, end_f + 1):\n",
    "        cap.set(cv2.CAP_PROP_POS_FRAMES, fidx)\n",
    "        ret, frame_bgr = cap.read()\n",
    "        if not ret:\n",
    "            print(f\"[경고] 프레임 {fidx} 읽기 실패.\")\n",
    "            continue\n",
    "        \n",
    "        for m_info in loaded_models:\n",
    "            m_name = m_info[\"model_name\"]\n",
    "            model_obj = m_info[\"model_obj\"]\n",
    "            cls_names = m_info[\"class_names\"]\n",
    "            score_thr_dict = m_info[\"score_thresholds\"]\n",
    "            \n",
    "            # 디텍션\n",
    "            detections = detect_objects(frame_bgr, model_obj,\n",
    "                                        cls_names, score_thr_dict)\n",
    "            # 시각화\n",
    "            vis_img = draw_detections(frame_bgr, detections,\n",
    "                                      box_color=BOX_COLOR,\n",
    "                                      thickness=THICKNESS,\n",
    "                                      font_scale=FONT_SCALE,\n",
    "                                      font_thickness=FONT_THICKNESS)\n",
    "            \n",
    "            # 저장\n",
    "            if save_dir is not None:\n",
    "                base_name = os.path.splitext(os.path.basename(video_path))[0]\n",
    "                out_fn = f\"{base_name}_frame{fidx:04d}_{m_name}.jpg\"\n",
    "                out_path = os.path.join(save_dir, out_fn)\n",
    "                cv2.imwrite(out_path, vis_img)\n",
    "                # print(f\"저장: {out_path}\")\n",
    "            \n",
    "            # 노트북 inline 표시\n",
    "            plt.figure(figsize=(8,6))\n",
    "            plt.title(f\"{m_name} / frame={fidx}\")\n",
    "            plt.imshow(cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB))\n",
    "            plt.axis(\"off\")\n",
    "            plt.show()\n",
    "            \n",
    "            # 결과 누적 (model_name, frame_idx, class_name, score, area 등)\n",
    "            for det in detections:\n",
    "                x1, y1, x2, y2 = det[\"bbox\"]\n",
    "                record = {\n",
    "                    \"model_name\": m_name,\n",
    "                    \"frame_idx\": fidx,\n",
    "                    \"class_name\": det[\"class_name\"],\n",
    "                    \"score\": det[\"score\"],\n",
    "                    \"area\": det[\"area\"],\n",
    "                    \"x1\": x1, \"y1\": y1, \"x2\": x2, \"y2\": y2\n",
    "                }\n",
    "                det_records.append(record)\n",
    "    \n",
    "    cap.release()\n",
    "    return det_records\n",
    "\n",
    "print(\"✅ [4] 부분 프레임 추론 & 결과 수집 함수 정의 완료.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "# # [5] 특정 초 주변 프레임만 추론 (예시)\n",
    "# \n",
    "# - time_in_seconds=10.0 (10초 지점)\n",
    "# - frame_range=2 (±2프레임, 총5프레임)\n",
    "# - 결과 이미지는 output_folder 내에 저장\n",
    "# - return된 det_records를 DataFrame으로 구성\n",
    "\n",
    "# %%\n",
    "time_in_seconds = 10.0\n",
    "frame_range_val = 2\n",
    "\n",
    "det_list = run_detection_for_time(\n",
    "    video_path=VIDEO_PATH,\n",
    "    time_in_seconds=time_in_seconds,\n",
    "    loaded_models=loaded_models,\n",
    "    frame_range=frame_range_val,\n",
    "    save_dir=output_folder  # None이면 저장 안 함\n",
    ")\n",
    "\n",
    "# det_list -> DataFrame\n",
    "df_det = pd.DataFrame(det_list)\n",
    "print(\"🔎 df_det.shape:\", df_det.shape)\n",
    "print(df_det.head(10))\n",
    "\n",
    "print(\"✅ [5] 실행 예시 완료. 위에 표시된 이미지를 확인하세요.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "# # [6] 데이터 분석 & 그래프 저장\n",
    "# \n",
    "# - df_det에는 [model_name, frame_idx, class_name, score, area, x1,y1,x2,y2]가 있음.\n",
    "# - 모델별 평균 score, 평균 area, 클래스별 count 등을 집계\n",
    "# - 그래프(bar/hist)로 저장 + 표시\n",
    "# - CSV 저장 예시\n",
    "\n",
    "# %%\n",
    "if df_det.empty:\n",
    "    print(\"⚠️ df_det가 비어있음. 추론된 결과가 없습니다.\")\n",
    "else:\n",
    "    # (A) 모델별 평균 score & 평균 area\n",
    "    df_model_group = df_det.groupby(\"model_name\").agg({\n",
    "        \"score\":\"mean\",\n",
    "        \"area\":\"mean\"\n",
    "    }).reset_index().rename(columns={\n",
    "        \"score\":\"avg_score\",\n",
    "        \"area\":\"avg_area\"\n",
    "    })\n",
    "    print(\"\\n=== 모델별 평균 스코어 & 평균 면적 ===\")\n",
    "    print(df_model_group)\n",
    "\n",
    "    # (B) 모델별 클래스별 count\n",
    "    df_model_cls = df_det.groupby([\"model_name\",\"class_name\"]).size().reset_index(name=\"count\")\n",
    "    print(\"\\n=== 모델별 클래스별 검출 개수 ===\")\n",
    "    print(df_model_cls)\n",
    "\n",
    "    # (C) 그래프 예시: 모델별 평균 스코어\n",
    "    plt.figure(figsize=(6,4))\n",
    "    plt.bar(df_model_group[\"model_name\"], df_model_group[\"avg_score\"])\n",
    "    plt.title(\"Average Confidence by Model\")\n",
    "    plt.xlabel(\"Model\")\n",
    "    plt.ylabel(\"Mean Score\")\n",
    "    plt.ylim([0,1])\n",
    "    graph_path = os.path.join(output_folder, \"avg_score_by_model.png\")\n",
    "    plt.savefig(graph_path, dpi=150)\n",
    "    plt.show()\n",
    "    print(f\"✅ 그래프 저장: {graph_path}\")\n",
    "\n",
    "    # (D) 그래프 예시: 모델별 평균 Area\n",
    "    plt.figure(figsize=(6,4))\n",
    "    plt.bar(df_model_group[\"model_name\"], df_model_group[\"avg_area\"])\n",
    "    plt.title(\"Average BBox Area by Model\")\n",
    "    plt.xlabel(\"Model\")\n",
    "    plt.ylabel(\"Mean Area (pixel^2)\")\n",
    "    graph_path2 = os.path.join(output_folder, \"avg_area_by_model.png\")\n",
    "    plt.savefig(graph_path2, dpi=150)\n",
    "    plt.show()\n",
    "    print(f\"✅ 그래프 저장: {graph_path2}\")\n",
    "\n",
    "    # (E) CSV 저장\n",
    "    csv_path_det = os.path.join(output_folder, \"detection_records.csv\")\n",
    "    df_det.to_csv(csv_path_det, index=False, encoding=\"utf-8-sig\")\n",
    "    print(f\"✅ 전체 디텍션 CSV 저장: {csv_path_det}\")\n",
    "\n",
    "    csv_path_model_group = os.path.join(output_folder, \"summary_by_model.csv\")\n",
    "    df_model_group.to_csv(csv_path_model_group, index=False)\n",
    "    print(f\"✅ 모델별 평균 스코어/면적 CSV 저장: {csv_path_model_group}\")\n",
    "\n",
    "    csv_path_model_cls = os.path.join(output_folder, \"summary_by_model_class.csv\")\n",
    "    df_model_cls.to_csv(csv_path_model_cls, index=False)\n",
    "    print(f\"✅ 모델별 클래스별 검출 개수 CSV 저장: {csv_path_model_cls}\")\n",
    "\n",
    "print(\"✅ [6] 데이터 분석 & 그래프 저장 단계 완료.\")\n",
    "print(f\"📁 최종 결과 디렉토리: {output_folder}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ltb",
   "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": 2
}