{ "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 = \"/DATA2/ltb/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", "VIDEO_DIR = ('/DATA2/ltb/mmyolo/it1_hd_val_sample/샘플박스/포스코_광주_고양_합침')\n", "SAVE_RESULT = ('/DATA2/ltb/mmyolo/it1_hd_val_sample/포스코_광주_고양_합침_conequip_helmet_head_cutmix_resize_rotate_1819_ff_base_1,1,1_epoch36_4090_01')\n", "os.makedirs(SAVE_RESULT, exist_ok=True)\n", "\n", "# VIDEO_EXTENSIONS = ['.mp4', '.avi', '.mov', '.mkv']\n", "# video_files = []\n", "# for ext in VIDEO_EXTENSIONS:\n", "# video_files.extend(glob.glob(os.path.join(VIDEO_DIR, f'*{ext}')))\n", "\n", "# existing_exps = [d for d in os.listdir(SAVE_RESULT) if d.startswith('exp')]\n", "# exp_numbers = [int(d.replace('exp', '')) for d in existing_exps if d.replace('exp', '').isdigit()]\n", "# next_exp_num = max(exp_numbers) + 1 if exp_numbers else 1\n", "# vis_folder = os.path.join(SAVE_RESULT, f'exp{next_exp_num}')\n", "# os.makedirs(vis_folder, exist_ok=True)\n", "\n", "# 클래스 이름 설정\n", "class_names_whh = ['worker', 'head', 'helmet']\n", "# class_names_ce = ['con_equip']\n", "# class_names_fall = ['Worker', 'background']\n", "\n", "# 디텍션 스코어 임계값\n", "class_score_thresholds_whh_init = {'worker': 0.0, 'head': 0.0, 'helmet': 0.0}\n", "# class_score_thresholds_ce_init = {'con_equip': 0.05}\n", "# class_score_thresholds_fall_init = {'Worker': 0.0}\n", "\n", "# # 중장비 IoU, 움직임 임계값, 안정 시간(초)\n", "# iou_threshold_for_filtering = 0.7 # <-- 중장비와 겹치는 IoU 임계값\n", "# movement_threshold = 500 # <-- 중장비 겹칠 때 '안정'으로 볼 움직임 범위\n", "# stable_time_con_equip = 0.2 # <-- 위 조건이 몇 초 지속되어야 필터링하는지 (중장비)\n", "\n", "# # Worker, helmet, head 움직임 필터링\n", "\n", "# still_move_thrshold = 10 # <= 10픽셀이면 거의 안 움직인다고 보기\n", "# stable_time_no_move = 180 # 2초간 안 움직이면 필터링\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 # 윈도우 내 통과비율 (아래 stable_time_window 참고)\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}\")\n" ] }, { "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] 헬퍼 함수 정의 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true } }, "outputs": [], "source": [ "# %% [셀 3]: 세 가지 디텍션 모델 + Pose 모델 로드\n", "import sys\n", "sys.path.append(\"/DATA2/ltb/mmdetection\")\n", "\n", "# [A] Worker/helmet/head 모델 (whh)\n", "config_file_whh = ('/DATA2/ltb/mmyolo/work_dirs/conequip_helmet_head_cutmix_resize_rotate_1819_ff_base_1,1,1_01/4090/conequip_helmet_head_cutmix_resize_rotate_1819_ff_base_1,1,1.py')\n", "checkpoint_file_whh = ('/DATA2/ltb/mmyolo/work_dirs/conequip_helmet_head_cutmix_resize_rotate_1819_ff_base_1,1,1_01/4090/epoch_36.pth')\n", "\n", "# # [B] 중장비(con_equip) 모델\n", "# config_file_ce = os.path.join(\n", "# BASE_DIR,\n", "# 'mmyolo/work_dirs/it1_960_conequip_01/yolov8_l_960_custom_only_conequip.py'\n", "# )\n", "# checkpoint_file_ce = os.path.join(\n", "# BASE_DIR,\n", "# 'mmyolo/work_dirs/it1_960_conequip_01/epoch_3.pth'\n", "# )\n", "\n", "# # [C] 쓰러짐(fall) 모델\n", "# config_file_fall = ( '/DATA2/ltb/mmyolo/work_dirs/it1_roboflow_fall_only_rotate_960_yolo_01/yolov8_l_960_custom_fall.py' )\n", "# checkpoint_file_fall = ( '/DATA2/ltb/mmyolo/work_dirs/it1_roboflow_fall_only_rotate_960_yolo_01/best_coco_bbox_mAP_epoch_13.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", "# det_model_ce = init_detector(config_file_ce, checkpoint_file_ce, device=device)\n", "# det_model_ce.dataset_meta = {'CLASSES': class_names_ce}\n", "# det_model_fall = init_detector(config_file_fall, checkpoint_file_fall, device=device)\n", "# det_model_fall.dataset_meta = {'CLASSES': class_names_fall}\n", "\n", "print(\"[모델 로드 경로 안내]\")\n", "print(f\" [INFO] whh config: {config_file_whh}\")\n", "print(f\" [INFO] whh weight: {checkpoint_file_whh}\")\n", "# print(f\" [INFO] ce config: {config_file_ce}\")\n", "# print(f\" [INFO] ce weight: {checkpoint_file_ce}\")\n", "# print(f\" [INFO] fall config: {config_file_fall}\")\n", "# print(f\" [INFO] fall weight: {checkpoint_file_fall}\")\n", "\n", "# (주석) Pose 모델은 현재 사용 안 하므로 주석 처리\n", "# pose_config = '/DATA2/ltb/root/kisa_fall_yt/weights/td-hm_res101_8xb64-210e_coco-256x192.py'\n", "# pose_checkpoint = '/DATA2/ltb/root/kisa_fall_yt/weights/td-hm_res101_8xb64-210e_coco-256x192-065d3625_20220926.pth'\n", "# print(f\" [INFO] Pose config: {pose_config}\")\n", "# print(f\" [INFO] Pose weight: {pose_checkpoint}\")\n", "\n", "# os.chdir(\"/DATA2/ltb/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", "\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':[x1,y1,x2,y2], 'score':..., 'label':..., 'source':..., 'track_id':...}, ...]\n", " \"\"\"\n", " output = frame.copy()\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", " # 라벨별 색상 설정\n", " if lbl == 'worker':\n", " color = (255, 0, 0) # 파랑\n", " # if 'source' in det:\n", " # if det['source'] == 'whh':\n", " # color = (255, 0, 0) # 파랑\n", " # elif det['source'] == 'fall':\n", " # color = (0, 165, 165) # 청록\n", " elif lbl == 'helmet':\n", " color = (0, 165, 255) # 주황\n", " # elif lbl == 'con_equip':\n", " # color = (0, 255, 0) # 녹색\n", " elif lbl == 'head':\n", " color = (165, 0, 165) # 보라\n", " else:\n", " color = (128, 128, 128)\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] 디텍션 및 시각화 함수 정의 완료.\")\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", "\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", "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", "\n", "def update_track_states(\n", " tracks, \n", " track_dict, \n", " # ce_dets, \n", " fps,\n", " # iou_thr, \n", " # move_thr, \n", " # stable_time_con_equip,\n", " stable_time_window, \n", " score_threshold,\n", " # 아래 2개는 새로 추가된 '멈춤' 필터링에 사용:\n", " # still_move_thr=10, # <-- 움직임이 이 픽셀 이하이면 '멈춤'으로 간주\n", " # stable_time_no_movement=2.0 # <-- '멈춤'이 이 초 이상 지속되면 필터링\n", "):\n", " \"\"\"\n", " [주요 기능 - 2가지 필터 로직 동작]\n", " 1) 중장비(con_equip)와 겹치는(overlap=True) 상태에서,\n", " move_thr 이하의 움직임으로 stable_time_con_equip(초) 이상 지속되면\n", " is_filtered = True (가시화 제외)\n", "\n", " 2) 일정시간( stable_time_no_movement 초 ) 동안\n", " 거의 움직임(still_move_thr 이하)이 없으면 is_filtered = True\n", "\n", " [추가로 슬라이딩 윈도우 기반 score_passed 여부도 갱신]\n", "\n", " Parameters:\n", " tracks: ByteTracker 업데이트 결과 (각 트랙 정보)\n", " track_dict: 트랙 정보를 저장하는 dict\n", " ce_dets: 중장비(con_equip) 디텍션 목록\n", " fps: 영상 FPS\n", " iou_thr: 중장비와 겹쳤다고 판단할 IoU 임계값\n", " move_thr: '중장비 위에서' 안정으로 볼 움직임 범위\n", " stable_time_con_equip: 중장비 위에서 move_thr 이하 움직임이 이 초 이상이면 필터\n", " stable_time_window: 슬라이딩 윈도우 크기(초) => score 통과 여부 계산용\n", " score_threshold: detection score 임계값\n", " still_move_thr: '멈춤(거의 안 움직임)'으로 볼 픽셀 이동 범위 (새로 추가)\n", " stable_time_no_movement: 위 멈춤상태가 이 초 이상 지속되면 필터 처리 (새로 추가)\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", " # ce_boxes = [d['bbox'] for d in ce_dets]\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", " # 'continuous_stable_count': 0, # <-- 중장비 겹침+안정 카운트\n", " # 'continuous_still_count': 0, # <-- 추가한 '멈춤' 카운트\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", " # # 1) 중장비와의 겹침(overlap) 여부 확인\n", " # overlap = False\n", " # box_now = [x1, y1, x2, y2]\n", " # for ce_box in ce_boxes:\n", " # iou_val = iou_calc(box_now, ce_box)\n", " # if iou_val >= iou_thr:\n", " # overlap = True\n", " # break\n", "\n", " # 2) center_now 계산 및 각종 카운트 처리\n", " last_center = track_dict[tid]['last_center']\n", "\n", " # center가 None이면(초기) 업데이트 후 다음 로직 진행\n", " if last_center is None:\n", " track_dict[tid]['last_center'] = center_now\n", " else:\n", " dist_val = distance(center_now, last_center)\n", "\n", " # # (2-1) 중장비 겹침(stable_time_con_equip) 관련 로직\n", " # if overlap:\n", " # # 움직임이 move_thr 이하이면 안정 카운트++\n", " # if dist_val <= move_thr:\n", " # track_dict[tid]['continuous_stable_count'] += 1\n", " # else:\n", " # track_dict[tid]['continuous_stable_count'] = 0\n", "\n", " # # 만약 stable_time_con_equip(초) * fps 이상 유지되었다면 필터\n", " # if track_dict[tid]['continuous_stable_count'] >= int(fps * stable_time_con_equip):\n", " # track_dict[tid]['is_filtered'] = True\n", " # else:\n", " # # 중장비와 겹치지 않는 동안은 겹침 카운트 초기화\n", " # track_dict[tid]['continuous_stable_count'] = 0\n", "\n", " # # (2-2) 추가: '멈춤(거의 안 움직임)' 로직\n", " # if dist_val <= still_move_thr:\n", " # track_dict[tid]['continuous_still_count'] += 1\n", " # else:\n", " # track_dict[tid]['continuous_still_count'] = 0\n", "\n", " # # 일정 시간( stable_time_no_movement ) 이상 멈춰 있으면 필터\n", " # if track_dict[tid]['continuous_still_count'] >= int(fps * stable_time_no_movement):\n", " # track_dict[tid]['is_filtered'] = True\n", "\n", " # 매번 last_center 갱신\n", " track_dict[tid]['last_center'] = center_now\n", "\n", " # 3) 슬라이딩 윈도우 기반 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\n", "\n" ] }, { "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] # (17, 2)\n", "# keypoint_scores = data_sample.pred_instances.keypoint_scores[0] # (17,)\n", "# bbox = data_sample.pred_instances.bboxes[0] # (4,)\n", "\n", "# # ... (중략: 쓰러짐 판단 로직) ...\n", "# fall_detected = False\n", "# # ... (임의 조건)\n", "\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] 쓰러짐 감지 함수 및 경고 시각화 함수 정의 완료.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true } }, "outputs": [], "source": [ "# %% [셀 9]: 전체 영상 처리 (앙상블) + 트래킹 및 필터링\n", "\n", "def process_video(\n", " video_path, \n", " models, \n", " score_thresh_inits, \n", " class_names_lists,\n", " # iou_threshold_for_filtering=0.7, \n", " # move_threshold=50,\n", " # stable_time_con_equip=1.0, \n", " stable_time_window=1.0,\n", " do_inference=True, \n", " do_save=True, \n", " resize_factor=1.0,\n", " # 아래 2개는 '멈춤' 필터링에 사용할 새 파라미터 (원하면 조절)\n", " # still_move_thr=10,\n", " # stable_time_no_movement=2.0\n", "):\n", " \"\"\"\n", " 영상에 대해 다음 작업을 수행:\n", " 1) whh/ce/fall 모델로 디텍션\n", " 2) ByteTracker로 Worker, helmet, head를 각각 추적\n", " 3) 중장비와 큰 IoU이면서 움직임이 없는 경우 (stable_time_con_equip 초 이상) => 필터 처리\n", " 4) '멈춤(거의 안 움직임)'도 일정 시간( stable_time_no_movement 초 ) 지속 => 필터\n", " 5) 슬라이딩 윈도우 기반 score 필터링\n", " 6) 최종 시각화 후 영상 저장\n", "\n", " Parameters\n", " ----------\n", " iou_threshold_for_filtering : float\n", " 중장비와 겹쳤다고 보는 IoU 임계값.\n", " move_threshold : float\n", " 중장비 위에서 안정으로 볼 픽셀 움직임 범위.\n", " stable_time_con_equip : float\n", " 해당 안정 상태가 몇 초 지속되면 필터링할지.\n", " stable_time_window : float\n", " 슬라이딩 윈도우 크기(초). 이 안에서 score_threshold 이상의 비율이 일정 이상이면 score_passed.\n", " still_move_thr : float\n", " '멈춤(거의 안 움직임)'으로 간주할 이동량(픽셀).\n", " stable_time_no_movement : float\n", " 위 '멈춤' 상태가 이 초 이상이면 필터링.\n", "\n", " Returns\n", " -------\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", " # 언팩: whh, ce, fall 모델\n", " whh_model = models[0]\n", " whh_score = score_thresh_inits[0]\n", " whh_cls = class_names_lists[0]\n", "\n", " # # 트래커 생성\n", " worker_tracker = create_worker_helmet_head_tracker()\n", " helmet_tracker = create_worker_helmet_head_tracker()\n", " head_tracker = create_worker_helmet_head_tracker()\n", "\n", " # 트랙별 상태 저장할 dict\n", " worker_dict = {}\n", " helmet_dict = {}\n", " head_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) 세 모델 디텍션\n", " all_detections = []\n", " # whh 모델\n", " d_whh = detect_objects(frame, whh_model, whh_score, whh_cls)\n", " for dw in d_whh:\n", " if dw['label'] == 'worker':\n", " dw['source'] = 'whh'\n", " all_detections.extend(d_whh)\n", "\n", " # # 중장비\n", " # d_ce = detect_objects(frame, ce_model, ce_score, ce_cls)\n", " # all_detections.extend(d_ce)\n", "\n", " # # fall 모델\n", " # d_fall = detect_objects(frame, fall_model, fall_score, fall_cls)\n", " # for df in d_fall:\n", " # if df['label'] == 'Worker':\n", " # df['source'] = 'fall'\n", " # all_detections.extend(d_fall)\n", "\n", " # (2) 분류별로 구분\n", " # ce_dets = [d for d in all_detections if d['label'] == 'con_equip']\n", " worker_dets = [d for d in all_detections if d['label'] == 'worker']\n", " helmet_dets = [d for d in all_detections if d['label'] == 'helmet']\n", " head_dets = [d for d in all_detections if d['label'] == 'head']\n", "\n", " # (3) ByteTracker에 입력\n", " # 형식 [x1, y1, x2, y2, score]\n", " worker_xyxyscore = []\n", " for w in worker_dets:\n", " x1, y1, x2, y2 = w['bbox']\n", " s = w['score']\n", " worker_xyxyscore.append([x1, y1, x2, y2, s])\n", "\n", " helmet_xyxyscore = []\n", " for h in helmet_dets:\n", " x1, y1, x2, y2 = h['bbox']\n", " s = h['score']\n", " helmet_xyxyscore.append([x1, y1, x2, y2, s])\n", "\n", " head_xyxyscore = []\n", " for hd in head_dets:\n", " x1, y1, x2, y2 = hd['bbox']\n", " s = hd['score']\n", " head_xyxyscore.append([x1, y1, x2, y2, s])\n", "\n", " # 트래커 업데이트\n", " online_worker = worker_tracker.update(np.array(worker_xyxyscore), [h_org, w_org])\n", " online_helmet = helmet_tracker.update(np.array(helmet_xyxyscore), [h_org, w_org])\n", " online_head = head_tracker.update(np.array(head_xyxyscore), [h_org, w_org])\n", "\n", " # (4) update_track_states: 중장비/움직임/슬라이딩윈도우/멈춤 로직 등\n", " update_track_states(\n", " online_worker, worker_dict, \n", " # ce_dets, \n", " fps,\n", " # iou_threshold_for_filtering, \n", " # move_threshold,\n", " # stable_time_con_equip, \n", " stable_time_window, score_threshold,\n", " # still_move_thr=still_move_thr,\n", " # stable_time_no_movement=stable_time_no_movement\n", " )\n", " update_track_states(\n", " online_helmet, helmet_dict, \n", " # ce_dets, \n", " fps,\n", " # iou_threshold_for_filtering, \n", " # move_threshold,\n", " # stable_time_con_equip, \n", " stable_time_window, score_threshold,\n", " # still_move_thr=still_move_thr,\n", " # stable_time_no_movement=stable_time_no_movement\n", " )\n", " update_track_states(\n", " online_head, head_dict, \n", " # ce_dets, \n", " fps,\n", " # iou_threshold_for_filtering, \n", " # move_threshold,\n", " # stable_time_con_equip, \n", " stable_time_window, score_threshold,\n", " # still_move_thr=still_move_thr,\n", " # stable_time_no_movement=stable_time_no_movement\n", " )\n", "\n", " # (5) 최종 Worker/helmet/head 선별\n", " # ==> is_filtered=False 이고, score_passed=True 인 경우만 통과\n", " # final_worker = worker_dets\n", " # final_helmet = helmet_dets\n", " final_worker = []\n", " for t in online_worker:\n", " if worker_dict[t.track_id].get('is_filtered', False):\n", " continue # <-- 추가: 필터 처리된 경우 skip\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", " final_worker.append({\n", " 'bbox': [x1, y1, x2, y2],\n", " 'score': t.score,\n", " 'label': 'worker',\n", " 'source': found_source,\n", " 'track_id': t.track_id\n", " })\n", "\n", " final_helmet = []\n", " for t in online_helmet:\n", " if helmet_dict[t.track_id].get('is_filtered', False):\n", " continue # <-- 추가\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 head_dict[t.track_id].get('is_filtered', False):\n", " continue # <-- 추가\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", " # final_ce = ce_dets # 중장비는 별도 추적 사용 안 함 (그냥 디텍션 결과)\n", "\n", " # 헬멧과 헤드 박스가 겹치면 -> 최종 라벨을 head로 덮어쓰기 (예시)\n", " # merged = []\n", " # merged.extend(final_helmet)\n", " # merged.extend(final_head)\n", " # for hm in merged:\n", " # if hm['label'] == 'helmet':\n", " # for hd in merged:\n", " # if hd['label'] == 'head':\n", " # iou_val = iou_calc(hm['bbox'], hd['bbox'])\n", " # if iou_val > 0.5:\n", " # hm['label'] = 'head'\n", " # break\n", "\n", " # (6) 최종 시각화\n", " all_final = []\n", " # all_final.extend(final_ce) # 중장비\n", " all_final.extend(final_worker) # Worker\n", " all_final.extend(final_helmet) # helmet\n", " all_final.extend(final_head) # head\n", "\n", " vis_frame = draw_detections(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", " 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" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# %% [셀 10]: 실행부 (영상 폴더 내 모든 파일 처리)\n", "video_inference_on = True\n", "video_save_on = True\n", "resize_factor_video = 1.0\n", "\n", "det_models = [\n", " det_model_whh, \n", " # det_model_ce, \n", " # det_model_fall\n", " ]\n", "\n", "scores_inits = [\n", " class_score_thresholds_whh_init,\n", " # class_score_thresholds_ce_init,\n", " # class_score_thresholds_fall_init\n", "]\n", "names_lists = [\n", " class_names_whh,\n", " # class_names_ce,\n", " # class_names_fall\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", " models=det_models,\n", " score_thresh_inits=scores_inits,\n", " class_names_lists=names_lists,\n", " # iou_threshold_for_filtering=iou_threshold_for_filtering,\n", " # move_threshold=movement_threshold,\n", " # stable_time_con_equip=stable_time_con_equip,\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", " # still_move_thr=still_move_thrshold, # <= 10픽셀이면 거의 안 움직인다고 보기\n", " # stable_time_no_movement=stable_time_no_move # 2초간 안 움직이면 필터링\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 }