Spaces:
Sleeping
Sleeping
File size: 4,976 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 | """影片 → keypoint cache(整條 pipeline 唯一需要 GPU 的步驟)。
批次模式 skip-existing:Colab 斷線重跑不重工。
每支影片各自 reset tracker,track id 不跨影片汙染。
"""
from __future__ import annotations
import hashlib
import subprocess
from pathlib import Path
from typing import Callable
import numpy as np
import pandas as pd
from ..config import Config
from ..io.cache import CACHE_COLUMNS, SCHEMA_VERSION, CacheMeta, write_cache
from ..io.video import iter_frames, probe
from .pose_tracker import PoseTracker
def _quick_sha1(path: str | Path, n_bytes: int = 1 << 20) -> str:
"""檔案前 1MB 的 sha1:足以偵測「換了影片但沒換 cache」的漂移,又不用讀全檔。"""
h = hashlib.sha1()
with open(path, "rb") as f:
h.update(f.read(n_bytes))
return h.hexdigest()
def _git_commit() -> str:
try:
return (
subprocess.run(
["git", "rev-parse", "--short", "HEAD"],
capture_output=True,
text=True,
check=True,
cwd=Path(__file__).resolve().parent,
).stdout.strip()
)
except Exception: # noqa: BLE001 - 沒 git(如 pip 安裝)不影響功能
return ""
def extract_video(
video_path: str | Path,
out_path: str | Path,
cfg: Config,
device: str | None = None,
tracker: PoseTracker | None = None,
progress_every: int = 300,
on_frame: Callable[[int, int], None] | None = None,
) -> CacheMeta:
"""單支影片 → cache parquet。可傳入共用 tracker(批次時避免重複載模型)。
``on_frame(frame_idx, total_frames)`` 選填,每幀呼叫一次,供呼叫端
(如 Gradio demo 的 ``gr.Progress``)回報逐幀進度;不影響既有的
``progress_every`` 主控台列印。
"""
video_path = Path(video_path)
info = probe(video_path)
if tracker is None:
tracker = PoseTracker(
model_name=cfg.model.name,
tracker_yaml=cfg.model.tracker,
conf=cfg.model.conf,
iou=cfg.model.iou,
device=device,
)
tracker.reset()
rows = []
n_frames = 0
for frame_idx, frame in iter_frames(video_path):
n_frames = frame_idx + 1
det = tracker.track_frame(frame, frame_idx)
t_ms = frame_idx / info.fps * 1000.0
for i in range(det.n):
rows.append(
{
"frame_idx": np.int32(frame_idx),
"t_ms": t_ms,
"track_id": np.int32(det.track_ids[i]),
"bbox_x1": det.boxes[i, 0],
"bbox_y1": det.boxes[i, 1],
"bbox_x2": det.boxes[i, 2],
"bbox_y2": det.boxes[i, 3],
"bbox_conf": det.box_conf[i],
"kpts_xy": det.kpts_xy[i].reshape(-1),
"kpts_conf": det.kpts_conf[i],
}
)
if on_frame is not None:
on_frame(frame_idx, info.n_frames)
if progress_every and frame_idx % progress_every == 0 and frame_idx > 0:
print(f" {video_path.name}: {frame_idx} 幀…")
df = pd.DataFrame(rows, columns=CACHE_COLUMNS)
meta = CacheMeta(
schema_version=SCHEMA_VERSION,
video_path=str(video_path),
video_sha1=_quick_sha1(video_path),
fps=info.fps,
width=info.width,
height=info.height,
n_frames=n_frames,
model_name=cfg.model.name,
ultralytics_version=PoseTracker.ultralytics_version(),
tracker_yaml=cfg.model.tracker,
conf=cfg.model.conf,
iou=cfg.model.iou,
device=str(device),
git_commit=_git_commit(),
)
write_cache(df, meta, out_path)
return meta
def extract_batch(
videos: list[str | Path],
out_dir: str | Path,
cfg: Config,
device: str | None = None,
skip_existing: bool = True,
) -> list[Path]:
"""批次抽取:輸出 {out_dir}/{影片檔名}.parquet;已存在即跳過(idempotent)。"""
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
tracker = PoseTracker(
model_name=cfg.model.name,
tracker_yaml=cfg.model.tracker,
conf=cfg.model.conf,
iou=cfg.model.iou,
device=device,
)
outputs = []
for i, video in enumerate(videos):
video = Path(video)
out_path = out_dir / f"{video.stem}.parquet"
outputs.append(out_path)
if skip_existing and out_path.exists():
print(f"[{i + 1}/{len(videos)}] {video.stem}: skip(已存在)")
continue
print(f"[{i + 1}/{len(videos)}] {video.stem}: 抽取中…")
extract_video(video, out_path, cfg, device=device, tracker=tracker)
return outputs
|