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

# detector radio choice -> (backend, overlay/selection style)
DETECTORS = {
    "ViTDet (boxes)": ("vitdet", "boxes"),
    "RF-DETR (boxes)": ("rfdetr", "boxes"),
    "RF-DETR (segments)": ("rfdetr", "segments"),
}


def _resolve(detector):
    return DETECTORS.get(detector, ("vitdet", "boxes"))


def _render(frame, people, selected, style):
    if style == "segments":
        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 == "segments" 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, [], "", [], 0)
    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, detector):
    """Detect step: chosen detector (ViTDet / RF-DETR), boxes + masks. No meshes here."""
    from pipeline.gpu import detect_frame
    if not video_path:
        return None, [], [], "Upload a clip first.", gr.update(choices=[], value=[])
    backend, style = _resolve(detector)
    people = detect_frame(video_path, idx, conf, backend)
    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 == "segments" else "box"
    msg = (f"Detected {len(people)} players with {detector}. Click a player's "
           f"{unit} to select (click again to deselect).")
    return annotated, people, [], msg, gr.update(choices=[], value=[])


# ============================================================================
# Stage (d): click players to select  (CPU)
# ============================================================================
def on_select_player(people, selected, clean_frame, cur_def, detector, evt: gr.SelectData):
    """Toggle the clicked player; refresh highlight + defender choices (box or segment)."""
    if not people:
        return None, selected or [], "Detect players first.", gr.update()
    _, style = _resolve(detector)
    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)


def on_detector_change(people, selected, clean_frame, detector):
    """Re-render current detections in the new style; hint to re-Detect on backend swap."""
    if not people:
        return gr.update(), "Detector set — click Detect players to apply."
    _, style = _resolve(detector)
    return _render(clean_frame, people, selected or [], style), \
        "Re-rendered. For a different backend, click Detect players to re-run."


# ============================================================================
# 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 str(attack_dir).startswith("+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)


def on_gen3js(placed, plane_x, attack_sign, defender_ids, masks):
    """Generate a clean three.js view of the current scene (uses the current plane)."""
    from pipeline.gpu import get_faces
    from pipeline import threed
    if not placed:
        return "<p style='color:#9a8bd0'>Build a 3D scene first, then generate.</p>"
    return threed.scene_html(placed, get_faces(), float(plane_x),
                             int(attack_sign), defender_ids or [], masks)


# ============================================================================
# UI
# ============================================================================
# Roboflow-flavored theme (violet primary ≈ Roboflow purple #7C3AED) + light CSS.
RF_PURPLE = "#7C3AED"
THEME = gr.themes.Soft(primary_hue=gr.themes.colors.violet,
                       neutral_hue=gr.themes.colors.slate,
                       font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"])
RF_CSS = """
.gradio-container {max-width: 1180px !important}
#rf-header {background: #7C3AED; color: #fff; padding: 18px 22px; border-radius: 14px;
            margin-bottom: 6px}
#rf-header h2 {color: #fff !important; margin: 0 0 4px 0}
#rf-header p {color: #EDE7FF !important; margin: 0; font-size: 0.92rem}
.gr-button-primary, button.primary {background: #7C3AED !important; border-color: #7C3AED !important}
"""

with gr.Blocks(title="VAR Offside Visualizer") as demo:
    gr.HTML(
        "<div id='rf-header'><h2>VAR-style Offside Visualizer</h2>"
        "<p>Upload → scrub → click 2 goal-parallel lines → Detect → "
        "click players to select → mark defenders → Build</p></div>"
    )
    gr.Markdown(
        "_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, then click to select
    with gr.Row():
        thr = gr.Slider(0.0, 0.95, value=0.3, step=0.05, label="Detection confidence")
        detector = gr.Radio(list(DETECTORS.keys()), value="ViTDet (boxes)",
                            label="Detector")
        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(["−X  ←", "+X  →"], value="−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()

    gen3js_btn = gr.Button("🎥 Generate clean 3D scene (three.js)")
    scene3js = gr.HTML()

    # --- 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, detector],
                     [detect_view, st_people, st_selected, select_status, defenders])
    detect_view.select(on_select_player,
                       [st_people, st_selected, st_frame, defenders, detector],
                       [detect_view, st_selected, select_status, defenders])
    detector.change(on_detector_change, [st_people, st_selected, st_frame, detector],
                    [detect_view, select_status])

    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])
    gen3js_btn.click(on_gen3js,
                     [st_placed, plane_slider, st_attack, st_defenders, st_masks],
                     [scene3js])


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