poolcoach-analyzer / scripts /harvest_frames.py
masterdanh's picture
feat(dataset): add ND1 frame harvest script for archive VODs
3ce4659
Raw History Blame Contribute Delete
12.4 kB
"""Harvest still frames from an archive VOD into the ND1 dataset -- design
doc 7_FarZone_Detection_BallID.md section 3.5, step 3 of the BRIEF of
28/08/2026.
WHY. Every fix in that design doc is a patch that buys frames back without
retraining: more pixels through the detector, a low-confidence recovery lane
for tracks that are already alive. The real repair is a detector fine-tuned
on THIS footage -- far balls of 20 px, balls hiding behind balls in a rack
seen down its own axis -- and that needs a labelled set. This script is the
first half of it: pull the frames out, record where every one came from, and
stop. Labelling and training are later sessions.
WHAT IT GUARANTEES.
- The archive is READ ONLY. ``D:\\Khoa luan`` is opened for reading and
nothing else; the whole tree is fingerprinted (every file's size and
mtime) before and after and the two are compared, so "no byte was written
to the archive" is a measurement in the manifest, not a promise.
- Every frame is traceable. The manifest records the source path, the
SHA-256 of the source file, the nominal fps and duration, and for each
frame its file name, its index and its container PTS. A frame whose
provenance is lost is a frame that cannot be used as training data.
- The output lands OUTSIDE both repositories, in
``D:\\poolcoach\\datasets\\nd1_frames\\<vod stem>\\``: broadcast footage is
not committed, and neither are its frames.
SAMPLING. One frame every ``--every`` seconds (2 s by default), taken on the
container's own PTS, which is what makes the count come out at
duration / every within a frame. ``--shots <shots.json>`` optionally ADDS
frames around the moments a shot list already knows about -- the strike and
the settle of every shot, and the seconds before a rack is broken, which is
failure mode 2 of section 1 and the hardest thing in the set to label. The
extra frames are snapped onto the same PTS grid and de-duplicated against
the uniform ones, and the manifest counts the two kinds separately.
RUNNING IT (the launcher owns the paths):
..\\poolcoach-cv-env\\Scripts\\python.exe scripts\\harvest_frames.py ^
--video "D:\\Khoa luan\\...\\rack_5245_5805.mp4" ^
--out "D:\\poolcoach\\datasets\\nd1_frames"
"""
from __future__ import annotations
import argparse
import hashlib
import json
import sys
import time
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
for _p in (ROOT / "src", ROOT):
if str(_p) not in sys.path:
sys.path.insert(0, str(_p))
EVERY_S = 2.0 # the default sampling period. A rack lasts minutes
# and the table barely changes between two
# neighbouring frames; 1 frame / 2 s is what
# section 3.5 asks for
JPEG_QUALITY = 95 # training data, not a thumbnail: the artefacts of
# a cheap JPEG land on exactly the 20 px balls the
# set exists to teach
SHOT_OFFSETS_S = (-1.5, -0.5, 0.0, 0.3, 0.8)
# around each shot's onset when a shot list is
# given: the rack still standing, the strike, and
# the first moments of the spread -- the states a
# uniform grid hits only by luck
HASH_CHUNK = 1 << 20
def sha256_of(path: Path) -> str:
"""The source file's SHA-256 -- the only way a frame harvested today can
be tied to the exact file it came from tomorrow."""
h = hashlib.sha256()
with open(path, "rb") as fh:
while True:
chunk = fh.read(HASH_CHUNK)
if not chunk:
break
h.update(chunk)
return h.hexdigest()
def fingerprint(root: Path) -> dict:
"""Every file under ``root`` as name -> (size, mtime_ns), so "the
archive was not written to" can be CHECKED after the run instead of
assumed."""
out = {}
for p in sorted(root.rglob("*")):
try:
st = p.stat()
except OSError:
continue
if p.is_file():
out[str(p.relative_to(root))] = [st.st_size, st.st_mtime_ns]
return out
def wanted_times(duration_s: float, every: float,
shots_path: Path | None) -> tuple[list[float], set[float]]:
"""The timestamps to harvest: the uniform grid, plus the shot-relative
ones when a shot list is given.
Returns (sorted times, the subset that came from the shot list). The
uniform grid is ``0, every, 2*every, ...`` strictly below the duration,
which is what makes the count land on ``duration / every`` within one
frame -- the acceptance number of step 3.
"""
grid: list[float] = []
t = 0.0
while t < duration_s:
grid.append(round(t, 3))
t += every
extra: set[float] = set()
if shots_path is not None:
manifest = json.loads(shots_path.read_text(encoding="utf-8"))
for shot in manifest.get("shots", []):
base = shot.get("t_onset_s")
if base is None:
continue
for off in SHOT_OFFSETS_S:
t = round(float(base) + off, 3)
if 0.0 <= t < duration_s:
extra.add(t)
extra -= set(grid)
return sorted(set(grid) | extra), extra
def harvest(video: Path, out_dir: Path, every: float,
shots_path: Path | None, limit: int | None) -> dict:
"""Decode the VOD ONCE from start to end, keeping the frame nearest each
wanted timestamp.
Sequential decoding rather than seeking: a seek on a long-GOP broadcast
file lands on the previous keyframe, so the PTS written into the
manifest would not be the PTS of the image saved next to it. Decoding
320 s costs seconds and buys an honest manifest.
"""
import cv2
cap = cv2.VideoCapture(str(video))
if not cap.isOpened():
raise SystemExit(f"cannot open {video}")
try:
fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
n_est = cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0
duration = n_est / fps if fps > 0 and n_est > 0 else 0.0
if duration <= 0:
raise SystemExit(f"{video} reports no duration (fps={fps}, "
f"frames={n_est}) -- it cannot be sampled")
times, from_shots = wanted_times(duration, every, shots_path)
if limit:
times = times[:limit]
out_dir.mkdir(parents=True, exist_ok=True)
frames: list[dict] = []
idx = -1
want = 0
t0 = time.perf_counter()
while want < len(times):
if not cap.grab():
break
idx += 1
t_s = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0
if t_s + 1e-9 < times[want]:
continue
ok, frame = cap.retrieve()
if not ok:
break
name = f"{video.stem}_{idx:06d}.jpg"
cv2.imwrite(str(out_dir / name), frame,
[cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
frames.append({"file": name, "frame_idx": idx,
"t_s": round(t_s, 3),
"t_wanted_s": times[want],
"source": "shots" if times[want] in from_shots
else "grid"})
want += 1
if len(frames) % 25 == 0:
print(f"[harvest] {len(frames)}/{len(times)} frames "
f"(t={t_s:.1f}s)", flush=True)
elapsed = time.perf_counter() - t0
finally:
cap.release()
return {"fps_nominal": round(fps, 3), "duration_s": round(duration, 3),
"n_frames_video": int(n_est), "every_s": every,
"n_wanted": len(times), "n_written": len(frames),
"n_from_shots": sum(1 for f in frames if f["source"] == "shots"),
"elapsed_s": round(elapsed, 1), "frames": frames}
def main(argv=None) -> int:
ap = argparse.ArgumentParser(
description="Harvest ND1 training frames from an archive VOD "
"(design doc 7 section 3.5). The archive is read only.")
ap.add_argument("--video", required=True, type=Path)
ap.add_argument("--out", required=True, type=Path,
help="the ND1 root; the frames go in <out>/<vod stem>/")
ap.add_argument("--every", type=float, default=EVERY_S,
help=f"seconds between frames (default {EVERY_S:g})")
ap.add_argument("--shots", type=Path, default=None,
help="an existing shots.json -- ADDS frames around every "
"shot onset (the rack before the break included)")
ap.add_argument("--limit", type=int, default=None,
help="stop after N frames (a smoke run)")
ap.add_argument("--archive-root", type=Path, default=None,
help="the read-only tree to fingerprint before and after "
"(default: the video's own drive root is NOT "
"guessed -- pass it explicitly to get the check)")
args = ap.parse_args(argv)
if not args.video.exists():
raise SystemExit(f"{args.video} does not exist")
out_dir = args.out / args.video.stem
print(f"[harvest] video : {args.video} "
f"({args.video.stat().st_size / 1e6:.1f} MB)")
print(f"[harvest] out : {out_dir}", flush=True)
fp_before = None
if args.archive_root is not None:
print(f"[harvest] fingerprinting {args.archive_root} ...", flush=True)
fp_before = fingerprint(args.archive_root)
print(f"[harvest] {len(fp_before)} files", flush=True)
print("[harvest] hashing the source ...", flush=True)
digest = sha256_of(args.video)
print(f"[harvest] sha256 {digest}", flush=True)
result = harvest(args.video, out_dir, args.every, args.shots, args.limit)
archive = None
if fp_before is not None:
fp_after = fingerprint(args.archive_root)
changed = sorted(k for k in set(fp_before) | set(fp_after)
if fp_before.get(k) != fp_after.get(k))
archive = {"root": str(args.archive_root), "n_files": len(fp_before),
"n_changed": len(changed), "changed": changed[:20]}
manifest = {
"dataset": "ND1",
"design_doc": "7_FarZone_Detection_BallID.md section 3.5",
"created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
"source": {"path": str(args.video), "name": args.video.name,
"size_bytes": args.video.stat().st_size,
"sha256": digest},
"sampling": {"every_s": args.every,
"shots": str(args.shots) if args.shots else None,
"shot_offsets_s": list(SHOT_OFFSETS_S)
if args.shots else None,
"limit": args.limit},
"jpeg_quality": JPEG_QUALITY,
"fps_nominal": result["fps_nominal"],
"duration_s": result["duration_s"],
"n_frames_video": result["n_frames_video"],
"n_wanted": result["n_wanted"],
"n_written": result["n_written"],
"n_from_shots": result["n_from_shots"],
"elapsed_s": result["elapsed_s"],
"archive_check": archive,
"frames": result["frames"],
}
(out_dir / "manifest.json").write_text(json.dumps(manifest, indent=1),
encoding="utf-8")
expected = result["duration_s"] / args.every
print(f"[harvest] {result['n_written']} frames written in "
f"{result['elapsed_s']:.1f}s "
f"({result['n_from_shots']} of them around a shot onset)")
print(f"[harvest] duration/every = {expected:.1f} -- written "
f"{result['n_written']}")
print(f"[harvest] manifest {out_dir / 'manifest.json'}")
if archive is not None:
if archive["n_changed"]:
print(f"[harvest] RED: {archive['n_changed']} file(s) under "
f"{archive['root']} changed -- {archive['changed'][:5]}")
return 7
print(f"[harvest] archive untouched ({archive['n_files']} files "
f"identical)")
else:
print("[harvest] NOTE: no --archive-root given, so nothing verified "
"that the source tree was left alone")
return 0
if __name__ == "__main__":
raise SystemExit(main())