Spaces:
Running
Running
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())
|