main_tinyYolo.py

import os
import cv2
import time
import torch
import argparse
import numpy as np

from Detection.Utils import ResizePadding
from Modules.Video.CameraLoader import CamLoader, CamLoader_Q
from Modules.Detect.DetectorLoader import TinyYOLOv3_onecls

from Modules.Pose.PoseEstimateLoader import SPPE_FastPose
from Modules.Visualize.fn import draw_single

from Track.Tracker import Detection, Tracker
from Modules.Action.ActionsEstLoader import TSSTG
from Modules.Video.s3_utils import upload_video, send_url
from datetime import datetime
from Modules.Video.merge_utils import merge_videos


#source = '../Data/test_video/test7.mp4'
#source = '../Data/falldata/Home/Videos/video (2).avi'  # hard detect
source = '../Data/falldata/Home/Videos/video (1).avi'
#source = 2


def preproc(image):
    """preprocess function for CameraLoader.
    """
    image = resize_fn(image)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image


def kpt2bbox(kpt, ex=20):
    """Get bbox that hold on all of the keypoints (x,y)
    kpt: array of shape `(N, 2)`,
    ex: (int) expand bounding box,
    """
    return np.array((kpt[:, 0].min() - ex, kpt[:, 1].min() - ex,
                     kpt[:, 0].max() + ex, kpt[:, 1].max() + ex))


if __name__ == '__main__':
    par = argparse.ArgumentParser(description='Human Fall Detection Demo.')
    par.add_argument('-C', '--camera', default=source,  # required=True,  # default=2,
                        help='Source of camera or video file path.')
    par.add_argument('--detection_input_size', type=int, default=384,
                        help='Size of input in detection model in square must be divisible by 32 (int).')
    par.add_argument('--pose_input_size', type=str, default='224x160',
                        help='Size of input in pose model must be divisible by 32 (h, w)')
    par.add_argument('--pose_backbone', type=str, default='resnet50',
                        help='Backbone model for SPPE FastPose model.')
    par.add_argument('--show_detected', default=False, action='store_true',
                        help='Show all bounding box from detection.')
    par.add_argument('--show_skeleton', default=True, action='store_true',
                        help='Show skeleton pose.')
    par.add_argument('--save_out', type=str, default='',
                        help='Save display to video file.')
    par.add_argument('--device', type=str, default='cuda',
                        help='Device to run model on cpu or cuda.')
    par.add_argument('--phone_number', type=str,
                        help='대쉬 없이 유저 전화번호 입력')
    args = par.parse_args()

    device = args.device

    fall_detected = False
    pre_fall_detected = False

    pre_detected1 = False
    pre_detected2 = False

    pre_detected_time = "null"
    detected_time = "null"

    # DETECTION MODEL.
    inp_dets = args.detection_input_size
    detect_model = TinyYOLOv3_onecls(inp_dets, device=device)

    # POSE MODEL.
    inp_pose = tuple(map(int, args.pose_input_size.split('x')))
    pose_model = SPPE_FastPose(args.pose_backbone, inp_pose[0], inp_pose[1], device=device)

    # Tracker.
    max_age = 5
    tracker = Tracker(max_age=max_age, n_init=3)

    # Actions Estimate.
    action_model = TSSTG()

    resize_fn = ResizePadding(inp_dets, inp_dets)

    cam_source = args.camera
    if type(cam_source) is str and os.path.isfile(cam_source):
        # Use loader thread with Q for video file.
        cam = CamLoader_Q(cam_source, queue_size=1000, preprocess=preproc).start()
    else:
        # Use normal thread loader for webcam.
        cam = CamLoader(int(cam_source) if cam_source.isdigit() else cam_source,
                        preprocess=preproc).start()

    #frame_size = cam.frame_size
    #scf = torch.min(inp_size / torch.FloatTensor([frame_size]), 1)[0]

    outvid = bool(args.save_out)
    if outvid:
        codec = cv2.VideoWriter_fourcc(*'mp4v')
        writer = cv2.VideoWriter(args.save_out, codec, 30, (inp_dets * 2, inp_dets * 2))


    fps_time = 0
    f = 0

    fps = 30  # 저장할 프레임 속도 (원 영상 fps와 동일하게 맞추면 좋음)
    clip_duration = 15  # 초 단위로 저장
    clip_index = 1
    recent_clips = []
    clip_start_time = time.time()
    height, width = inp_dets * 2, inp_dets * 2
    fourcc = cv2.VideoWriter_fourcc(*'avc1')
    os.makedirs('OUTPUT', exist_ok=True)

    current_clip_filename = os.path.join('OUTPUT', f'output_{clip_index:03d}.mp4')
    video_clip_writer = cv2.VideoWriter(current_clip_filename, fourcc, fps, (width, height))


    while cam.grabbed():
        f += 1
        frame = cam.getitem()
        image = frame.copy()

        # Detect humans bbox in the frame with detector model.
        detected = detect_model.detect(frame, need_resize=False, expand_bb=10)

        # Predict each tracks bbox of current frame from previous frames information with Kalman filter.
        tracker.predict()
        # Merge two source of predicted bbox together.
        for track in tracker.tracks:
            det = torch.tensor([track.to_tlbr().tolist() + [0.5, 1.0, 0.0]], dtype=torch.float32)
            detected = torch.cat([detected, det], dim=0) if detected is not None else det

        detections = []  # List of Detections object for tracking.
        if detected is not None:
            #detected = non_max_suppression(detected[None, :], 0.45, 0.2)[0]
            # Predict skeleton pose of each bboxs.
            poses = pose_model.predict(frame, detected[:, 0:4], detected[:, 4])

            # Create Detections object.
            detections = [Detection(kpt2bbox(ps['keypoints'].numpy()), # tlbr
                                    np.concatenate((ps['keypoints'].numpy(),ps['kp_score'].numpy()), axis=1), # keypoints
                                    ps['kp_score'].mean().numpy()) for ps in poses] # confidence

            # VISUALIZE.
            if args.show_detected:
                for bb in detected[:, 0:5]:
                    frame = cv2.rectangle(frame, (bb[0], bb[1]), (bb[2], bb[3]), (0, 0, 255), 1)

        # Update tracks by matching each track information of current and previous frame or
        # create a new track if no matched.
        tracker.update_one(detections)
        action_probs_text = []

        # Predict Actions of each track.
        for i, track in enumerate(tracker.tracks):
            if not track.is_confirmed():
                continue

            track_id = track.track_id
            bbox = track.to_tlbr().astype(int)
            center = track.get_center().astype(int)

            action = 'pending..'
            clr = (0, 255, 0)

            # Use 30 frames time-steps to prediction.
            if len(track.keypoints_list) == 30:
                pts = np.array(track.keypoints_list, dtype=np.float32)
                out = action_model.predict(pts, frame.shape[:2])
                action_idx = out[0].argmax()
                action_name = action_model.class_names[out[0].argmax()]
                action_confidence = out[0][action_idx]

                action = '{}: {:.2f}%'.format(action_name, out[0].max() * 100)
                # 클래스별 확률 저장
                for idx, class_name in enumerate(action_model.class_names):
                    action_probs_text.append(f"{class_name}: {out[0][idx] * 100:.1f}%")

                if action_name == 'Fall Down':
                    clr = (255, 0, 0)
                    fall_detected = True
                    pre_detected1 = True
                    detected_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

                elif action_name == 'Lying Down':
                    clr = (255, 200, 0)

            # VISUALIZE.
            if track.time_since_update == 0:
                if args.show_skeleton:
                    frame = draw_single(frame, track.keypoints_list[-1])
                frame = cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 1)
                frame = cv2.putText(frame, str(track_id), (center[0], center[1]), cv2.FONT_HERSHEY_SIMPLEX,
                                    0.4, (255, 0, 0), 2)
                frame = cv2.putText(frame, action, (bbox[0] + 5, bbox[1] + 15), cv2.FONT_HERSHEY_SIMPLEX,
                                    0.4, clr, 1)
        # Show Frame.
        frame = cv2.resize(frame, (0, 0), fx=2., fy=2.)
        frame = cv2.putText(frame, '%d, FPS: %f' % (f, 1.0 / (time.time() - fps_time)),
                            (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
        for i, prob_text in enumerate(action_probs_text):
            frame = cv2.putText(frame, '%s' % prob_text, (10, 35 + i * 12),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 255), 1)

        frame = frame[:, :, ::-1]
        fps_time = time.time()

        if outvid:
            writer.write(frame)

        now = time.time()
        if args.phone_number is not None:
            if now - clip_start_time >= clip_duration:
                print(f"[INFO] 클립 종료 - 인덱스: {clip_index}")
                print(f"[INFO] 최근 클립 리스트: {recent_clips}")
                print(f"[INFO] pre_detected1: {pre_detected1}, pre_detected2: {pre_detected2}")
                print(f"[INFO] fall_detected: {fall_detected}, detected_time: {detected_time}")
                video_clip_writer.release()

                recent_clips.append(current_clip_filename)
                if len(recent_clips) > 3:
                    os.remove(recent_clips.pop(0))  # 오래된 영상 삭제

                if pre_detected2:
                    print("[Fall] 이전에 T였음. 영상 병합 시작...")
                    merged_path = os.path.join('OUTPUT', f'merged_{clip_index:03d}.mp4')
                    merge_videos(recent_clips[-3:], merged_path)

                    try:
                        s3_url = upload_video(merged_path)
                        print(f"[S3] 병합 영상 업로드 완료: {s3_url}")
                        send_url(s3_url,args.phone_number, pre_fall_detected, pre_detected_time)
                    except Exception as e:
                        print(f"[Error] 병합 영상 업로드 실패: {e}")
                    else:
                        print("[Info] 감지되지 않음 → 병합 X")

                clip_index += 1
                pre_detected2 = pre_detected1
                pre_detected1 = False
                pre_fall_detected = fall_detected
                pre_detected_time = detected_time
                fall_detected = False
                detected_time = "null"

                current_clip_filename = os.path.join('OUTPUT', f'output_{clip_index:03d}.mp4')
                video_clip_writer = cv2.VideoWriter(current_clip_filename, fourcc, fps, (width, height))

                clip_start_time = now

        video_clip_writer.write(frame)

        cv2.imshow('frame', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # Clear resource.
    cam.stop()
    if outvid:
        writer.release()
    video_clip_writer.release()
    if args.phone_number is not None:
        try:
            s3_url = upload_video(current_clip_filename)
            print(f"[S3] 마지막 클립 업로드 완료: {s3_url}")
            send_url(s3_url, args.phone_number, fall_detected, detected_time)
        except Exception as e:
            print(f"[Error] 마지막 클립 S3 업로드 실패:", e)
    cv2.destroyAllWindows()