Spaces:
Sleeping
Sleeping
| """Render overlay.mp4 per cú — lát A2 phần 3 (13/08/2026, BRIEF A2). | |
| Tái dụng cách vẽ ``render_overlay`` của ``scripts/broadcast/track_p0.py`` | |
| (bản Danh đã xem và chốt 13/08 vì trực quan — đuôi track + header trạng | |
| thái), nâng cấp theo design §3.2: vẽ MỌI bi (không chỉ cue), màu theo số | |
| bi (BallID hai tầng chạy per-frame trên chính crop bbox), chấm sự kiện | |
| va chạm suy ra. Script P0 giữ NGUYÊN không sửa — đây là code path mới | |
| cho app (nếp broadcast.py). | |
| Render chạy TRONG job analyze, TRƯỚC khi worker xoá clip (bài học lần đo | |
| A1: clip đã mất thì overlay không render nổi hồi tố). Detection per frame | |
| KHÔNG detect lại: ``RecordingModel`` bọc YOLO quanh MỘT lần gọi | |
| ``analyze_clip`` và ghi lại đúng những bbox mà tầng số đã dùng (nếp | |
| "adapter quanh một lần gọi, không sửa broadcast.py" — BG26/BG28/BG31); | |
| frame đọc lại từ chính clip còn trên đĩa. | |
| Chữ trên frame là ASCII không dấu — cv2.putText (font Hershey) không vẽ | |
| được tiếng Việt có dấu; phần chữ có dấu sống ở FE/JSON. | |
| Encode hai bước: cv2 VideoWriter mp4v (như P0) rồi ffmpeg re-encode H.264 | |
| yuv420p +faststart — mp4v mở được bằng player thường nhưng <video> của | |
| trình duyệt KHÔNG phát được; ffmpeg đã là điều kiện sẵn của luồng video | |
| đa cú (cut_clip). ffmpeg vắng → giữ bản mp4v + warning (player thường | |
| vẫn mở — gate A2.2 vế "ngoài app" vẫn sống). | |
| Import cv2 LƯỜI trong hàm (nếp poolcoach_cv) — venv app import module | |
| này được cho unit test phần thuần. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import shutil | |
| import subprocess | |
| from pathlib import Path | |
| import numpy as np | |
| from poolcoach_cv.broadcast import BALL_CLASSES, BALL_R_M | |
| TRAIL_N = 90 # số điểm track vẽ đuôi cue ball — như P0 | |
| # Màu bi theo số (BGR) — khớp bảng COLORS của FE (app.js) để hai nơi kể | |
| # cùng một câu chuyện màu; None/không nhận ra → xám. | |
| BALL_BGR = { | |
| 1: (24, 197, 245), 2: (225, 105, 65), 3: (47, 52, 224), | |
| 4: (191, 63, 125), 5: (19, 127, 240), 6: (70, 158, 46), | |
| 7: (32, 32, 139), 8: (40, 40, 40), 9: (24, 197, 245), | |
| } | |
| CUE_BGR = (0, 220, 220) # như P0 | |
| UNKNOWN_BGR = (180, 180, 180) | |
| TRAIL_BGR = (0, 220, 0) | |
| COLLISION_BGR = (0, 215, 255) # vàng gold — khớp dấu ✕ trên FE | |
| class RecordingModel: | |
| """Bọc YOLO model: forward nguyên ``predict``, ghi lại per-frame | |
| ``[(cls_name, conf, (x1, y1, x2, y2)), ...]`` theo ĐÚNG thứ tự frame | |
| mà analyze_clip đưa vào (batch tuần tự). Chỉ dùng quanh MỘT lần gọi | |
| analyze_clip — model gốc không đổi gì.""" | |
| def __init__(self, model) -> None: | |
| self._model = model | |
| self.frames: list[list[tuple]] = [] | |
| def predict(self, imgs, **kw): | |
| results = self._model.predict(imgs, **kw) | |
| seq = imgs if isinstance(imgs, (list, tuple)) else [imgs] | |
| for res in results[:len(seq)]: | |
| names = res.names | |
| dets = [] | |
| for b in res.boxes: | |
| x1, y1, x2, y2 = (float(v) for v in b.xyxy[0]) | |
| dets.append((names[int(b.cls)], float(b.conf), | |
| (x1, y1, x2, y2))) | |
| self.frames.append(dets) | |
| return results | |
| def _uncomp_xy(pts: np.ndarray, height_comp: dict | None) -> np.ndarray: | |
| """Đảo bù độ cao tâm bi: toạ độ P trong JSON (đã bù) → X' vị trí map | |
| mặt vải — chính là chỗ bi HIỆN RA trên frame khi chiếu ngược qua | |
| homography. Nghịch đảo đóng của camera.HeightCompensation.apply: | |
| X' = P·(1+k) − k·C, k = R/(h−R). Cờ off/thiếu camera → giữ nguyên.""" | |
| if not height_comp or not height_comp.get("on"): | |
| return pts | |
| h = height_comp.get("h_m") | |
| cx, cy = height_comp.get("cx_m"), height_comp.get("cy_m") | |
| if h is None or cx is None or cy is None: | |
| return pts | |
| k = BALL_R_M / (h - BALL_R_M) | |
| return np.asarray(pts, dtype=np.float64) * (1.0 + k) \ | |
| - k * np.array([cx, cy]) | |
| def _reencode_h264(src: Path, dst: Path) -> str | None: | |
| """mp4v → H.264 yuv420p +faststart để <video> trình duyệt phát được. | |
| Trả None khi OK; message warning khi ffmpeg vắng/hỏng (giữ bản mp4v).""" | |
| exe = shutil.which("ffmpeg") | |
| if exe is None: | |
| os.replace(src, dst) | |
| return ("máy không có ffmpeg — overlay giữ codec mp4v (player " | |
| "thường mở được, trình duyệt có thể không phát)") | |
| try: | |
| r = subprocess.run( | |
| [exe, "-y", "-v", "error", "-i", str(src), "-c:v", "libx264", | |
| "-preset", "veryfast", "-crf", "23", "-pix_fmt", "yuv420p", | |
| "-movflags", "+faststart", str(dst)], | |
| capture_output=True, text=True, timeout=300) | |
| except subprocess.TimeoutExpired: | |
| os.replace(src, dst) | |
| return "ffmpeg re-encode overlay quá 300s — overlay giữ codec mp4v" | |
| if r.returncode != 0 or not dst.exists(): | |
| os.replace(src, dst) | |
| err = (r.stderr or "?").strip().splitlines()[-1][:120] | |
| return f"ffmpeg re-encode overlay lỗi ({err}) — giữ codec mp4v" | |
| src.unlink(missing_ok=True) | |
| return None | |
| def render_shot_overlay(clip_path, corners_px, result: dict, | |
| frames_dets: list[list[tuple]], out_path, | |
| label: str = "", beat=None) -> str | None: | |
| """Render overlay MỘT cú từ clip còn trên đĩa + detection đã ghi. | |
| ``result``: JSON kết quả analyze (track/collisions/height_comp/ | |
| suspect_not_shot — toạ độ bàn, t_s tương đối frame đầu clip). | |
| ``frames_dets``: từ RecordingModel — index = thứ tự frame. | |
| ``label``: chữ đầu header (ASCII), ví dụ "shot 03". ``beat``: callback | |
| gọi định kỳ để giữ heartbeat trong job dài. Trả None khi trơn, message | |
| warning (tiếng Việt) khi có chuyện nhưng file vẫn ra được; ném khi | |
| không ra nổi file. | |
| """ | |
| import cv2 | |
| from poolcoach_cv.ballid import identify_balls | |
| from poolcoach_cv.broadcast import TABLE_L_M, TABLE_W_M | |
| from poolcoach_cv.homography import TableHomography, orient_corners | |
| out_path = Path(out_path) | |
| th = TableHomography(orient_corners(np.asarray(corners_px, | |
| dtype=np.float64)), | |
| table_w=TABLE_W_M, table_l=TABLE_L_M) | |
| hc = result.get("height_comp") | |
| track = result.get("track") or [] | |
| trail_t = np.array([float(p["t_s"]) for p in track]) | |
| trail_px = (th.table_to_px(_uncomp_xy( | |
| np.array([[float(p["x_m"]), float(p["y_m"])] for p in track]), hc)) | |
| if track else np.empty((0, 2))) | |
| cols = result.get("collisions") or [] | |
| col_t = [float(c["t_s"]) for c in cols] | |
| col_px = (th.table_to_px(_uncomp_xy( | |
| np.array([[float(c["x_m"]), float(c["y_m"])] for c in cols]), hc)) | |
| if cols else np.empty((0, 2))) | |
| suspect = bool((result.get("suspect_not_shot") or {}).get("flagged")) | |
| cap = cv2.VideoCapture(str(clip_path)) | |
| if not cap.isOpened(): | |
| raise ValueError("Không mở lại được clip để render overlay.") | |
| # temp CẠNH file đích (không phải system temp): os.replace cần cùng ổ | |
| # đĩa, và mkstemp giữ fd mở làm Windows khoá file không xoá được | |
| out_path.parent.mkdir(parents=True, exist_ok=True) | |
| tmp = out_path.with_name(out_path.stem + ".raw.mp4") | |
| vw = None | |
| try: | |
| fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 | |
| n_total = len(frames_dets) | |
| i = -1 | |
| t_first = None | |
| while True: | |
| ok, frame = cap.read() | |
| if not ok: | |
| break | |
| i += 1 | |
| t = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0 | |
| # POS_MSEC đọc SAU read() là PTS frame KẾ — xấp xỉ đủ cho vẽ; | |
| # số đo thật vẫn nằm ở JSON, overlay là lớp mắt người kiểm | |
| if t_first is None: | |
| t_first = t | |
| t_rel = t - t_first | |
| if vw is None: | |
| h, w = frame.shape[:2] | |
| vw = cv2.VideoWriter(str(tmp), | |
| cv2.VideoWriter_fourcc(*"mp4v"), | |
| fps, (w, h)) | |
| if beat is not None and i % 32 == 0: | |
| beat() | |
| dets = frames_dets[i] if i < n_total else [] | |
| balls = [(cls, conf, box) for cls, conf, box in dets | |
| if cls in BALL_CLASSES] | |
| others = [b for b in balls if b[0] != "Cue"] | |
| cue = next((b for b in balls if b[0] == "Cue"), None) | |
| assigns = [] | |
| if others: | |
| try: | |
| assigns, _wb = identify_balls( | |
| frame, [b[2] for b in others], | |
| cue[2] if cue else None) | |
| except Exception: # noqa: BLE001 — màu hỏng ≠ mất overlay | |
| assigns = [(None, 0.0)] * len(others) | |
| # mọi bi: khung màu theo số + nhãn số (design §3.2) | |
| for (cls, _conf, (x1, y1, x2, y2)), (num, _nc) in zip( | |
| others, assigns or [(None, 0.0)] * len(others)): | |
| col = BALL_BGR.get(num, UNKNOWN_BGR) | |
| cv2.rectangle(frame, (int(x1), int(y1)), | |
| (int(x2), int(y2)), col, 2) | |
| if num is not None: | |
| cv2.putText(frame, str(num), (int(x1), int(y1) - 4), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.55, col, 2) | |
| if cue is not None: | |
| x1, y1, x2, y2 = cue[2] | |
| cv2.rectangle(frame, (int(x1), int(y1)), | |
| (int(x2), int(y2)), CUE_BGR, 2) | |
| # đuôi track cue ball (như P0) — điểm đã đo tới thời điểm này | |
| m = trail_t <= t_rel + 1e-6 | |
| pts = trail_px[m][-TRAIL_N:] | |
| if len(pts) >= 2: | |
| cv2.polylines(frame, [pts.astype(np.int32)], False, | |
| TRAIL_BGR, 2) | |
| if len(pts): | |
| cv2.circle(frame, tuple(pts[-1].astype(int)), 14, | |
| TRAIL_BGR, 2) | |
| # chấm sự kiện va chạm — hiện từ lúc xảy ra tới hết clip | |
| for tc, (px, py), c in zip(col_t, col_px, cols): | |
| if t_rel + 1e-6 < tc: | |
| continue | |
| p = (int(px), int(py)) | |
| s = 9 | |
| cv2.line(frame, (p[0] - s, p[1] - s), (p[0] + s, p[1] + s), | |
| COLLISION_BGR, 2) | |
| cv2.line(frame, (p[0] - s, p[1] + s), (p[0] + s, p[1] - s), | |
| COLLISION_BGR, 2) | |
| cv2.putText(frame, f"{tc:.2f}s {c.get('contact') or ''}", | |
| (p[0] + s + 3, p[1] - 3), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLLISION_BGR, 1) | |
| # header trạng thái (ASCII — Hershey không có dấu) | |
| n_seen = int((trail_t <= t_rel + 1e-6).sum()) | |
| head = (f"{label} t={t_rel:6.2f}s frame {i + 1}" | |
| f" track {n_seen}/{len(trail_t)}") | |
| if suspect: | |
| head += " [NGHI KHONG PHAI CU DANH]" | |
| cv2.rectangle(frame, (0, 0), (frame.shape[1], 30), (0, 0, 0), -1) | |
| cv2.putText(frame, head, (8, 22), cv2.FONT_HERSHEY_SIMPLEX, | |
| 0.6, (0, 0, 255) if suspect else (0, 220, 0), 2) | |
| vw.write(frame) | |
| finally: | |
| if vw is not None: | |
| vw.release() | |
| cap.release() | |
| if vw is None: | |
| tmp.unlink(missing_ok=True) | |
| raise ValueError("Clip không đọc được frame nào để render overlay.") | |
| return _reencode_h264(tmp, out_path) | |