{ "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", "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", "# 레지스트리 등록\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", "MAIN_OUTPUT_DIR = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/모델비교_아웃풋들\"\n", "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n", "print(f\"🚀 디바이스: {device}\")\n", "\n", "# 4개 모델 정보\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", "CLASS_COLORS = {\n", " \"worker\": (0, 255, 0), # 초록\n", " \"helmet\": (255, 0, 0), # 파랑\n", " \"head\": (0, 0, 255), # 빨강\n", " # 그 외 클래스는 기본 회색\n", "}\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)\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", "# %%\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 이미지 (numpy)\n", " model: mmdet/mmyolo detector\n", " class_names: [\"worker\",\"helmet\",\"head\", ...]\n", " class_score_thresholds: {\"worker\":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_name\": cname,\n", " \"width\": w,\n", " \"height\": h,\n", " \"area\": area\n", " })\n", " return detections\n", "\n", "\n", "def draw_detections_advanced(frame_bgr, detections, class_colors, font_scale=0.5, font_thick=1):\n", " \"\"\"\n", " - 클래스별 색상 다르게\n", " - 텍스트 (class, score, width, height, area)를 겹치지 않도록 여러 줄로 표시\n", " - 텍스트를 bbox 상단부에 표시하되, 2줄 이상\n", " 1줄: f\"{class_name} {score:.2f}\"\n", " 2줄: f\"W={w:.0f},H={h:.0f},Area={area:.0f}\"\n", " \"\"\"\n", " output = frame_bgr.copy()\n", " for det in detections:\n", " x1,y1,x2,y2 = map(int, det[\"bbox\"])\n", " cname = det[\"class_name\"]\n", " sc = det[\"score\"]\n", " w = det[\"width\"]\n", " h = det[\"height\"]\n", " area = det[\"area\"]\n", " \n", " color = class_colors.get(cname, (128,128,128)) # 기본 회색\n", " \n", " # bbox\n", " cv2.rectangle(output, (x1,y1), (x2,y2), color, 2)\n", " \n", " # 여러 줄 텍스트\n", " line1 = f\"{cname} {sc:.2f}\"\n", " line2 = f\"W={w:.0f},H={h:.0f},A={area:.0f}\"\n", " \n", " # 텍스트 높이를 측정해, 위로 2줄 쌓기\n", " txt_size1, _ = cv2.getTextSize(line1, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thick)\n", " txt_size2, _ = cv2.getTextSize(line2, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thick)\n", " \n", " # 보통 bbox 상단 외부에 표시 (y1 - 5 => 첫 줄, 그 위에 두 번째 줄)\n", " y_text_line1 = y1 - 5\n", " y_text_line2 = y_text_line1 - (txt_size1[1] + 5)\n", " \n", " # 만약 y_text_line2 < 0 인 경우, 화면 밖으로 가므로 bbox 내부(바로 아래)에 표시\n", " if y_text_line2 < 0:\n", " # bbox 내부로 위치 조정\n", " y_text_line2 = y2 + txt_size2[1] + 5\n", " y_text_line1 = y_text_line2 + txt_size1[1] + 5\n", " \n", " # line1\n", " cv2.putText(output, line1, (x1, y_text_line1),\n", " cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, font_thick, cv2.LINE_AA)\n", " # line2\n", " cv2.putText(output, line2, (x1, y_text_line2),\n", " cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, font_thick, cv2.LINE_AA)\n", " return output\n", "\n", "\n", "print(\"✅ [3] 모델 로드 & 디텍션 함수 정의 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [markdown]\n", "# # [4] 부분 프레임 추론 & 2x2 격자 이미지 생성\n", "# \n", "# - 각 프레임에 대해 4개 모델 디텍션 이미지를 만든 뒤, 2x2로 합쳐 한 장에 저장\n", "# - (딕셔너리 형태로) 디텍션 정보는 리스트로 모아서 return\n", "# \n", "# - 2x2 합치기 로직: make_image_grid(images, rows=2, cols=2)\n", "\n", "# %%\n", "def make_image_grid(img_list, rows=2, cols=2):\n", " \"\"\"\n", " img_list: [img1, img2, img3, img4]\n", " - 2x2로 배치\n", " - 단, img의 크기가 달라지면 resize 필요 (여기서는 동일 크기라고 가정)\n", " \"\"\"\n", " if len(img_list) < rows*cols:\n", " # 부족하면 빈 이미지로 채울 수도 있음\n", " # 여기서는 단순 처리\n", " blank = np.zeros_like(img_list[0])\n", " while len(img_list) < rows*cols:\n", " img_list.append(blank)\n", " \n", " h, w = img_list[0].shape[:2]\n", " grid = np.zeros((h*rows, w*cols, 3), dtype=np.uint8)\n", " idx = 0\n", " for r in range(rows):\n", " for c in range(cols):\n", " y1 = r*h\n", " y2 = y1 + h\n", " x1 = c*w\n", " x2 = x1 + w\n", " grid[y1:y2, x1:x2] = img_list[idx]\n", " idx += 1\n", " return grid\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\n", " - frame_index ± frame_range만 추론\n", " - 4개 모델 결과 이미지를 2x2로 합쳐 저장\n", " - detection info 반환\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_f = int(time_in_seconds * fps)\n", " start_f = max(0, target_f - frame_range)\n", " end_f = min(total_frames-1, target_f + frame_range)\n", " \n", " print(f\"🎯 중심초={time_in_seconds}, frame_range=±{frame_range}\")\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", " # 4개 모델별 시각화된 이미지를 담는 리스트\n", " model_vis_list = []\n", " # 모델별 간단 통계(검출개수, 평균 score, 평균 area)\n", " model_summaries = []\n", " \n", " for m_info in loaded_models:\n", " m_name = m_info[\"model_name\"]\n", " m_obj = m_info[\"model_obj\"]\n", " cls_list = m_info[\"class_names\"]\n", " thr_dict = m_info[\"score_thresholds\"]\n", " \n", " # 디텍션\n", " detections = detect_objects(frame_bgr, m_obj, cls_list, thr_dict)\n", " # 시각화 (클래스별 색상)\n", " vis_img = draw_detections_advanced(frame_bgr, detections, class_colors=CLASS_COLORS,\n", " font_scale=0.5, font_thick=1)\n", " \n", " # 모델별 이미지 저장용\n", " model_vis_list.append(vis_img)\n", " \n", " # detection info -> det_records\n", " for det in detections:\n", " rec = {\n", " \"model_name\": m_name,\n", " \"frame_idx\": fidx,\n", " \"class_name\": det[\"class_name\"],\n", " \"score\": det[\"score\"],\n", " \"width\": det[\"width\"],\n", " \"height\": det[\"height\"],\n", " \"area\": det[\"area\"]\n", " }\n", " det_records.append(rec)\n", " \n", " # 모델별 간단 통계\n", " n_det = len(detections)\n", " mean_score = np.mean([d[\"score\"] for d in detections]) if n_det>0 else 0.0\n", " mean_area = np.mean([d[\"area\"] for d in detections]) if n_det>0 else 0.0\n", " model_summaries.append((m_name, n_det, mean_score, mean_area))\n", " \n", " # 이제 model_vis_list(4장 예상)를 2x2로 합쳐 한 장 이미지\n", " if len(model_vis_list) > 0:\n", " grid_img = make_image_grid(model_vis_list, rows=2, cols=2)\n", " \n", " # 합쳐진 이미지 상단에 \"프레임 fidx\" 표시, 그리고 모델별 간략 요약\n", " # (원한다면 더 예쁜 UI로도 가능)\n", " summary_text = f\"Frame {fidx}\\n\"\n", " for (m_name, n_det, m_sc, m_ar) in model_summaries:\n", " summary_text += f\"{m_name}: #Box={n_det}, avgSc={m_sc:.2f}, avgArea={m_ar:.0f}\\n\"\n", " \n", " # 그림자(혹은 검정 배경) 위에 텍스트를 표시할 수도 있지만, 여기서는 간단히\n", " y_offset = 30\n", " for line in summary_text.split(\"\\n\"):\n", " cv2.putText(grid_img, line, (20, y_offset),\n", " cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2, cv2.LINE_AA)\n", " y_offset += 30\n", " \n", " if save_dir:\n", " base_name = os.path.splitext(os.path.basename(video_path))[0]\n", " out_fn = f\"{base_name}_frame{fidx:04d}_4model_grid.jpg\"\n", " out_path = os.path.join(save_dir, out_fn)\n", " cv2.imwrite(out_path, grid_img)\n", " # 노트북에 표시\n", " plt.figure(figsize=(10,8))\n", " plt.imshow(cv2.cvtColor(grid_img, cv2.COLOR_BGR2RGB))\n", " plt.title(f\"Frame {fidx} (4-model comparison)\")\n", " plt.axis('off')\n", " plt.show()\n", " \n", " cap.release()\n", " return det_records\n", "\n", "print(\"✅ [4] 부분 프레임 추론 & 2x2 격자 이미지 생성 함수 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [markdown]\n", "# # [5] 실제 실행 예시\n", "# \n", "# 예) time_in_seconds=10.0, frame_range=2\n", "# => 8~12프레임 (fps*10 근처)\n", "# => 각 프레임마다 4개 모델 결과를 2x2로 합쳐 1장에 저장\n", "# => 검출 정보는 df로 관리\n", "\n", "# %%\n", "# (1) 모델 로드\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", "# (2) 특정 초 주변 프레임 처리\n", "time_in_seconds = 10.0\n", "frame_range = 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,\n", " save_dir=output_folder\n", ")\n", "\n", "# (3) DataFrame 변환\n", "df_det = pd.DataFrame(det_list)\n", "print(\"df_det.shape:\", df_det.shape)\n", "print(df_det.head(10))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [markdown]\n", "# # [6] 추가 데이터/그래프 정리\n", "# \n", "# - 모델별 평균 Score, 평균 Area, 클래스별 count\n", "# - 그래프(bar chart) 저장\n", "# - CSV 저장\n", "\n", "# %%\n", "if df_det.empty:\n", " print(\"⚠️ df_det empty. 디텍션 결과가 없습니다.\")\n", "else:\n", " # 모델별 통계\n", " df_model_stats = df_det.groupby(\"model_name\").agg({\n", " \"score\":\"mean\",\n", " \"area\":\"mean\"\n", " }).reset_index()\n", " df_model_stats.rename(columns={\"score\":\"avg_score\",\"area\":\"avg_area\"}, inplace=True)\n", " \n", " print(\"=== 모델별 평균 Score / Area ===\")\n", " print(df_model_stats)\n", " \n", " df_cls_count = df_det.groupby([\"model_name\",\"class_name\"]).size().reset_index(name=\"count\")\n", " print(\"=== 모델별 클래스별 검출 개수 ===\")\n", " print(df_cls_count)\n", " \n", " # 그래프: 모델별 avg_score\n", " plt.figure(figsize=(6,4))\n", " plt.bar(df_model_stats[\"model_name\"], df_model_stats[\"avg_score\"])\n", " plt.title(\"Average Score by Model\")\n", " plt.xlabel(\"Model\")\n", " plt.ylabel(\"Mean Score\")\n", " plt.ylim(0,1)\n", " save_path1 = os.path.join(output_folder, \"avg_score_by_model.png\")\n", " plt.savefig(save_path1, dpi=150)\n", " plt.show()\n", " \n", " # 그래프: 모델별 avg_area\n", " plt.figure(figsize=(6,4))\n", " plt.bar(df_model_stats[\"model_name\"], df_model_stats[\"avg_area\"])\n", " plt.title(\"Average BBox Area by Model\")\n", " plt.xlabel(\"Model\")\n", " plt.ylabel(\"Mean Area (pixels^2)\")\n", " save_path2 = os.path.join(output_folder, \"avg_area_by_model.png\")\n", " plt.savefig(save_path2, dpi=150)\n", " plt.show()\n", " \n", " # CSV 저장\n", " csv_det = os.path.join(output_folder, \"detection_records.csv\")\n", " df_det.to_csv(csv_det, index=False, encoding=\"utf-8-sig\")\n", " print(f\"✅ 검출 디테일 CSV 저장: {csv_det}\")\n", " \n", " csv_model_stats = os.path.join(output_folder, \"summary_by_model.csv\")\n", " df_model_stats.to_csv(csv_model_stats, index=False)\n", " print(f\"✅ 모델별 요약 CSV 저장: {csv_model_stats}\")\n", " \n", " csv_cls_count = os.path.join(output_folder, \"summary_by_model_class.csv\")\n", " df_cls_count.to_csv(csv_cls_count, index=False)\n", " print(f\"✅ 모델별 클래스별 개수 CSV 저장: {csv_cls_count}\")\n", "\n", "print(\"🎉 모든 프로세스 완료. 결과를 확인하세요.\")\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 }