{ "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) 추가: 쓰러짐 감지 전용 모델 및 helmet IoU threshold 등 추가 변수 설정\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", "\n", "# mmpose 관련 임포트\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", "# yolox 관련 임포트\n", "from yolox.tracker.byte_tracker import BYTETracker\n", "from yolox.tracking_utils.timer import Timer\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 = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/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", "# 결과 영상 저장 폴더\n", "SAVE_RESULT = '/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/val_sample/output_val_samples/포스코_광주_고양_합침_whh8_fall_2'\n", "os.makedirs(SAVE_RESULT, exist_ok=True)\n", "\n", "# 디바이스 설정 (cuda:0가 있으면 사용)\n", "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n", "\n", "# ------------------------------------------------------------------\n", "# [전역 변수 설정]\n", "# 기존 디텍션 모델 (Worker, helmet, head)\n", "class_names_whh = ['Worker', 'helmet', 'head']\n", "class_score_thresholds_whh = {'Worker': 0.0, 'helmet': 0.0, 'head': 0.0}\n", "\n", "# 쓰러짐 감지 전용 모델 (fall) - Worker 클래스와 앙상블\n", "# (아래 경로는 요청하신 fall 전용 모델 config 및 weight)\n", "fall_model_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", "fall_model_weight = \"/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", "# fall 모델은 Worker 디텍션으로 간주\n", "fall_class_names = ['Worker']\n", "class_score_thresholds_fall = {'Worker': 0.0}\n", "\n", "# 슬라이딩 윈도우 관련 (디텍션 score 안정성 판정)\n", "stable_time_window = 1.0 # (초) 기본 슬라이딩 윈도우 시간\n", "score_threshold = 0.1 # 디텍션 score 임계값 (기본)\n", "\n", "# 쓰러짐 경고 슬라이딩 윈도우 (예: 10프레임 중 5프레임 이상 신뢰도 통과 시 쓰러짐 경고 지속)\n", "FALL_SLIDING_WINDOW_FRAMES = 10 \n", "FALL_REQUIRED_COUNT = 5\n", "FALL_SCORE_THRESHOLD = 0.3 # 쓰러짐 판단에 사용할 score 임계값\n", "\n", "# Helmet 착용 여부 판단 IoU 임계값 (Worker와 helmet의 IoU가 60% 미만이면 No Helmet 경고)\n", "helmet_iou_threshold = 0.6\n", "\n", "# 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", "# 비디오 리사이즈 인자 (필요시)\n", "resize_factor_video = 1.0\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}\")\n", "print(f\"[INFO] FALL_SLIDING_WINDOW_FRAMES: {FALL_SLIDING_WINDOW_FRAMES}, FALL_REQUIRED_COUNT: {FALL_REQUIRED_COUNT}\")\n", "print(f\"[INFO] Helmet IoU Threshold: {helmet_iou_threshold}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true } }, "outputs": [], "source": [ "# %% [셀 2]: 헬퍼 함수 정의 (IoU, 거리 계산 등)\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] 헬퍼 함수 (IoU, distance) 정의 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true } }, "outputs": [], "source": [ "# %% [셀 3]: 디텍션 및 Pose 모델 로드\n", "# ------------------------------------------------------------------\n", "# [A] Worker, helmet, head 모델 (whh)\n", "config_file_whh = '/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_file_whh = '/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", "\n", "os.chdir(mmyolo_dir)\n", "det_model_whh = init_detector(config_file_whh, checkpoint_file_whh, device=device)\n", "det_model_whh.dataset_meta = {'CLASSES': class_names_whh}\n", "print(\"[셀 3] whh 모델 (Worker, helmet, head) 로드 완료.\")\n", "\n", "# ------------------------------------------------------------------\n", "# [B] 쓰러짐 감지 전용 모델 (fall) 로드\n", "det_model_fall = init_detector(fall_model_config, fall_model_weight, device=device)\n", "det_model_fall.dataset_meta = {'CLASSES': fall_class_names}\n", "print(\"[셀 3] fall 감지 전용 모델 로드 완료.\")\n", "\n", "# ------------------------------------------------------------------\n", "# [C] Pose 모델 로드 (Worker에 한해서 pose-estimating)\n", "pose_config = '/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/kisa_fall_yt/weights/td-hm_res101_8xb64-210e_coco-256x192.py'\n", "pose_checkpoint = '/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/kisa_fall_yt/weights/td-hm_res101_8xb64-210e_coco-256x192-065d3625_20220926.pth'\n", "os.chdir(\"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/mmpose\")\n", "init_default_scope('mmpose')\n", "pose_estimator = init_pose_estimator(pose_config, pose_checkpoint, device=device)\n", "visualizer = VISUALIZERS.build(pose_estimator.cfg.visualizer)\n", "visualizer.set_dataset_meta(pose_estimator.dataset_meta)\n", "print(\"[셀 3] Pose 모델 로드 완료.\")\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와 텍스트를 시각화하여 반환.\n", " - 객체명, 트랙 번호 등은 여러 줄로 표시하여 서로 겹치지 않도록 함.\n", " - Worker: 텍스트는 bbox 상단에 가로로 출력하며,\n", " fall_alert가 True이면 박스와 텍스트는 빨간색, 그렇지 않으면 초록색으로 표시.\n", " - helmet, head: 텍스트는 bbox 상단에 가로로 출력 (기존 vertical 출력 대신)\n", " \"\"\"\n", " output = frame.copy()\n", " global_warnings = [] # 영상 상단 전역 경고용 (필요시 활용)\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", " track_id = det.get('track_id', None)\n", " # 기본 색상 설정: Worker는 조건에 따라 결정, helmet과 head는 기존 색상 유지\n", " if lbl == 'Worker':\n", " # Worker: fall_alert가 True이면 빨간색, 아니면 초록색\n", " if det.get('fall_alert', False):\n", " color = (0, 0, 255) # 빨강\n", " else:\n", " color = (0, 255, 0) # 초록\n", " elif lbl == 'helmet':\n", " color = (0, 128, 255) # 주황\n", " elif lbl == 'head':\n", " color = (255, 0, 255) # 마젠타\n", " else:\n", " color = (128, 128, 128)\n", " \n", " # bbox 그리기\n", " cv2.rectangle(output, (x1, y1), (x2, y2), color, 2)\n", " \n", " # 텍스트 내용 (여러 줄로: 객체명, 트랙번호, score, fall 경고 등)\n", " text_lines = []\n", " text_lines.append(f\"{lbl}\")\n", " if track_id is not None:\n", " text_lines.append(f\"ID: {track_id}\")\n", " text_lines.append(f\"{sc:.2f}\")\n", " if det.get('fall_alert', False):\n", " text_lines.append(\"Fall Detected\")\n", " \n", " # 모든 객체는 텍스트를 bbox 상단에 가로로 출력 (여러 줄로)\n", " offset_y = y1 - 5\n", " for line in text_lines:\n", " (text_w, text_h), _ = cv2.getTextSize(line, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 2)\n", " cv2.putText(output, line, (x1, offset_y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, 2)\n", " offset_y -= (text_h + 5)\n", " \n", " # (글로벌 경고 메시지 출력이 필요하다면 아래 코멘트 해제)\n", " # if global_warnings:\n", " # warning_text = \" / \".join(global_warnings)\n", " # cv2.putText(output, warning_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,0,255), 3)\n", " \n", " return output\n", "\n", "print(\"[셀 4] 디텍션 및 시각화 함수 정의 완료.\")\n" ] }, { "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] 여백 추가 함수 정의 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true } }, "outputs": [], "source": [ "# %% [셀 6]: 단일 프레임 디텍션 예시 함수 (옵션)\n", "def show_detection_single_frame(video_path, frame_idx, det_models, score_thresh_inits, class_names_list, scale_factor=1.0, font_scale=0.5):\n", " \"\"\"\n", " 주어진 영상의 특정 프레임(frame_idx)에 대해,\n", " 각 모델의 디텍션 결과를 시각화하는 예시 함수.\n", " \"\"\"\n", " cap = cv2.VideoCapture(video_path)\n", " if not cap.isOpened():\n", " print(f\"[오류] 영상 열기 실패: {video_path}\")\n", " return\n", " cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)\n", " ret, frame = cap.read()\n", " cap.release()\n", " if not ret:\n", " print(\"[오류] 해당 프레임 없음\")\n", " return\n", "\n", " frame = add_margin_frame(frame, 0, 0)\n", " all_detections = []\n", " for model, score_thresh_init, cls_list in zip(det_models, score_thresh_inits, class_names_list):\n", " dets = detect_objects(frame, model, score_thresh_init, cls_list)\n", " all_detections.extend(dets)\n", " out_frame = draw_detections(frame, all_detections, font_scale)\n", " if scale_factor != 1.0:\n", " oh, ow = out_frame.shape[:2]\n", " out_frame = cv2.resize(out_frame, (int(ow * scale_factor), int(oh * scale_factor)))\n", " plt.figure(figsize=(10,10))\n", " plt.imshow(cv2.cvtColor(out_frame, cv2.COLOR_BGR2RGB))\n", " plt.axis('off')\n", " plt.show()\n", "\n", "print(\"[셀 6] 단일 프레임 디텍션 예시 함수 정의 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true } }, "outputs": [], "source": [ "# %% [셀 7]: ByteTracker 기반 추적 및 슬라이딩 윈도우 업데이트 함수\n", "def create_worker_tracker():\n", " \"\"\"\n", " Worker, helmet, head 등 추적을 위한 ByteTracker 생성 함수.\n", " \"\"\"\n", " return BYTETracker(tracking_args)\n", "\n", "def update_worker_track_states(tracks, track_dict, fps, stable_time_window, score_threshold,\n", " fall_window_size, fall_required_count, fall_score_threshold):\n", " \"\"\"\n", " [기능]\n", " - 슬라이딩 윈도우 방식으로 score 통과 여부를 업데이트.\n", " (기존 쓰러짐 sliding window 코드는 제거하여, Pose 기반 판정을 사용하도록 함)\n", " \n", " Parameters:\n", " tracks: ByteTracker 업데이트 결과 (각 트랙 정보)\n", " track_dict: 트랙 정보를 저장하는 dict\n", " fps: 영상 FPS\n", " stable_time_window: score 판정을 위한 슬라이딩 윈도우 (초)\n", " score_threshold: 디텍션 score 임계값\n", " fall_window_size, fall_required_count, fall_score_threshold: (현재 사용하지 않음)\n", " \"\"\"\n", " stable_window_size = int(fps * stable_time_window)\n", " required_pass_count = int(stable_window_size * 0.67)\n", "\n", " for t in tracks:\n", " tid = t.track_id\n", " if tid not in track_dict:\n", " track_dict[tid] = {\n", " 'last_center': None,\n", " 'score_history': [],\n", " 'score_passed': False\n", " }\n", " # 현재 bbox 정보 (tlwh: x, y, w, h)\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", " _ = distance(center_now, last_center)\n", " track_dict[tid]['last_center'] = center_now\n", " \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", " 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\n", "\n", "print(\"[셀 7] ByteTracker 기반 추적 및 슬라이딩 윈도우 업데이트 함수 정의 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [셀 8]: Pose 추론 및 쓰러짐 감지 함수 정의\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", " # COCO 기준 keypoint 인덱스 예시\n", " LEFT_SHOULDER = 5\n", " RIGHT_SHOULDER = 6\n", " LEFT_HIP = 11\n", " RIGHT_HIP = 12\n", " LEFT_ANKLE = 15\n", " RIGHT_ANKLE = 16\n", " LEFT_KNEE = 13\n", " RIGHT_KNEE = 14\n", "\n", " left_shoulder = keypoints[LEFT_SHOULDER]\n", " right_shoulder = keypoints[RIGHT_SHOULDER]\n", " left_hip = keypoints[LEFT_HIP]\n", " right_hip = keypoints[RIGHT_HIP]\n", " left_ankle = keypoints[LEFT_ANKLE]\n", " right_ankle = keypoints[RIGHT_ANKLE]\n", "\n", " # 대략적인 몸통 길이 계산\n", " len_factor_left = np.hypot(left_shoulder[0] - left_hip[0], left_shoulder[1] - left_hip[1])\n", " len_factor_right = np.hypot(right_shoulder[0] - right_hip[0], right_shoulder[1] - right_hip[1])\n", " len_factor = max(len_factor_left, len_factor_right)\n", "\n", " xmin, ymin, xmax, ymax = bbox\n", " dx = xmax - xmin\n", " dy = ymax - ymin\n", " difference = dy - dx\n", "\n", " fall_detected = False\n", " if ((left_shoulder[1] > left_ankle[1] - len_factor and left_hip[1] > left_ankle[1] - (len_factor / 2)) or\n", " (right_shoulder[1] > right_ankle[1] - len_factor and right_hip[1] > right_ankle[1] - (len_factor / 2)) or\n", " (difference < 0)):\n", " if (keypoint_scores[LEFT_KNEE] >= 0.5 or keypoint_scores[RIGHT_KNEE] >= 0.5 or\n", " keypoint_scores[LEFT_ANKLE] >= 0.5 or keypoint_scores[RIGHT_ANKLE] >= 0.5):\n", " fall_detected = True\n", "\n", " return fall_detected, bbox\n", "\n", "def pose_inference_for_workers(frame, worker_dets, pose_estimator, visualizer):\n", " \"\"\"\n", " Worker 디텍션 bbox들에 대해 Pose 모델 추론을 수행하고,\n", " 각 인물에 대해 쓰러짐 감지를 수행하여, 스켈레톤 시각화 이미지와\n", " 각 Worker에 대한 fall 감지 결과(True/False 리스트)를 반환.\n", " \"\"\"\n", " if len(worker_dets) == 0:\n", " return frame, []\n", "\n", " # Worker bbox들 추출\n", " bboxes = [d['bbox'] for d in worker_dets]\n", " bboxes = np.array(bboxes)\n", " results = inference_topdown(pose_estimator, frame, bboxes, bbox_format='xyxy')\n", " data_samples = merge_data_samples(results)\n", " \n", " # 스켈레톤 시각화\n", " visualizer.add_datasample(\n", " 'result',\n", " frame,\n", " data_sample=data_samples,\n", " draw_gt=False,\n", " draw_heatmap=False,\n", " draw_bbox=False,\n", " show=False,\n", " wait_time=0,\n", " out_file=None,\n", " kpt_thr=0.3\n", " )\n", " out_img = visualizer.get_image()\n", " \n", " # 각 인물별로 쓰러짐 감지 수행 및 fall 결과 리스트 생성\n", " fall_flags = []\n", " if hasattr(data_samples, 'pred_instances') and len(data_samples.pred_instances) > 0:\n", " # data_samples.pred_instances에 있는 순서와 worker_dets 순서가 동일하다고 가정\n", " for i in range(len(data_samples.pred_instances)):\n", " single_data_sample = data_samples.clone()\n", " single_data_sample.pred_instances = single_data_sample.pred_instances[i:i+1]\n", " fall_detected, _ = fall_detection_single(single_data_sample)\n", " fall_flags.append(fall_detected)\n", " # (원한다면 여기서 각 인물에 대해 개별 경고 표시를 추가할 수 있음)\n", " else:\n", " # 만약 결과가 없다면 모두 False로 처리\n", " fall_flags = [False] * len(worker_dets)\n", " \n", " return out_img, fall_flags\n", "\n", "print(\"[셀 8] Pose 추론 및 쓰러짐 감지 함수 정의 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [셀 9]: 전체 영상 처리 (앙상블, 트래킹, 경고, Pose 등)\n", "def process_video(video_path, \n", " det_models, # [det_model_whh, det_model_fall]\n", " score_thresh_inits, # [class_score_thresholds_whh, class_score_thresholds_fall]\n", " class_names_lists, # [class_names_whh, fall_class_names]\n", " stable_time_window=stable_time_window,\n", " fall_window_size=FALL_SLIDING_WINDOW_FRAMES,\n", " fall_required_count=FALL_REQUIRED_COUNT,\n", " score_threshold=score_threshold,\n", " fall_score_threshold=FALL_SCORE_THRESHOLD,\n", " helmet_iou_thr=helmet_iou_threshold,\n", " do_inference=True, \n", " do_save=True, \n", " resize_factor=resize_factor_video):\n", " \"\"\"\n", " 영상에 대해 아래 작업을 수행:\n", " 1) whh 및 fall 모델을 이용해 디텍션 수행 (Worker, helmet, head)\n", " - fall 모델 결과는 Worker와 통합하여 앙상블함.\n", " 2) ByteTracker를 통한 트래킹 및 슬라이딩 윈도우 기반 score 업데이트\n", " 3) Worker 디텍션에 대해 Pose 추론을 수행하여 쓰러짐 감지(결과는 fall_alert로 반영)\n", " 4) 최종 결과(디텍션 bbox, 트랙 id, 경고 등)를 영상에 시각화 후 저장.\n", " - Worker: fall_alert가 True이면 박스와 텍스트는 빨간색, 그렇지 않으면 초록색.\n", " - helmet, head: 텍스트는 bbox 상단에 가로로 출력.\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", " # 디텍션 모델 언팩: det_model_whh와 det_model_fall\n", " whh_model, fall_model = det_models\n", " whh_score = score_thresh_inits[0]\n", " fall_score = score_thresh_inits[1]\n", " whh_cls = class_names_lists[0]\n", " fall_cls = class_names_lists[1]\n", "\n", " # 트래커 생성 (Worker, helmet, head 각각)\n", " worker_tracker = create_worker_tracker()\n", " helmet_tracker = create_worker_tracker()\n", " head_tracker = create_worker_tracker()\n", "\n", " worker_dict = {}\n", " helmet_dict = {}\n", " head_dict = {}\n", "\n", " print(f\"[INFO] 영상 처리 시작 -> {video_path}\")\n", " start_total = time.time()\n", "\n", " while True:\n", " st = time.time()\n", " ret, frame = cap.read()\n", " if not ret:\n", " break\n", "\n", " # 1) 디텍션 수행\n", " # whh 모델: Worker, helmet, head 디텍션\n", " d_whh = detect_objects(frame, whh_model, whh_score, whh_cls)\n", " for d in d_whh:\n", " d['source'] = 'whh'\n", " # fall 모델: fall 전용 모델 -> Worker 디텍션으로 간주\n", " d_fall = detect_objects(frame, fall_model, fall_score, fall_cls)\n", " for d in d_fall:\n", " d['source'] = 'fall'\n", " # Worker 디텍션: whh와 fall 결과를 통합 (label이 'Worker'인 것)\n", " worker_dets = [d for d in d_whh if d['label'] == 'Worker'] + [d for d in d_fall if d['label'] == 'Worker']\n", " # helmet, head 디텍션은 whh 모델 결과에서만 사용\n", " helmet_dets = [d for d in d_whh if d['label'] == 'helmet']\n", " head_dets = [d for d in d_whh if d['label'] == 'head']\n", "\n", " # 2) 트래커 업데이트\n", " def prepare_xyxyscore(dets):\n", " xyxyscore = []\n", " for d in dets:\n", " x1, y1, x2, y2 = d['bbox']\n", " s = d['score']\n", " xyxyscore.append([x1, y1, x2, y2, s])\n", " return np.array(xyxyscore)\n", " \n", " worker_xyxyscore = prepare_xyxyscore(worker_dets)\n", " helmet_xyxyscore = prepare_xyxyscore(helmet_dets)\n", " head_xyxyscore = prepare_xyxyscore(head_dets)\n", " \n", " online_worker = worker_tracker.update(worker_xyxyscore, [h_org, w_org])\n", " online_helmet = helmet_tracker.update(helmet_xyxyscore, [h_org, w_org])\n", " online_head = head_tracker.update(head_xyxyscore, [h_org, w_org])\n", " \n", " # 3) 트랙 상태 업데이트 (score sliding window만 적용)\n", " update_worker_track_states(online_worker, worker_dict, fps, stable_time_window, score_threshold,\n", " fall_window_size, FALL_REQUIRED_COUNT, FALL_SCORE_THRESHOLD)\n", " update_worker_track_states(online_helmet, helmet_dict, fps, stable_time_window, score_threshold,\n", " fall_window_size, FALL_REQUIRED_COUNT, FALL_SCORE_THRESHOLD)\n", " update_worker_track_states(online_head, head_dict, fps, stable_time_window, score_threshold,\n", " fall_window_size, FALL_REQUIRED_COUNT, FALL_SCORE_THRESHOLD)\n", " \n", " # 4) 최종 Worker, helmet, head 결과 생성\n", " final_worker = []\n", " for t in online_worker:\n", " if not worker_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", " for wd in worker_dets:\n", " if iou_calc(box_now, wd['bbox']) > 0.5:\n", " found_source = wd.get('source', None)\n", " break\n", " # Worker entry에 fall_alert 플래그는 이후 Pose 추론 결과로 업데이트할 예정\n", " worker_entry = {\n", " 'bbox': box_now,\n", " 'score': t.score,\n", " 'label': 'Worker',\n", " 'source': found_source,\n", " 'track_id': t.track_id\n", " }\n", " final_worker.append(worker_entry)\n", " \n", " final_helmet = []\n", " for t in online_helmet:\n", " if not helmet_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", " final_helmet.append({\n", " 'bbox': [x1, y1, x2, y2],\n", " 'score': t.score,\n", " 'label': 'helmet',\n", " 'track_id': t.track_id\n", " })\n", " \n", " final_head = []\n", " for t in online_head:\n", " if not head_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", " final_head.append({\n", " 'bbox': [x1, y1, x2, y2],\n", " 'score': t.score,\n", " 'label': 'head',\n", " 'track_id': t.track_id\n", " })\n", " \n", " # 5) Worker에 대해 Pose 추론 수행 (쓰러짐 감지용)\n", " # pose_inference_for_workers는 (이미지, fall_flags 리스트)를 반환함\n", " pose_frame, fall_flags = pose_inference_for_workers(frame, final_worker, pose_estimator, visualizer)\n", " # fall_flags의 순서는 final_worker 순서와 동일하다고 가정하여 업데이트\n", " for idx_w, worker_entry in enumerate(final_worker):\n", " worker_entry['fall_alert'] = fall_flags[idx_w] if idx_w < len(fall_flags) else False\n", " \n", " # 6) 최종 시각화를 위해 모든 디텍션 결과 병합 (Worker, helmet, head)\n", " all_final = []\n", " all_final.extend(final_worker)\n", " all_final.extend(final_helmet)\n", " all_final.extend(final_head)\n", " vis_frame = draw_detections(pose_frame, all_final, 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", "\n", " cap.release()\n", " if writer is not None:\n", " writer.release()\n", " if frame_times:\n", " avg_fps = 1/np.mean(frame_times)\n", " print(f\"[INFO] 평균 FPS: {avg_fps:.2f}\")\n", " print(f\"[INFO] 영상 처리 완료 -> {video_path}\")\n", " print(f\"[INFO] 결과 영상 저장 -> {output_video_path}\")\n", "\n", "print(\"[셀 9] 전체 영상 처리 함수 정의 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [셀 10]: 실행부 (영상 폴더 내 모든 파일 처리)\n", "video_inference_on = True\n", "video_save_on = True\n", "\n", "# 두 디텍션 모델 (whh와 fall)을 리스트로 전달\n", "det_models = [det_model_whh, det_model_fall]\n", "scores_inits = [class_score_thresholds_whh, class_score_thresholds_fall]\n", "names_lists = [class_names_whh, fall_class_names]\n", "\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", " det_models=det_models,\n", " score_thresh_inits=scores_inits,\n", " class_names_lists=names_lists,\n", " stable_time_window=stable_time_window,\n", " fall_window_size=FALL_SLIDING_WINDOW_FRAMES,\n", " fall_required_count=FALL_REQUIRED_COUNT,\n", " score_threshold=score_threshold,\n", " fall_score_threshold=FALL_SCORE_THRESHOLD,\n", " helmet_iou_thr=helmet_iou_threshold,\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] 전체 영상 처리 완료.\")\n" ] } ], "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 }