main.py

import os
import cv2
import time
import torch
import argparse
import numpy as np
import zipfile
import tempfile
from datetime import datetime

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

from PoseEstimateLoader import SPPE_FastPose
from fn import draw_single

from Track.Tracker import Detection, Tracker
from ActionsEstLoader import TSSTG

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


def process_video(video_path, args, detect_model, pose_model, action_model, tracker, resize_fn):
    """Process a single video file for fall detection.
    
    Args:
        video_path: Path to the video file
        args: Command line arguments
        detect_model: Detection model
        pose_model: Pose estimation model
        action_model: Action recognition model
        tracker: Tracker instance
        resize_fn: Resize function for preprocessing
        
    Returns:
        bool: Whether a fall was detected
    """
    fall_detected = False
    video_name = os.path.basename(str(video_path))
    
    # Initialize video capture
    if type(video_path) is str and os.path.isfile(video_path):
        cam = CamLoader_Q(video_path, queue_size=1000, preprocess=preproc).start()
    else:
        cam = CamLoader(int(video_path) if video_path.isdigit() else video_path,
                        preprocess=preproc).start()

    # Initialize video writer if needed
    outvid = False
    if args.save_out != '':
        try:
            output_path = os.path.join(args.save_out, f'processed_{video_name}')
            # Try different codecs in order of preference
            codecs = [
                ('mp4v', 'mp4'),  # MPEG-4
                ('XVID', 'avi'),  # XVID
                ('MJPG', 'avi'),  # Motion JPEG
                ('X264', 'mp4'),  # H.264
            ]
            
            writer = None
            for codec_name, ext in codecs:
                try:
                    codec = cv2.VideoWriter_fourcc(*codec_name)
                    output_path = os.path.splitext(output_path)[0] + '.' + ext
                    writer = cv2.VideoWriter(output_path, codec, 30, (args.detection_input_size * 2, args.detection_input_size * 2))
                    if writer.isOpened():
                        outvid = True
                        print(f"Using codec: {codec_name}")
                        break
                except Exception as e:
                    print(f"Failed to initialize {codec_name} codec: {str(e)}")
                    continue
            
            if not outvid:
                print("Warning: Could not initialize video writer. Processing will continue without saving output video.")
        except Exception as e:
            print(f"Error setting up video writer: {str(e)}")
            print("Processing will continue without saving output video.")

    fps_time = 0
    f = 0
    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:
            # 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()),
                                    np.concatenate((ps['keypoints'].numpy(),
                                                    ps['kp_score'].numpy()), axis=1),
                                    ps['kp_score'].mean().numpy()) for ps in poses]

            # 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(detections)

        # 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_name = action_model.class_names[out[0].argmax()]
                action = '{}: {:.2f}%'.format(action_name, out[0].max() * 100)
                if action_name == 'Fall Down':
                    clr = (255, 0, 0)
                    fall_detected = True
                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_COMPLEX,
                                    0.4, (255, 0, 0), 2)
                frame = cv2.putText(frame, action, (bbox[0] + 5, bbox[1] + 15), cv2.FONT_HERSHEY_COMPLEX,
                                    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.5, (0, 255, 0), 1)
        frame = frame[:, :, ::-1]
        fps_time = time.time()

        if outvid:
            writer.write(frame)

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

    # Clear resource.
    cam.stop()
    if outvid:
        writer.release()
    if args.show_detected or args.show_skeleton:
        cv2.destroyAllWindows()
        
    return fall_detected, video_name

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))

def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 'True', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'False', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

if __name__ == '__main__':
    par = argparse.ArgumentParser(description='Human Fall Detection Demo.')
    par.add_argument('-C', '--camera', default=source,
                        help='Source of camera or video file path.')
    par.add_argument('--zip_file', type=str, default='',
                        help='Path to zip file containing videos to process.')
    par.add_argument('--detection_input_size', type=int, default=416,
                        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 or directory for multiple videos.')
    par.add_argument('--device', type=str, default='cuda',
                        help='Device to run model on cpu or cuda.')
    par.add_argument('--use_onnx', type=str2bool, default=True,
                        help='Use ONNX model instead of PyTorch model.')
    args = par.parse_args()

    # Initialize models
    device = args.device
    use_onnx = args.use_onnx
    inp_dets = args.detection_input_size
    detect_model = TinyYOLOv3_onecls(inp_dets, device=device, use_onnx=use_onnx)

    inp_pose = args.pose_input_size.split('x')
    inp_pose = (int(inp_pose[0]), int(inp_pose[1]))
    pose_model = SPPE_FastPose(args.pose_backbone, inp_pose[0], inp_pose[1], device=device, use_onnx=use_onnx)

    # Initialize tracker
    max_age = 30
    tracker = Tracker(max_age=max_age, n_init=3)

    # Initialize action model
    action_model = TSSTG(use_onnx=use_onnx)

    resize_fn = ResizePadding(inp_dets, inp_dets)

    # Create output directory if needed
    if args.save_out and not os.path.exists(args.save_out):
        os.makedirs(args.save_out)

    # Process videos
    if args.zip_file:
        # Process videos from zip file
        with zipfile.ZipFile(args.zip_file, 'r') as zip_ref:
            with tempfile.TemporaryDirectory() as temp_dir:
                # Extract all videos to temporary directory
                zip_ref.extractall(temp_dir)
                
                # Process each video file
                for root, _, files in os.walk(temp_dir):
                    for file in files:
                        if file.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
                            video_path = os.path.join(root, file)
                            fall_detected, video_name = process_video(
                                video_path, args, detect_model, pose_model, 
                                action_model, tracker, resize_fn
                            )
                            
                            # Record results
                            with open('fall_detection_results.txt', 'a') as f:
                                time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                                result_str = f'{time_str} | {video_name}: {"Fall Detected" if fall_detected else "No Fall"}\n'
                                f.write(result_str)
    else:
        # Process single video
        fall_detected, video_name = process_video(
            args.camera, args, detect_model, pose_model,
            action_model, tracker, resize_fn
        )
        
        # Record results
        with open('fall_detection_results.txt', 'a') as f:
            time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            result_str = f'{time_str} | {video_name}: {"Fall Detected" if fall_detected else "No Fall"}\n'
            f.write(result_str)