Spaces:
Running
Running
Download scripts/harvest_frames.py from masterdanh/poolcoach-analyzer: direct link, hf CLI and curl.
- Browser
- Download file 12.4 kB
-
https://huggingface.co/spaces/masterdanh/poolcoach-analyzer/resolve/main/scripts/harvest_frames.py
- Command line
-
hf download hf://spaces/masterdanh/poolcoach-analyzer/scripts/harvest_frames.py
-
curl -L -o harvest_frames.py https://huggingface.co/spaces/masterdanh/poolcoach-analyzer/resolve/main/scripts/harvest_frames.py
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()) | |