| |
| from __future__ import annotations |
|
|
| import argparse |
| from collections import deque |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| from fall_detection.features import normalize_pose_sequence |
| from fall_detection.models import PoseGRU, PoseTCN |
| from fall_detection.pose_extractor import MediaPipePoseExtractor |
|
|
| CONNECTIONS = [ |
| (11, 12), (11, 13), (13, 15), (12, 14), (14, 16), |
| (11, 23), (12, 24), (23, 24), (23, 25), (25, 27), |
| (24, 26), (26, 28), (27, 31), (28, 32), |
| ] |
|
|
|
|
| def draw_pose(frame: np.ndarray, pose: np.ndarray | None) -> None: |
| if pose is None: |
| return |
| height, width = frame.shape[:2] |
| for first, second in CONNECTIONS: |
| if min(pose[first, 3], pose[second, 3]) < 0.35: |
| continue |
| p1 = (int(pose[first, 0] * width), int(pose[first, 1] * height)) |
| p2 = (int(pose[second, 0] * width), int(pose[second, 1] * height)) |
| cv2.line(frame, p1, p2, (70, 230, 70), 2) |
| for point in pose: |
| if point[3] >= 0.35: |
| cv2.circle(frame, (int(point[0] * width), int(point[1] * height)), 2, (0, 255, 255), -1) |
|
|
|
|
| def load_model(path: Path, device: torch.device) -> tuple[nn.Module, dict]: |
| checkpoint = torch.load(path, map_location=device, weights_only=False) |
| if checkpoint.get("architecture", "gru") == "tcn": |
| model = PoseTCN( |
| input_size=checkpoint["input_size"], |
| channels=checkpoint["tcn_channels"], |
| dropout=checkpoint["dropout"], |
| ).to(device) |
| else: |
| model = PoseGRU( |
| input_size=checkpoint["input_size"], |
| hidden_size=checkpoint["hidden_size"], |
| num_layers=checkpoint["num_layers"], |
| dropout=checkpoint["dropout"], |
| ).to(device) |
| model.load_state_dict(checkpoint["state_dict"]) |
| model.eval() |
| return model, checkpoint |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Run fall detection on a video") |
| parser.add_argument("--input", type=Path, required=True) |
| parser.add_argument("--model", type=Path, default=Path("artifacts/experiments/urfd/pose_tcn/model.pt")) |
| parser.add_argument("--output", type=Path) |
| parser.add_argument("--target-fps", type=float, default=10.0) |
| parser.add_argument("--smoothing", type=int, default=5) |
| parser.add_argument("--alert-consecutive", type=int, default=3) |
| args = parser.parse_args() |
| output = args.output or args.input.with_name(args.input.stem + "_detected.mp4") |
| output.parent.mkdir(parents=True, exist_ok=True) |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model, checkpoint = load_model(args.model, device) |
| sequence_length = int(checkpoint["sequence_length"]) |
| feature_mean = np.asarray(checkpoint["feature_mean"], dtype=np.float32) |
| feature_std = np.asarray(checkpoint["feature_std"], dtype=np.float32) |
| threshold = float(checkpoint["threshold"]) |
|
|
| capture = cv2.VideoCapture(str(args.input)) |
| if not capture.isOpened(): |
| raise SystemExit(f"Cannot open {args.input}") |
| source_fps = capture.get(cv2.CAP_PROP_FPS) or 25.0 |
| width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| writer = cv2.VideoWriter( |
| str(output), cv2.VideoWriter_fourcc(*"mp4v"), source_fps, (width, height) |
| ) |
| sample_every = max(1, int(round(source_fps / args.target_fps))) |
| pose_window: deque[np.ndarray] = deque(maxlen=sequence_length) |
| probability_window: deque[float] = deque(maxlen=args.smoothing) |
| probability = 0.0 |
| consecutive = 0 |
| frame_index = 0 |
|
|
| with MediaPipePoseExtractor() as extractor: |
| while True: |
| ok, frame = capture.read() |
| if not ok: |
| break |
| current_pose = None |
| display_pose = None |
| if frame_index % sample_every == 0: |
| is_wide_urfd = frame.shape[1] / frame.shape[0] > 2.2 |
| pose_frame = frame[:, frame.shape[1] // 2 :] if is_wide_urfd else frame |
| current_pose = extractor.process_rgb(cv2.cvtColor(pose_frame, cv2.COLOR_BGR2RGB)) |
| if current_pose is None: |
| current_pose = pose_window[-1].copy() if pose_window else np.zeros((33, 4), dtype=np.float32) |
| pose_window.append(current_pose) |
| display_pose = current_pose.copy() |
| if is_wide_urfd: |
| display_pose[:, 0] = 0.5 + 0.5 * display_pose[:, 0] |
| if len(pose_window) >= sequence_length: |
| features = normalize_pose_sequence( |
| np.stack(pose_window), checkpoint["visibility_threshold"] |
| ) |
| feature_indices = np.asarray( |
| checkpoint.get("feature_indices", np.arange(features.shape[1])), |
| dtype=np.int64, |
| ) |
| features = features[:, feature_indices] |
| features = (features - feature_mean) / feature_std |
| tensor = torch.from_numpy(features[None]).float().to(device) |
| with torch.no_grad(): |
| probability = float(torch.sigmoid(model(tensor)).item()) |
| probability_window.append(probability) |
| probability = float(np.mean(probability_window)) |
| consecutive = consecutive + 1 if probability >= threshold else 0 |
|
|
| draw_pose(frame, display_pose) |
| alert = consecutive >= args.alert_consecutive |
| color = (0, 0, 255) if alert else (0, 180, 0) |
| status = "CANH BAO: TE NGA" if alert else "BINH THUONG" |
| cv2.rectangle(frame, (10, 10), (385, 78), (0, 0, 0), -1) |
| cv2.putText(frame, status, (20, 38), cv2.FONT_HERSHEY_SIMPLEX, 0.72, color, 2) |
| cv2.putText( |
| frame, |
| f"P(fall)={probability:.2f} threshold={threshold:.2f}", |
| (20, 66), |
| cv2.FONT_HERSHEY_SIMPLEX, |
| 0.50, |
| (255, 255, 255), |
| 1, |
| ) |
| writer.write(frame) |
| frame_index += 1 |
|
|
| capture.release() |
| writer.release() |
| print(f"Saved annotated video to {output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
