app.py

"""
VAR-style offside visualizer — Gradio app.

Flow:
  upload -> scrub to the frame -> draw 2 goal-parallel lines on that frame
  -> Detect players (GPU, once/frame) -> click players to select
  -> mark defenders -> Build the 3D scene with a draggable offside plane.

The GPU is touched ONLY inside pipeline.gpu (detect_frame / reconstruct_selected);
every other callback here runs on cached numpy and stays on the CPU.
"""

import os

# headless GL + CUDA fragmentation hygiene must be set before any heavy import
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import numpy as np
import gradio as gr

from pipeline.video import probe_video, grab_frame
from pipeline.overlay import (annotate_detections, draw_lines, draw_masks,
                              pick_box, pick_mask)
from pipeline.autolines import propose_lines
from pipeline import geometry as G

SEG_STYLE = "Segments"


def _render(frame, people, selected, style):
    """Draw detections in the chosen style (boxes vs segments)."""
    if style == SEG_STYLE:
        return draw_masks(frame, people, selected)
    return annotate_detections(frame, people, selected)


def _pick(people, x, y, style):
    return pick_mask(people, x, y) if style == SEG_STYLE else pick_box(people, x, y)


# ============================================================================
# Stage (a): upload + frame scrubbing  (pure CPU)
# ============================================================================
def on_upload(video_path):
    """New clip: size the slider, show frame 0, reset lines/frame state."""
    if not video_path:
        return (gr.update(maximum=1, value=0), None, "Upload a clip to begin.",
                0, None, [], "")
    n, fps = probe_video(video_path)
    n_max = max(n - 1, 0)
    frame = grab_frame(video_path, 0)
    return (
        gr.update(maximum=max(n_max, 1), value=0),
        frame,
        f"{n} frames @ {fps:.1f} fps — scrub to the moment the ball is played, "
        "then auto-detect or click the 2 goal-parallel lines.",
        n_max,
        frame,        # st_frame: clean copy for line redraws
        [],           # st_lines reset
        "",           # line_status reset
        [],           # st_families reset
        0,            # st_fam_idx reset
    )


def on_scrub(video_path, idx):
    """Show the new frame and reset any half-drawn / proposed lines on it."""
    if not video_path:
        return None, None, [], "", [], 0
    frame = grab_frame(video_path, int(idx))
    return frame, frame, [], "", [], 0


def step_frame(idx, delta, n_max):
    return max(0, min(int(n_max), int(idx) + delta))


# ============================================================================
# Stage (b): draw 2 goal-parallel lines on the scrubbed frame  (CPU)
# ============================================================================
def on_line_click(line_pts, clean_frame, evt: gr.SelectData):
    """Collect 4 clicks = 2 lines; always redraw from the clean frame (no drift).

    A manual click clears any auto-proposed families (so Flip stops applying).
    """
    pts = list(line_pts) if line_pts else []
    if len(pts) >= 4:
        pts = []
    pts.append([float(evt.index[0]), float(evt.index[1])])
    img = draw_lines(clean_frame, pts)
    status = {1: "line 1: 1/2", 2: "line 1 set", 3: "line 2: 1/2",
              4: "both lines set ✓ — now Detect players"}.get(len(pts), "")
    return pts, img, status, [], 0


def on_auto_lines(clean_frame):
    """Propose the 2 goal-parallel lines from detected pitch lines."""
    if clean_frame is None:
        return [], None, "Scrub to a frame first.", [], 0
    fams = propose_lines(clean_frame)
    if not fams:
        return [], gr.update(), ("No clear pitch lines found — draw the 2 "
                                 "goal-parallel lines by hand."), [], 0
    pts = fams[0]["pts"]
    img = draw_lines(clean_frame, pts)
    status = (f"Proposed goal-parallel lines (direction 1/{len(fams)}). If the "
              "offside axis looks wrong, click **Flip line direction**, or just "
              "click the frame to redraw by hand.")
    return pts, img, status, fams, 0


def on_flip_lines(families, fam_idx, clean_frame):
    """Switch the proposal to the other detected line-family."""
    if not families:
        return gr.update(), gr.update(), "Run **Auto-detect lines** first.", gr.update()
    if len(families) < 2:
        return (families[0]["pts"], draw_lines(clean_frame, families[0]["pts"]),
                "Only one line direction was found — redraw by hand if it's wrong.", 0)
    new_idx = (int(fam_idx) + 1) % len(families)
    pts = families[new_idx]["pts"]
    return (pts, draw_lines(clean_frame, pts),
            f"Line direction {new_idx + 1}/{len(families)}.", new_idx)


# ============================================================================
# Stage (c): detect (GPU, cached) + show boxes  (CPU after the call)
# ============================================================================
def on_detect(video_path, idx, conf, style):
    """Detect step: RF-DETR-Seg only (boxes + masks). No mesh reconstruction here."""
    from pipeline.gpu import detect_frame
    if not video_path:
        return None, [], [], "Upload a clip first.", gr.update(choices=[], value=[])
    people = detect_frame(video_path, idx, conf)
    if not people:
        return (None, [], [], "No players detected — lower the confidence slider.",
                gr.update(choices=[], value=[]))
    annotated = _render(grab_frame(video_path, idx), people, [], style)
    unit = "silhouette" if style == SEG_STYLE else "box"
    msg = (f"Detected {len(people)} players. Click a player's {unit} to select "
           "(click again to deselect).")
    return annotated, people, [], msg, gr.update(choices=[], value=[])


def on_restyle(people, selected, clean_frame, style):
    """Re-render the current detection when the box/segment toggle changes."""
    if not people or clean_frame is None:
        return gr.update()
    return _render(clean_frame, people, selected or [], style)


# ============================================================================
# Stage (d): click players to select  (CPU)
# ============================================================================
def on_select_player(people, selected, clean_frame, cur_def, style, evt: gr.SelectData):
    """Toggle the player clicked; refresh highlight + defender choices (box or segment)."""
    if not people:
        return None, selected or [], "Detect players first.", gr.update()
    hit = _pick(people, evt.index[0], evt.index[1], style)
    sel = list(selected or [])
    if hit is not None:
        sel.remove(hit) if hit in sel else sel.append(hit)
    sel = sorted(sel)
    img = _render(clean_frame, people, sel, style)
    msg = (f"Selected players: {sel}. Mark defenders below, then Build."
           if sel else "Click a player to select.")
    keep_def = [d for d in (cur_def or []) if d in sel]
    return img, sel, msg, gr.update(choices=sel, value=keep_def)


# ============================================================================
# Stage (e)+(f): place players + build the Plotly scene with a draggable plane
# ============================================================================
def on_build(video_path, idx, people_det, selected_ids, line_pts, flip_up,
             attack_dir, defender_ids):
    from pipeline.gpu import reconstruct_selected, get_faces
    if not selected_ids:
        return None, "Click at least one player to select.", gr.update(), None, +1, [], {}
    if not line_pts or len(line_pts) < 4:
        return None, "Draw 2 goal-parallel lines (4 points) on the frame first.", \
               gr.update(), None, +1, [], {}

    # Reconstruct ONLY the selected players' boxes (the heavy GPU step).
    selected_ids = sorted(int(i) for i in selected_ids)
    boxes = [people_det[i]["bbox"] for i in selected_ids]
    recon = reconstruct_selected(video_path, idx, boxes)
    if not recon:
        return None, "Reconstruction returned no meshes.", gr.update(), None, +1, [], {}
    # Key meshes back to their original detection ids (recon order == boxes order).
    people = {selected_ids[k]: recon[k] for k in range(len(recon))}
    faces = get_faces()
    h, w = grab_frame(video_path, idx).shape[:2]
    focal = people[selected_ids[0]]["focal_length"]
    gdir = G.goal_dir_from_lines(line_pts, focal, w, h)

    placed = G.place_players(people, selected_ids, gdir, flip_up=flip_up)
    med = float(np.median([placed[i][:, 2].max() for i in placed]))
    masks = {i: G.non_arm_mask(people[i]) for i in selected_ids}  # exclude arms/hands

    attack_sign = +1 if attack_dir == "toward +X" else -1
    dset = [int(d) for d in (defender_ids or [])]
    plane_x = G.offside_plane_x(placed, attack_sign, dset, masks)
    fig = G.build_scene(placed, faces, plane_x, attack_sign, dset, masks)

    allX = np.vstack(list(placed.values()))[:, 0]
    x0, x1 = float(allX.min() - 4), float(allX.max() + 4)
    plane_update = gr.update(minimum=x0, maximum=x1, value=float(plane_x),
                             visible=True, label="Drag the offside plane (X, m)")

    warn = "  ⚠ heights look wrong — toggle 'flip up'." if med < 1.0 else ""
    return (fig, f"Median player height {med:.2f} m (expect ~1.7–1.9).{warn}",
            plane_update, placed, attack_sign, dset, masks)


def on_plane(placed, plane_x, attack_sign, defender_ids, masks):
    """Re-render the scene at a new plane X — pure CPU on the cached placement."""
    from pipeline.gpu import get_faces
    if not placed:
        return gr.update()
    return G.build_scene(placed, get_faces(), float(plane_x),
                         int(attack_sign), defender_ids or [], masks)


# ============================================================================
# UI
# ============================================================================
with gr.Blocks(title="VAR Offside Visualizer") as demo:
    gr.Markdown(
        "## VAR-style Offside Visualizer\n"
        "Upload → scrub → **click 2 goal-parallel lines** → Detect → "
        "**click players to select** → mark defenders → Build.\n\n"
        "_Scale comes from reconstructed body height, so positions are approximate "
        "metres — good for relative offside ordering, not sub-10 cm calls._"
    )

    # session state
    st_nmax = gr.State(0)        # last valid frame index
    st_frame = gr.State(None)    # clean RGB of the current frame
    st_people = gr.State([])     # slim detections
    st_lines = gr.State([])      # clicked / proposed line points
    st_families = gr.State([])   # auto-proposed line families
    st_fam_idx = gr.State(0)     # which proposed family is active
    st_selected = gr.State([])   # player ids selected by clicking
    st_placed = gr.State(None)   # placed meshes after build
    st_attack = gr.State(+1)
    st_defenders = gr.State([])
    st_masks = gr.State({})      # per-player non-arm vertex masks

    video = gr.Video(label="1. Upload match clip")
    status = gr.Markdown()

    with gr.Row():
        frame_slider = gr.Slider(0, 1, value=0, step=1,
                                 label="2. Scrub to the offside frame")
    with gr.Row():
        prev_btn = gr.Button("◀ prev frame")
        next_btn = gr.Button("next frame ▶")

    # Stage (b): lines drawn on the scrubbed frame — auto-proposed or by hand
    frame_view = gr.Image(label="3. Goal-parallel lines: auto-detect or click 4 points",
                          interactive=True)
    with gr.Row():
        auto_lines_btn = gr.Button("✨ Auto-detect lines")
        flip_lines_btn = gr.Button("↔ Flip line direction")
    line_status = gr.Markdown()

    # Stage (c)/(d): detect (RF-DETR-Seg), then click to select
    with gr.Row():
        thr = gr.Slider(0.0, 0.95, value=0.4, step=0.05, label="Detection confidence")
        style = gr.Radio(["Bounding boxes", "Segments"], value="Bounding boxes",
                         label="Detection style")
        detect_btn = gr.Button("4. Detect players (GPU)", variant="primary")
    detect_view = gr.Image(label="5. Click a player to select (click again to deselect)",
                           interactive=True)
    select_status = gr.Markdown()

    # Stage (e)/(f)
    defenders = gr.CheckboxGroup(choices=[], label="6. Defenders (incl. GK) — sets the offside line")
    with gr.Row():
        flip = gr.Checkbox(False, label="flip up (if players are upside-down)")
        attack = gr.Radio(["toward +X", "toward -X"], value="toward +X",
                          label="Attacking direction")
    build_btn = gr.Button("7. Build 3D scene + offside line", variant="primary")

    scene = gr.Plot(label="3D scene")
    plane_slider = gr.Slider(-10, 10, value=0, step=0.05, visible=False,
                             label="Drag the offside plane (X, m)")
    build_status = gr.Markdown()

    # --- wiring ---
    video.change(on_upload, [video],
                 [frame_slider, frame_view, status, st_nmax, st_frame, st_lines,
                  line_status, st_families, st_fam_idx])
    frame_slider.change(on_scrub, [video, frame_slider],
                        [frame_view, st_frame, st_lines, line_status,
                         st_families, st_fam_idx])
    prev_btn.click(lambda i, m: step_frame(i, -1, m), [frame_slider, st_nmax], [frame_slider])
    next_btn.click(lambda i, m: step_frame(i, +1, m), [frame_slider, st_nmax], [frame_slider])

    frame_view.select(on_line_click, [st_lines, st_frame],
                      [st_lines, frame_view, line_status, st_families, st_fam_idx])
    auto_lines_btn.click(on_auto_lines, [st_frame],
                         [st_lines, frame_view, line_status, st_families, st_fam_idx])
    flip_lines_btn.click(on_flip_lines, [st_families, st_fam_idx, st_frame],
                         [st_lines, frame_view, line_status, st_fam_idx])

    detect_btn.click(on_detect, [video, frame_slider, thr, style],
                     [detect_view, st_people, st_selected, select_status, defenders])
    detect_view.select(on_select_player,
                       [st_people, st_selected, st_frame, defenders, style],
                       [detect_view, st_selected, select_status, defenders])
    style.change(on_restyle, [st_people, st_selected, st_frame, style], [detect_view])

    build_btn.click(
        on_build,
        [video, frame_slider, st_people, st_selected, st_lines, flip, attack, defenders],
        [scene, build_status, plane_slider, st_placed, st_attack, st_defenders, st_masks])
    plane_slider.change(on_plane,
                        [st_placed, plane_slider, st_attack, st_defenders, st_masks],
                        [scene])


if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)