"""BallID — gán SỐ bi 9-ball hai tầng theo MÀU, không train (06/08/2026). Thực thi `PoolCoach_CV_BallID_Design.md` (chốt 06/08): dataset pix2pockets KHÔNG có nhãn số nên tầng 1 là baseline màu thuần numpy; khi có bộ ảnh VN gán nhãn sẽ swap classifier train cùng chữ ký (crop → vector cost 9 số), tầng 2 giữ nguyên. - **Tầng 1** (`score_balls`): mỗi bbox bi → vector cost 9 số. Crop co 60% bbox quanh tâm (né nỉ + bóng) → chuẩn hoá trắng theo median crop bi CUE (đèn vàng quán VN ám cue bao nhiêu thì ám các bi khác bấy nhiêu; không cue → bỏ, cờ ``wb=False``) → tách pixel TRẮNG/MÀU → cost = khoảng cách hue circular (trọng số S/V thấp hơn H) tới palette + phạt lệch ``white_ratio``; bi 8 chấm theo V thấp thay vì hue. - **Tầng 2** (`assign_numbers`): ma trận cost → gán greedy TẤT ĐỊNH với ràng buộc mỗi số ≤ 1 bi (duyệt cặp cost tăng dần; hoà cost → index nhỏ trước); cost > ``cost_max`` → không gán (``None``). ``number_conf`` = nghịch biến cost + độ tách với ứng viên nhì. Thuần numpy như `homography.py` — KHÔNG import cv2/torch → unit test chạy trong cả `poolcoach-env` lẫn `poolcoach-cv-env`. **Thứ tự kênh màu là BGR (quy ước OpenCV)** — mọi ảnh/crop vào đây đều BGR uint8 (hoặc float [0,255]). Lẫn RGB↔BGR làm bi 2 (xanh dương) thành cam và ngược lại — test synthetic viết màu BGR tường minh để khoá đúng điều này. """ from __future__ import annotations import numpy as np _EPS = 1e-12 # --------------------------------------------------------------------------- # BẢNG HẰNG DUY NHẤT của BallID — palette + mọi ngưỡng (design §2). # # Căn cứ: màu bi 9-ball chuẩn (1 vàng · 2 xanh dương · 3 đỏ · 4 tím · 5 cam · # 6 xanh lá · 7 đỏ nâu · 8 đen · 9 sọc trắng-vàng), giá trị H/S/V ước lượng # tay từ bộ bi tiêu chuẩn dưới đèn trung tính. Quy ước HSV NỘI BỘ của module: # H ∈ [0,360)°, S/V ∈ [0,1] (KHÔNG phải hue OpenCV [0,180) — module không # dùng cv2). KHÔNG chỉnh các số này sau khi đã thấy số đo ảnh thật: thấy số # rồi muốn chỉnh là vòng 2, ghi HANDOFF để Cowork quyết (design §4). BALLID = { # Palette màu đại diện pixel MÀU (median HSV của crop): số → (h, s, v). # Bi 8 KHÔNG có mục ở đây — đen không có hue tin được, chấm riêng theo # V thấp bằng bộ ngưỡng v8/s8 bên dưới. Bi 9 cùng hue vàng với bi 1 — # phân biệt bằng white_ratio (sọc trắng), không phải bằng màu. "palette": { 1: (52.0, 0.85, 0.85), # vàng 2: (225.0, 0.80, 0.55), # xanh dương 3: (5.0, 0.90, 0.70), # đỏ 4: (285.0, 0.65, 0.45), # tím 5: (27.0, 0.90, 0.80), # cam 6: (135.0, 0.75, 0.45), # xanh lá 7: (350.0, 0.75, 0.40), # đỏ nâu (hue sát bi 3, tối + xỉn hơn) 9: (52.0, 0.85, 0.85), # phần vàng của sọc }, # Khoảng white_ratio kỳ vọng [lo, hi] — lệch ra ngoài phạt tuyến tính # theo khoảng cách tới biên (design §2: trơn ~0–0.15, sọc 9 ~0.3–0.6). "white_range": {"solid": (0.0, 0.15), "eight": (0.0, 0.10), "stripe": (0.30, 0.60)}, # Trọng số cost: hue nặng nhất, S/V phụ (design §2); white ngang hue để # cặp 1↔9 (cùng hue) tách được bằng sọc. "w_hue": 1.0, "w_sat": 0.25, "w_val": 0.25, "w_white": 1.0, # Bi 8: cost = w8_val·max(0, v−v8_max) + w8_sat·max(0, s−s8_max) + phạt # trắng. Bi đen chuẩn v≈0.1–0.3; hệ số dốc để bi màu (v≈0.5+) không với # tới số 8 dưới ngưỡng cost_max. "v8_max": 0.35, "s8_max": 0.40, "w8_val": 3.0, "w8_sat": 1.0, # Pixel TRẮNG = sáng cao + bão hoà thấp (sọc/bóng loé): V ≥ white_v_min # VÀ S ≤ white_s_max. Sau WB, trắng thật S ≈ 0.0–0.15; ám vàng KHÔNG WB # đẩy S lên ~0.5 → trắng "biến mất" — chính là cái đối chứng âm đo. "white_v_min": 0.65, "white_s_max": 0.25, # Crop co 60% kích thước bbox quanh tâm — né nỉ bàn + bóng đổ (design §2). "crop_shrink": 0.60, # Dưới tỉ lệ pixel MÀU này thì median màu hết nghĩa (bi gần toàn trắng — # vd bi trắng thứ hai bị hạ từ cue xuống ball) → dùng median MỌI pixel; # white_ratio cao sẽ tự đẩy cost vượt cost_max → number=None. "min_color_frac": 0.05, # Tầng 2: cost gán > cost_max → không gán (None); margin_ref chuẩn hoá # độ tách với ứng viên nhì khi tính conf. "cost_max": 0.60, "margin_ref": 0.15, # number_conf = w_conf_cost·(1 − cost/cost_max) # + w_conf_margin·min(1, margin/margin_ref) ∈ [0,1] "w_conf_cost": 0.6, "w_conf_margin": 0.4, } NUMBERS = (1, 2, 3, 4, 5, 6, 7, 8, 9) # cột j của ma trận cost ↔ số j+1 # ------------------------------------------------------------------ tầng 1 def _bgr_to_hsv(bgr: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """(N,3) BGR ∈ [0,1] → (H ∈ [0,360), S ∈ [0,1], V ∈ [0,1]). Công thức hexcone chuẩn; KHÔNG theo scale hue [0,180) của OpenCV. """ b, g, r = bgr[..., 0], bgr[..., 1], bgr[..., 2] v = bgr.max(axis=-1) c = v - bgr.min(axis=-1) s = np.where(v > _EPS, c / np.maximum(v, _EPS), 0.0) safe_c = np.maximum(c, _EPS) h = np.select( [v == r, v == g], [(g - b) / safe_c % 6.0, (b - r) / safe_c + 2.0], default=(r - g) / safe_c + 4.0, ) * 60.0 h = np.where(c > _EPS, h % 360.0, 0.0) return h, s, v def _hue_circular_median(h: np.ndarray) -> float: """Median hue có xử lý vòng tròn — bi đỏ/nâu vắt qua mốc 0°/360°. Median thẳng trên {358°, 2°} ra ~180° (cyan!) — sai chết người với bi 3/7. Cách làm: xoay cụm hue về quanh 180° theo hướng trung bình vector (atan2), lấy median trong hệ đã xoay, xoay ngược lại. """ rad = np.deg2rad(h) mean_ang = np.rad2deg(np.arctan2(np.sin(rad).mean(), np.cos(rad).mean())) shift = 180.0 - mean_ang return float((np.median((h + shift) % 360.0) - shift) % 360.0) def _crop(img: np.ndarray, box) -> np.ndarray | None: """Crop quanh tâm bbox, co ``crop_shrink`` kích thước; None nếu rỗng. bbox = (x1, y1, x2, y2) pixel trên CHÍNH ảnh đưa vào (không quy đổi scale ở đây — người gọi tự bảo đảm bbox đo trên cùng ảnh). """ x1, y1, x2, y2 = (float(v) for v in box) cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0 hw = (x2 - x1) * BALLID["crop_shrink"] / 2.0 hh = (y2 - y1) * BALLID["crop_shrink"] / 2.0 h_img, w_img = img.shape[:2] xa, xb = max(0, round(cx - hw)), min(w_img, round(cx + hw)) ya, yb = max(0, round(cy - hh)), min(h_img, round(cy + hh)) if xb - xa < 1 or yb - ya < 1: return None return img[ya:yb, xa:xb] def _features(crop_bgr: np.ndarray, wb_scale: np.ndarray | None ) -> tuple[float, float, float, float]: """Crop BGR → (hue_median, sat_median, val_median, white_ratio). ``wb_scale`` (3,) nhân vào 3 kênh BGR (chuẩn hoá trắng theo cue) trước khi đọc màu; None = không WB. Median lấy trên pixel MÀU (không trắng); hue dùng median vòng tròn (xem `_hue_circular_median`). """ px = crop_bgr.reshape(-1, 3).astype(np.float64) if wb_scale is not None: px = np.clip(px * wb_scale, 0.0, 255.0) h, s, v = _bgr_to_hsv(px / 255.0) white = (v >= BALLID["white_v_min"]) & (s <= BALLID["white_s_max"]) white_ratio = float(white.mean()) colored = ~white if colored.mean() < BALLID["min_color_frac"]: colored = np.ones_like(white) # gần toàn trắng — median mọi pixel return (_hue_circular_median(h[colored]), float(np.median(s[colored])), float(np.median(v[colored])), white_ratio) def _white_penalty(white_ratio: float, lo: float, hi: float) -> float: """Khoảng cách tuyến tính của white_ratio tới đoạn [lo, hi] (0 nếu trong).""" return max(0.0, lo - white_ratio, white_ratio - hi) def _cost_vector(h: float, s: float, v: float, white_ratio: float) -> np.ndarray: """Đặc trưng một crop → vector cost 9 số (index j ↔ số j+1). Đây là CHỮ KÝ swap classifier (design §2): sau này thay bằng model trả −log p, tầng 2 giữ nguyên. """ P = BALLID cost = np.empty(9, dtype=np.float64) for j, num in enumerate(NUMBERS): if num == 8: pen = _white_penalty(white_ratio, *P["white_range"]["eight"]) cost[j] = (P["w8_val"] * max(0.0, v - P["v8_max"]) + P["w8_sat"] * max(0.0, s - P["s8_max"]) + P["w_white"] * pen) continue h_ref, s_ref, v_ref = P["palette"][num] d_h = abs(h - h_ref) % 360.0 d_h = min(d_h, 360.0 - d_h) / 180.0 rng = P["white_range"]["stripe" if num == 9 else "solid"] cost[j] = (P["w_hue"] * d_h + P["w_sat"] * abs(s - s_ref) + P["w_val"] * abs(v - v_ref) + P["w_white"] * _white_penalty(white_ratio, *rng)) return cost def score_balls(img_bgr: np.ndarray, ball_boxes, cue_box=None ) -> tuple[np.ndarray, bool]: """Tầng 1: ảnh BGR + bbox các bi (+ bbox cue) → ma trận cost (N, 9). Trả (cost, wb): ``wb=True`` khi có chuẩn hoá trắng theo crop cue (cue hợp lệ trong ảnh); ``False`` khi không cue → đọc màu thô (design §2). Bbox hỏng/ngoài ảnh → hàng cost = +inf → tầng 2 tự trả ``None``. """ img = np.asarray(img_bgr) if img.ndim != 3 or img.shape[2] != 3: raise ValueError(f"img_bgr phải là (H,W,3) BGR, nhận shape {img.shape}.") wb_scale = None if cue_box is not None: cue_crop = _crop(img, cue_box) if cue_crop is not None: med = np.median(cue_crop.reshape(-1, 3).astype(np.float64), axis=0) if np.all(med > _EPS): # cue trắng ám đèn thế nào thì scale kênh về trung tính thế ấy wb_scale = med.mean() / med wb = wb_scale is not None cost = np.full((len(ball_boxes), 9), np.inf, dtype=np.float64) for i, box in enumerate(ball_boxes): crop = _crop(img, box) if crop is None: continue cost[i] = _cost_vector(*_features(crop, wb_scale)) return cost, wb # ------------------------------------------------------------------ tầng 2 def assign_numbers(cost: np.ndarray) -> list[tuple[int | None, float]]: """Tầng 2: ma trận cost (N,9) → [(số 1–9 | None, number_conf)] mỗi bi. Greedy tất định (design §2, chốt thay Hungarian có chủ đích): duyệt cặp (bi, số) theo cost tăng dần, gán khi cả bi lẫn số còn trống; hoà cost → (index bi, index số) nhỏ trước — sort tuple bảo đảm tất định. Cost gán > ``cost_max`` (hoặc không còn số trống, tức N > 9) → ``None``. ``number_conf`` ∈ [0,1] nghịch biến cost + margin với ứng viên nhì TRONG HÀNG của bi đó (margin âm — bị ràng buộc ép nhận số không phải rẻ nhất — kẹp về 0 = "hai số tranh nhau, không chắc"). """ c = np.asarray(cost, dtype=np.float64) if c.ndim != 2 or c.shape[1] != 9: raise ValueError(f"cost phải là (N,9), nhận {c.shape}.") P = BALLID pairs = sorted( (c[i, j], i, j) for i in range(c.shape[0]) for j in range(9) if np.isfinite(c[i, j]) and c[i, j] <= P["cost_max"] ) ball_of: dict[int, int] = {} # j → i num_of: dict[int, int] = {} # i → j for cost_ij, i, j in pairs: if i in num_of or j in ball_of: continue ball_of[j] = i num_of[i] = j out: list[tuple[int | None, float]] = [] for i in range(c.shape[0]): j = num_of.get(i) if j is None: out.append((None, 0.0)) continue others = np.delete(c[i], j) runner_up = float(others.min()) if np.isfinite(others).any() else np.inf margin = max(0.0, runner_up - float(c[i, j])) conf = (P["w_conf_cost"] * max(0.0, 1.0 - float(c[i, j]) / P["cost_max"]) + P["w_conf_margin"] * min(1.0, margin / P["margin_ref"])) out.append((NUMBERS[j], round(min(1.0, conf), 3))) return out def identify_balls(img_bgr: np.ndarray, ball_boxes, cue_box=None ) -> tuple[list[tuple[int | None, float]], bool]: """Trọn gói hai tầng: ảnh + bbox → [(number|None, conf)] theo thứ tự ``ball_boxes``, kèm cờ ``wb``. Không bi nào → danh sách rỗng.""" if not len(ball_boxes): return [], cue_box is not None cost, wb = score_balls(img_bgr, ball_boxes, cue_box) return assign_numbers(cost), wb