Hugging Face's logo

Spaces:
betty0
/
fall-detection-pose
Sleeping

App Files Community
fall-detection-pose
File size: 5,014 Bytes
11fab85
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
"""標註影片輸出:骨架 + track id + 狀態標籤 + ALARM 警示。



狀態來源是規則引擎的 per-frame debug 紀錄(單一事實來源:

畫面上看到的狀態就是引擎當下判定的狀態,而非事後由事件區間反推),

ALARM 橫幅出現的時刻即真實告警延遲,demo 不美化。

"""

from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np
import pandas as pd

from ..config import Config
from ..events.schema import FallEvent
from ..io.cache import N_KPTS
from ..io.video import H264VideoWriter, iter_frames, probe

# COCO 17 骨架連線(肢段)
SKELETON = [
    (5, 7), (7, 9),      # 左臂
    (6, 8), (8, 10),     # 右臂
    (5, 6),              # 肩線
    (5, 11), (6, 12),    # 軀幹
    (11, 12),            # 髖線
    (11, 13), (13, 15),  # 左腿
    (12, 14), (14, 16),  # 右腿
    (0, 5), (0, 6),      # 頭-肩
]

STATE_COLORS = {  # BGR
    "UPRIGHT": (80, 200, 80),
    "FALLING": (0, 165, 255),
    "FALLEN": (0, 60, 255),
    "ALARM": (0, 0, 255),
}


def _draw_person(

    frame: np.ndarray,

    kpts_xy: np.ndarray,

    kpts_conf: np.ndarray,

    bbox: tuple,

    track_id: int,

    state: str,

    kpt_conf_min: float,

) -> None:
    color = STATE_COLORS.get(state, (200, 200, 200))
    x1, y1, x2, y2 = (int(v) for v in bbox)
    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

    pts = kpts_xy.reshape(N_KPTS, 2)
    visible = kpts_conf >= kpt_conf_min
    for a, b in SKELETON:
        if visible[a] and visible[b]:
            pa = tuple(int(v) for v in pts[a])
            pb = tuple(int(v) for v in pts[b])
            cv2.line(frame, pa, pb, color, 2)
    for i in range(N_KPTS):
        if visible[i]:
            cv2.circle(frame, tuple(int(v) for v in pts[i]), 3, color, -1)

    label = f"id {track_id}  {state}"
    (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 2)
    ty = max(y1 - 8, th + 4)
    cv2.rectangle(frame, (x1, ty - th - 6), (x1 + tw + 6, ty + 4), color, -1)
    cv2.putText(
        frame, label, (x1 + 3, ty), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 2
    )


def _draw_alarm_banner(frame: np.ndarray, track_ids: list[int]) -> None:
    h, w = frame.shape[:2]
    overlay = frame.copy()
    cv2.rectangle(overlay, (0, 0), (w, 46), (0, 0, 220), -1)
    cv2.addWeighted(overlay, 0.75, frame, 0.25, 0, frame)
    ids = ",".join(str(t) for t in sorted(set(track_ids)))
    cv2.putText(
        frame,
        f"FALL ALARM  (track {ids})",
        (12, 32),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.9,
        (255, 255, 255),
        2,
    )


def annotate_video(

    video_path: str | Path,

    cache_df: pd.DataFrame,

    fps: float,

    cfg: Config,

    events: list[FallEvent],

    debug_records: list[dict],

    out_path: str | Path,

) -> Path:
    """輸出標註影片(H.264)。



    Args:

        cache_df: keypoint cache rows(骨架繪製來源)。

        events: 引擎輸出的事件(用於畫面右上角事件計數)。

        debug_records: ``run_engine(collect_debug=True)`` 的 per-frame 狀態。

    """
    out_path = Path(out_path)
    info = probe(video_path)

    state_by_frame_track: dict[tuple[int, int], str] = {
        (d["frame_idx"], d["track_id"]): d["state"] for d in debug_records
    }
    rows_by_frame: dict[int, list] = {}
    for row in cache_df.itertuples(index=False):
        rows_by_frame.setdefault(int(row.frame_idx), []).append(row)

    with H264VideoWriter(out_path, info.fps, info.width, info.height) as writer:
        for frame_idx, frame in iter_frames(video_path):
            t_s = frame_idx / fps
            alarm_tracks: list[int] = []
            for row in rows_by_frame.get(frame_idx, []):
                tid = int(row.track_id)
                state = state_by_frame_track.get((frame_idx, tid), "UPRIGHT")
                if tid < 0:
                    state = "-"  # 未指派 track 的偵測:只畫框不標狀態
                _draw_person(
                    frame,
                    np.asarray(row.kpts_xy),
                    np.asarray(row.kpts_conf),
                    (row.bbox_x1, row.bbox_y1, row.bbox_x2, row.bbox_y2),
                    tid,
                    state,
                    cfg.model.kpt_conf_min,
                )
                if state == "ALARM":
                    alarm_tracks.append(tid)
            if alarm_tracks:
                _draw_alarm_banner(frame, alarm_tracks)
            n_done = sum(1 for e in events if e.end_time_s <= t_s)
            cv2.putText(
                frame,
                f"frame {frame_idx}  t={t_s:6.2f}s  events={n_done}/{len(events)}",
                (10, info.height - 12),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (255, 255, 255),
                1,
            )
            writer.write(frame)
    return out_path