File size: 10,932 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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
"""``fdp`` 命令列入口。



子命令(依 pipeline 順序):

  extract   影片(或目錄)→ keypoint cache parquet【唯一 GPU 步驟】

  detect    cache → events.json(純 CPU,毫秒級;--debug 另輸出特徵 JSONL)

  annotate  影片 + cache → 標註影片(H.264)+ events.json

  pipeline  extract → detect → annotate 一條龍(--source 0 可用 webcam 錄一段再處理)

  bench     影片 → FPS benchmark(可攜:任何機器都能補跑一列,不綁 Colab)



重依賴(torch/ultralytics/cv2)一律延遲到子命令內 import:

``fdp detect`` 在只裝核心依賴的環境(無 GPU)也能跑。

"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv"}


def _write_debug_jsonl(path: str | Path, records: list[dict]) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        for rec in records:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")


def _run_rules(cache_path: str, config_path: str, collect_debug: bool):
    from .config import load_config
    from .io.cache import read_cache
    from .rules import run_engine

    cfg = load_config(config_path)
    df, meta = read_cache(cache_path)
    events, debug = run_engine(df, meta.fps, cfg, collect_debug=collect_debug)
    return cfg, df, meta, events, debug


def cmd_extract(args: argparse.Namespace) -> int:
    """影片 → keypoint cache。--source 可為單一影片或裝滿影片的目錄。"""
    from .config import load_config
    from .inference.extract import extract_batch, extract_video

    cfg = load_config(args.config)
    if args.model:
        cfg.model.name = args.model
    src = Path(args.source)
    if src.is_dir():
        videos = sorted(p for p in src.iterdir() if p.suffix.lower() in VIDEO_EXTS)
        if not videos:
            print(f"{src} 內沒有影片", file=sys.stderr)
            return 1
        extract_batch(videos, args.out, cfg, device=args.device)
    else:
        out = Path(args.out)
        out_path = out / f"{src.stem}.parquet" if out.suffix == "" else out
        meta = extract_video(src, out_path, cfg, device=args.device)
        print(f"cache → {out_path}({meta.n_frames} 幀,model={meta.model_name})")
    return 0


def cmd_detect(args: argparse.Namespace) -> int:
    """cache → events.json;--debug 另存 per-frame 特徵 JSONL(失敗分析用)。"""
    from .events.schema import write_events_json

    cfg, _, meta, events, debug = _run_rules(args.cache, args.config, args.debug is not None)
    write_events_json(args.out, events, source=meta.video_path, fps=meta.fps)
    if args.debug:
        _write_debug_jsonl(args.debug, debug)
    print(f"{Path(args.cache).stem}: {len(events)} 個事件 → {args.out}")
    for ev in events:
        print(
            f"  tracks={ev.track_ids} frames=[{ev.start_frame},{ev.end_frame}] "
            f"t=[{ev.start_time_s:.2f},{ev.end_time_s:.2f}]s rules={ev.rules_fired}"
        )
    return 0


def cmd_annotate(args: argparse.Namespace) -> int:
    """影片 + cache → 標註影片(H.264);同時輸出 events.json。"""
    from .events.schema import write_events_json
    from .viz.annotate import annotate_video

    cfg, df, meta, events, debug = _run_rules(args.cache, args.config, True)
    out = annotate_video(args.video, df, meta.fps, cfg, events, debug, args.out)
    if args.events_out:
        write_events_json(args.events_out, events, source=str(args.video), fps=meta.fps)
    print(f"標註影片 → {out}({len(events)} 個事件)")
    return 0


def _record_webcam(index: int, duration_s: float, out_path: Path) -> Path:
    """webcam 錄一段到暫存影片(離線 pipeline 的 webcam 模式)。"""
    import cv2

    cap = cv2.VideoCapture(index)
    if not cap.isOpened():
        raise RuntimeError(f"無法開啟 webcam {index}")
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    out_path.parent.mkdir(parents=True, exist_ok=True)
    writer = cv2.VideoWriter(str(out_path), cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
    n = int(duration_s * fps)
    print(f"webcam 錄製 {duration_s:.0f}s({n} 幀)…")
    for _ in range(n):
        ok, frame = cap.read()
        if not ok:
            break
        writer.write(frame)
    cap.release()
    writer.release()
    return out_path


def cmd_pipeline(args: argparse.Namespace) -> int:
    """extract → detect → annotate 一條龍,輸出標註影片與 events.json。"""
    import tempfile

    from .config import load_config
    from .events.schema import write_events_json
    from .inference.extract import extract_video
    from .rules import run_engine
    from .viz.annotate import annotate_video

    cfg = load_config(args.config)
    if args.model:
        cfg.model.name = args.model

    if str(args.source).isdigit():
        video = _record_webcam(
            int(args.source), args.duration, Path(tempfile.mkstemp(suffix=".mp4")[1])
        )
    else:
        video = Path(args.source)

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    cache_path = out_dir / f"{video.stem}.parquet"

    extract_video(video, cache_path, cfg, device=args.device)
    from .io.cache import read_cache

    df, meta = read_cache(cache_path)
    events, debug = run_engine(df, meta.fps, cfg, collect_debug=True)
    events_path = out_dir / f"{video.stem}.events.json"
    write_events_json(events_path, events, source=str(video), fps=meta.fps)
    if args.debug:
        _write_debug_jsonl(out_dir / f"{video.stem}.debug.jsonl", debug)
    annotated = annotate_video(
        video, df, meta.fps, cfg, events, debug, out_dir / f"{video.stem}_annotated.mp4"
    )
    print(f"完成:{annotated}{events_path}({len(events)} 個事件)")
    return 0


def cmd_bench(args: argparse.Namespace) -> int:
    """影片 → FPS benchmark(純推論 + 端到端 FPS、p50/p95 延遲)。



    可攜:任何機器都能對同一支(或任一支)影片補跑一列,不綁定 Colab——

    --model 可重複指定多次,一次跑完整個模型清單。

    """
    import platform

    from .bench.benchmark import benchmark, load_frames

    frames = load_frames(args.video, n_frames=args.n_frames)
    print(f"{args.video}: 載入 {len(frames)} 幀(要求 {args.n_frames})")

    results = []
    for model_name in args.model:
        r = benchmark(
            frames,
            model_name=model_name,
            device=args.device,
            quantize=args.quantize,
            n_runs=args.n_runs,
            warmup=args.warmup,
        )
        results.append(r.to_dict())
        print(
            f"  {model_name}: 純推論 {r.pure_inference_fps} FPS、端到端 {r.end_to_end_fps} FPS、"
            f"p50 {r.p50_latency_ms}ms、p95 {r.p95_latency_ms}ms"
        )

    if args.out:
        import torch
        import ultralytics

        payload = {
            "video": str(args.video),
            "platform": platform.platform(),
            "torch_version": torch.__version__,
            "ultralytics_version": ultralytics.__version__,
            "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
            "results": results,
        }
        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
        Path(args.out).write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"→ {args.out}")
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="fdp", description=__doc__)
    sub = parser.add_subparsers(dest="command", required=True)

    p = sub.add_parser("extract", help="影片 → keypoint cache(GPU)")
    p.add_argument("--source", required=True, help="影片檔或影片目錄")
    p.add_argument("--out", required=True, help="輸出 parquet 或目錄")
    p.add_argument("--config", default="config.yaml")
    p.add_argument("--model", default=None, help="覆寫 config 的模型名(如 yolo26s-pose.pt)")
    p.add_argument("--device", default=None, help="cuda:0 / cpu(預設由 ultralytics 自選)")
    p.set_defaults(func=cmd_extract)

    p = sub.add_parser("detect", help="cache → events.json(純 CPU)")
    p.add_argument("--cache", required=True)
    p.add_argument("--out", required=True, help="events.json 輸出路徑")
    p.add_argument("--config", default="config.yaml")
    p.add_argument("--debug", default=None, help="per-frame 特徵 JSONL 輸出路徑")
    p.set_defaults(func=cmd_detect)

    p = sub.add_parser("annotate", help="影片 + cache → 標註影片(H.264)")
    p.add_argument("--video", required=True)
    p.add_argument("--cache", required=True)
    p.add_argument("--out", required=True, help="標註影片輸出路徑")
    p.add_argument("--events-out", default=None, help="events.json 輸出路徑(可選)")
    p.add_argument("--config", default="config.yaml")
    p.set_defaults(func=cmd_annotate)

    p = sub.add_parser("pipeline", help="extract → detect → annotate 一條龍")
    p.add_argument("--source", required=True, help="影片檔,或 webcam 索引(如 0)")
    p.add_argument("--out-dir", default="outputs")
    p.add_argument("--config", default="config.yaml")
    p.add_argument("--model", default=None)
    p.add_argument("--device", default=None)
    p.add_argument("--duration", type=float, default=10.0, help="webcam 錄製秒數")
    p.add_argument("--debug", action="store_true", help="輸出 per-frame 特徵 JSONL")
    p.set_defaults(func=cmd_pipeline)

    p = sub.add_parser("bench", help="影片 → FPS benchmark(可攜,任何機器都能補跑)")
    p.add_argument("--video", required=True, help="固定用來計時的影片")
    p.add_argument("--model", action="append", required=True, help="模型名,可重複指定多次")
    p.add_argument("--device", default=None, help="cuda:0 / cpu(預設由 ultralytics 自選)")
    p.add_argument("--quantize", default=None, help="16 或 fp16 啟用 FP16 推論(僅 GPU 有意義)")
    p.add_argument("--n-frames", type=int, default=300, dest="n_frames")
    p.add_argument("--n-runs", type=int, default=3, dest="n_runs")
    p.add_argument("--warmup", type=int, default=20)
    p.add_argument("--out", default=None, help="bench.json 輸出路徑(可選)")
    p.set_defaults(func=cmd_bench)

    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())