scripts/gaming_detect_bench.py
Rama main · réplica de lectura del repositorio original.
"""Measure Gaming / Split-Screen detection on real footage.
Two stages, so the decision rules can be tuned without re-running the models:
measure: for every video given (or every .mp4 in a folder), cut a few
windows and save each one's person tracks, TalkNet scores and a
mid-window frame under <out>/raw/.
decide: apply gaming.detect's rules to the saved windows: per window the
faces and the streamer TalkNet picks, per video the webcam the
windows agree on, and the layout each window gets. Writes
bench.json and one PNG per window: the frame with every candidate
(red, with its speaking share and confidence), the webcam (green)
and the game band (blue), next to a preview of the 1080x1920
result.
Check the pictures by eye: is the green box the streamer's webcam, and is
chat outside the blue band?
python scripts/gaming_detect_bench.py measure <folder or files> --out <dir>
python scripts/gaming_detect_bench.py decide --out <dir>
"""
import argparse
import json
import pickle
import subprocess
import sys
import tempfile
import time
from collections import defaultdict
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
import cv2 # noqa: E402
import numpy as np # noqa: E402
from core.binaries import ffmpeg # noqa: E402
from gaming import detect, layout # noqa: E402
from video.capture import video_capture # noqa: E402
from video.tracker import _ASD_MAX_TRACKS, _ASD_MIN_FACE_SAMPLES # noqa: E402
def cut(src: Path, start: float, length: float, dst: Path) -> bool:
"""A re-encoded window, so audio and video start on the same frame."""
r = subprocess.run(
[ffmpeg(), "-v", "error", "-y", "-ss", str(start), "-i", str(src), "-t", str(length),
"-c:v", "libx264", "-preset", "ultrafast", "-crf", "23", "-c:a", "aac", str(dst)],
capture_output=True, timeout=600,
)
return r.returncode == 0 and dst.exists() and dst.stat().st_size > 0
def mid_frame(clip: Path, t: float):
with video_capture(clip, required=False) as cap:
if cap is None:
return None
cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000)
ok, frame = cap.read()
return frame if ok else None
def measure(args):
files = []
for i in args.inputs:
files += sorted(i.glob("*.mp4")) if i.is_dir() else [i]
files = [f for f in files if f.stat().st_size > 0]
raw = args.out / "raw"
raw.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryDirectory() as tmp:
for f in files:
for start in (float(s) for s in args.starts.split(",")):
key = f"{f.stem}_{int(start)}"
if (raw / f"{key}.pkl").exists():
continue
clip = Path(tmp) / f"{key}.mp4"
if not cut(f, start, args.length, clip):
print(f"{key}: could not cut")
continue
t0 = time.perf_counter()
tracks, w, h, fps, duration, n = detect.sample_tracks(clip, args.model)
t_tracks = time.perf_counter() - t0
# Score the longest-seen faces whatever their presence, so the
# presence bar can be tuned in `decide`.
ids = sorted((tid for tid, tr in tracks.items() if len(tr.seen) >= _ASD_MIN_FACE_SAMPLES),
key=lambda tid: -len(tracks[tid].times))[:_ASD_MAX_TRACKS]
t1 = time.perf_counter()
scored = detect.speaking_scores(clip, tracks, ids, duration, fps) if ids else None
t_asd = time.perf_counter() - t1
frame = mid_frame(clip, args.length / 2)
if frame is not None:
cv2.imwrite(str(raw / f"{key}.jpg"), frame)
with open(raw / f"{key}.pkl", "wb") as fh:
pickle.dump({"video": f.stem, "start": start, "tracks": tracks, "w": w, "h": h,
"duration": duration, "n": n, "scored": scored,
"time_tracks": t_tracks, "time_talknet": t_asd}, fh)
print(f"{key}: {len(tracks)} tracks, {len(ids)} scored, "
f"{t_tracks:.1f}s tracks + {t_asd:.1f}s TalkNet", flush=True)
def preview(frame, p: layout.Plan):
"""The 1080x1920 result, roughly, drawn with OpenCV (the real render is
FFmpeg; blur fills are left dark here)."""
out = np.zeros((layout.OUT_H, layout.OUT_W, 3), dtype=np.uint8)
for e in p.elements:
sx, sy, sw, sh = e.src
dx, dy, dw, dh = e.dest
region = frame[sy:sy + sh, sx:sx + sw]
if e.fit == "contain":
scale = min(dw / sw, dh / sh)
fw, fh = max(2, int(sw * scale)), max(2, int(sh * scale))
top = {"top": 0, "bottom": dh - fh}.get(e.anchor, (dh - fh) // 2)
out[dy + top:dy + top + fh, dx + (dw - fw) // 2:dx + (dw - fw) // 2 + fw] = cv2.resize(region, (fw, fh))
else:
tile = cv2.resize(region, (dw, dh))
if e.shift:
tile = np.vstack([np.zeros((e.shift, dw, 3), dtype=np.uint8), tile[:dh - e.shift]])
out[dy:dy + dh, dx:dx + dw] = tile
return out
def annotate(frame, finding: detect.ClipFinding, cam, p: layout.Plan, kind: str):
img = frame.copy()
h, w = img.shape[:2]
for tid, face in finding.faces.items():
x, y, bw, bh = face.box
cv2.rectangle(img, (int(x * w), int(y * h)), (int((x + bw) * w), int((y + bh) * h)), (0, 0, 255), 2)
conf = "-" if face.confidence is None else f"{face.confidence:.1f}"
cv2.putText(img, f"#{tid} p={face.presence:.2f} s={face.speaking_share:.2f} c={conf}",
(int(x * w) + 4, int(y * h) + 26), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
if cam:
x, y, bw, bh = cam
cv2.rectangle(img, (int(x * w), int(y * h)), (int((x + bw) * w), int((y + bh) * h)), (0, 255, 0), 5)
gx, gy, gw, gh = p.element("game").src
cv2.rectangle(img, (gx, gy), (gx + gw - 1, gy + gh - 1), (255, 128, 0), 4)
cv2.putText(img, kind, (20, h - 30), cv2.FONT_HERSHEY_SIMPLEX, 1.6, (255, 255, 255), 4)
scale = layout.OUT_H / h / 2
left = cv2.resize(img, (int(w * scale), int(h * scale)))
right = cv2.resize(preview(frame, p), (layout.OUT_W // 2, layout.OUT_H // 2))
pad = np.zeros((right.shape[0], left.shape[1], 3), dtype=np.uint8)
top = (right.shape[0] - left.shape[0]) // 2
pad[top:top + left.shape[0]] = left
return np.hstack([pad, right])
def decide(args):
raw = args.out / "raw"
windows = defaultdict(list)
for pkl in sorted(raw.glob("*.pkl")):
with open(pkl, "rb") as fh:
d = pickle.load(fh)
d["key"] = pkl.stem
windows[d["video"]].append(d)
results = {}
for video, ws in windows.items():
ws.sort(key=lambda d: d["start"])
findings = []
for d in ws:
ids = set(detect.candidates(d["tracks"], d["n"]))
scored = d["scored"]
if scored is not None:
scores, loud = scored
scores = {tid: s for tid, s in scores.items() if tid in ids}
scored = (scores, loud) if scores else None
findings.append(detect.judge(d["tracks"], d["n"], d["w"], d["h"], scored)
if ids and d["w"] else detect.ClipFinding(reason="nobody on screen for most of the clip"))
cam = detect.video_cam(findings)
for d, finding in zip(ws, findings):
if cam is not None and detect.on_screen(cam, finding):
kind, box = "split", cam
elif finding.camera is not None and not finding.camera.overlay:
kind, box = "standard", None
else:
kind, box = "fill", None
p = layout.plan(d["w"] or 1920, d["h"] or 1080, {"preset": "split", "cam": list(box) if box else None})
results[d["key"]] = {
"video": video, "kind": kind, "video_cam": [round(v, 3) for v in cam] if cam else None,
"streamer": finding.streamer, "reason": finding.reason,
"faces": {str(t): {"box": [round(v, 3) for v in f.box], "presence": round(f.presence, 2),
"share": round(f.speaking_share, 2),
"confidence": None if f.confidence is None else round(f.confidence, 2),
"overlay": f.overlay} for t, f in finding.faces.items()},
"time": round(d["time_tracks"] + d["time_talknet"], 1),
}
faces = ", ".join(f"#{t} p={f.presence:.2f} s={f.speaking_share:.2f} "
f"c={'-' if f.confidence is None else round(f.confidence, 2)}"
f"{' ov' if f.overlay else ''}" for t, f in finding.faces.items())
print(f"{d['key']:18} {kind:8} streamer={finding.streamer} [{faces}]")
frame_path = raw / f"{d['key']}.jpg"
if frame_path.exists() and d["w"]:
frame = cv2.imread(str(frame_path))
cv2.imwrite(str(args.out / f"{d['key']}.png"), annotate(frame, finding, box, p, kind))
print(f"== {video}: webcam {results[ws[-1]['key']]['video_cam']}")
(args.out / "bench.json").write_text(json.dumps(results, indent=1))
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
sub = ap.add_subparsers(dest="stage", required=True)
m = sub.add_parser("measure")
m.add_argument("inputs", nargs="+", type=Path)
m.add_argument("--out", type=Path, required=True)
m.add_argument("--starts", default="15,75,135", help="window starts, seconds")
m.add_argument("--length", type=float, default=40.0)
m.add_argument("--model", default="yolov8n-pose.pt")
d = sub.add_parser("decide")
d.add_argument("--out", type=Path, required=True)
args = ap.parse_args()
measure(args) if args.stage == "measure" else decide(args)
if __name__ == "__main__":
main()