Reorder flow: draw lines on scrubbed frame, split Detect, click-to-select players

app.py

@@ -1,13 +1,13 @@
 """
 VAR-style offside visualizer — Gradio app.
 
-Pipeline (matches the working Colab notebook):
-  upload video -> scrub to the offside frame -> detect players (GPU, once/frame)
-  -> select players -> click 2 goal-parallel lines -> place 3D poses on a field
-  -> Plotly scene with a draggable offside plane.
+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.reconstruct_frame(); every other
-callback in this file runs on cached numpy and stays on the CPU.
+callback here runs on cached numpy and stays on the CPU.
 """
 
 import os
@@ -20,7 +20,7 @@ import numpy as np
 import gradio as gr
 
 from pipeline.video import probe_video, grab_frame
-from pipeline.overlay import annotate_detections, draw_lines
+from pipeline.overlay import annotate_detections, draw_lines, pick_box
 from pipeline import geometry as G
 
 
@@ -28,65 +28,87 @@ from pipeline import geometry as G
 # Stage (a): upload + frame scrubbing  (pure CPU)
 # ============================================================================
 def on_upload(video_path):
-    """New clip: size the slider, show frame 0, report length, remember max idx."""
+    """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
+        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),
-        grab_frame(video_path, 0),
-        f"{n} frames @ {fps:.1f} fps — scrub to the moment the ball is played.",
+        frame,
+        f"{n} frames @ {fps:.1f} fps — scrub to the moment the ball is played, "
+        "then click 2 goal-parallel lines on the frame.",
         n_max,
+        frame,        # st_frame: clean copy for line redraws
+        [],           # st_lines reset
+        "",           # line_status reset
     )
 
 
 def on_scrub(video_path, idx):
+    """Show the new frame and reset any half-drawn lines on it."""
     if not video_path:
-        return None
-    return grab_frame(video_path, int(idx))
+        return None, None, [], ""
+    frame = grab_frame(video_path, int(idx))
+    return frame, frame, [], ""
 
 
 def step_frame(idx, delta, n_max):
-    """Clamp prev/next within [0, slider max]."""
     return max(0, min(int(n_max), int(idx) + delta))
 
 
 # ============================================================================
-# Stage (b)+(c): detect (GPU, cached) + populate selection  (CPU after the call)
+# 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)."""
+    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
+
+
+# ============================================================================
+# Stage (c): detect (GPU, cached) + show boxes  (CPU after the call)
 # ============================================================================
 def on_detect(video_path, idx, bbox_thr):
-    """The one GPU step. Imported lazily so the app boots without the model."""
+    """The one GPU step. Lazy-imported so the app boots without the model."""
     from pipeline.gpu import reconstruct_frame
     if not video_path:
-        return None, gr.update(choices=[], value=[]), gr.update(choices=[], value=[]), [], "Upload a clip first."
+        return None, [], [], "Upload a clip first.", gr.update(choices=[], value=[])
     people = reconstruct_frame(video_path, idx, bbox_thr)
     if not people:
-        return (None, gr.update(choices=[], value=[]), gr.update(choices=[], value=[]),
-                [], "No players detected — try a lower threshold.")
-    annotated = annotate_detections(grab_frame(video_path, idx), people)
-    choices = list(range(len(people)))
-    msg = (f"Detected {len(people)} players. Tick the ones to analyze, mark defenders, "
-           "then click 2 goal-parallel lines on the image above.")
-    return (annotated,
-            gr.update(choices=choices, value=[]),
-            gr.update(choices=choices, value=[]),
-            people, msg)
+        return (None, [], [], "No players detected — lower the confidence slider.",
+                gr.update(choices=[], value=[]))
+    annotated = annotate_detections(grab_frame(video_path, idx), people, [])
+    msg = (f"Detected {len(people)} players. Click a box to select a player "
+           "(click again to deselect).")
+    return annotated, people, [], msg, gr.update(choices=[], value=[])
 
 
 # ============================================================================
-# Stage (d): click two goal-parallel lines (4 points)  (CPU)
+# Stage (d): click players to select  (CPU)
 # ============================================================================
-def on_line_click(line_pts, frame_rgb, evt: gr.SelectData):
-    """Collect 4 clicks = 2 lines; redraw dots + segments. Resets after a full set."""
-    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(frame_rgb, pts)
-    status = {1: "line 1: 1/2", 2: "line 1 set", 3: "line 2: 1/2",
-              4: "both lines set — build the scene"}.get(len(pts), "")
-    return pts, img, status
+def on_select_player(people, selected, clean_frame, cur_def, evt: gr.SelectData):
+    """Toggle the player whose box was clicked; refresh highlight + defender choices."""
+    if not people:
+        return None, selected or [], "Detect players first.", gr.update()
+    hit = pick_box(people, evt.index[0], evt.index[1])
+    sel = list(selected or [])
+    if hit is not None:
+        sel.remove(hit) if hit in sel else sel.append(hit)
+    sel = sorted(sel)
+    img = annotate_detections(clean_frame, people, sel)
+    msg = (f"Selected players: {sel}. Mark defenders below, then Build."
+           if sel else "Click a box to select a player.")
+    keep_def = [d for d in (cur_def or []) if d in sel]
+    return img, sel, msg, gr.update(choices=sel, value=keep_def)
 
 
 # ============================================================================
@@ -96,9 +118,10 @@ def on_build(video_path, idx, bbox_thr, selected_ids, line_pts, flip_up,
              attack_dir, defender_ids):
     from pipeline.gpu import reconstruct_frame, get_faces
     if not selected_ids:
-        return None, "Select at least one player.", gr.update(), None, +1, []
+        return None, "Click at least one player to select.", gr.update(), None, +1, []
     if not line_pts or len(line_pts) < 4:
-        return None, "Click 2 goal-parallel lines (4 points) first.", gr.update(), None, +1, []
+        return None, "Draw 2 goal-parallel lines (4 points) on the frame first.", \
+               gr.update(), None, +1, []
 
     people = reconstruct_frame(video_path, idx, bbox_thr)
     faces = get_faces()
@@ -112,7 +135,6 @@ def on_build(video_path, idx, bbox_thr, selected_ids, line_pts, flip_up,
     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)
-
     fig = G.build_scene(placed, faces, plane_x, attack_sign, dset)
 
     allX = np.vstack(list(placed.values()))[:, 0]
@@ -121,8 +143,8 @@ def on_build(video_path, idx, bbox_thr, selected_ids, line_pts, flip_up,
                              visible=True, label="Drag the offside plane (X, m)")
 
     warn = "  ⚠ heights look wrong — toggle 'flip up'." if med < 1.0 else ""
-    msg = f"Median player height {med:.2f} m (expect ~1.7–1.9).{warn}"
-    return fig, msg, plane_update, placed, attack_sign, dset
+    return (fig, f"Median player height {med:.2f} m (expect ~1.7–1.9).{warn}",
+            plane_update, placed, attack_sign, dset)
 
 
 def on_plane(placed, plane_x, attack_sign, defender_ids):
@@ -130,9 +152,8 @@ def on_plane(placed, plane_x, attack_sign, defender_ids):
     from pipeline.gpu import get_faces
     if not placed:
         return gr.update()
-    fig = G.build_scene(placed, get_faces(), float(plane_x),
-                        int(attack_sign), defender_ids or [])
-    return fig
+    return G.build_scene(placed, get_faces(), float(plane_x),
+                         int(attack_sign), defender_ids or [])
 
 
 # ============================================================================
@@ -141,49 +162,52 @@ def on_plane(placed, plane_x, attack_sign, defender_ids):
 with gr.Blocks(title="VAR Offside Visualizer") as demo:
     gr.Markdown(
         "## VAR-style Offside Visualizer\n"
-        "Upload a clip, scrub to the moment the ball is played, then "
-        "detect → select → click 2 goal-parallel lines → build the 3D scene.\n\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_people = gr.State([])     # slim detections for the current frame
+    st_frame = gr.State(None)    # clean RGB of the current frame
+    st_people = gr.State([])     # slim detections
     st_lines = gr.State([])      # clicked line points
-    st_placed = gr.State(None)   # placed meshes (field frame) after build
-    st_attack = gr.State(+1)     # attack_sign
-    st_defenders = gr.State([])  # defender ids used for the verdict
+    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([])
 
-    # --- stage (a) ---
     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 ▶")
-    frame_view = gr.Image(label="Current frame", interactive=False)
 
-    # --- stage (b)/(c) ---
+    # Stage (b): lines are drawn directly on the scrubbed frame
+    frame_view = gr.Image(label="3. Click 2 goal-parallel lines here (4 points)",
+                          interactive=True)
+    line_status = gr.Markdown()
+
+    # Stage (c)/(d): detect, then click boxes to select
     with gr.Row():
         thr = gr.Slider(0.0, 0.95, value=0.85, step=0.05, label="Detection confidence")
-        detect_btn = gr.Button("3. Detect players (GPU)", variant="primary")
-    detect_view = gr.Image(
-        label="Detected players — click 2 goal-parallel lines here (4 points)",
-        interactive=True)
-    line_status = gr.Markdown()
-    sel = gr.CheckboxGroup(choices=[], label="4. Select players to analyze (indices)")
+        detect_btn = gr.Button("4. Detect players (GPU)", variant="primary")
+    detect_view = gr.Image(label="5. Click a player's box to select (click again to deselect)",
+                           interactive=True)
+    select_status = gr.Markdown()
 
-    # --- stage (e)/(f) controls ---
+    # 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")
-    defenders = gr.CheckboxGroup(
-        choices=[], label="Defenders (incl. GK) — sets the offside line")
-    build_btn = gr.Button("5. Build 3D scene + offside line", variant="primary")
+    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,
@@ -191,21 +215,24 @@ with gr.Blocks(title="VAR Offside Visualizer") as demo:
     build_status = gr.Markdown()
 
     # --- wiring ---
-    video.change(on_upload, [video], [frame_slider, frame_view, status, st_nmax])
-    frame_slider.change(on_scrub, [video, frame_slider], [frame_view])
-    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])
+    video.change(on_upload, [video],
+                 [frame_slider, frame_view, status, st_nmax, st_frame, st_lines, line_status])
+    frame_slider.change(on_scrub, [video, frame_slider],
+                        [frame_view, st_frame, st_lines, line_status])
+    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])
 
     detect_btn.click(on_detect, [video, frame_slider, thr],
-                     [detect_view, sel, defenders, st_people, status])
-    detect_view.select(on_line_click, [st_lines, detect_view],
-                       [st_lines, detect_view, line_status])
+                     [detect_view, st_people, st_selected, select_status, defenders])
+    detect_view.select(on_select_player, [st_people, st_selected, st_frame, defenders],
+                       [detect_view, st_selected, select_status, defenders])
 
     build_btn.click(
         on_build,
-        [video, frame_slider, thr, sel, st_lines, flip, attack, defenders],
+        [video, frame_slider, thr, st_selected, st_lines, flip, attack, defenders],
         [scene, build_status, plane_slider, st_placed, st_attack, st_defenders])
     plane_slider.change(on_plane, [st_placed, plane_slider, st_attack, st_defenders],
                         [scene])

pipeline/overlay.py

@@ -7,17 +7,33 @@ import cv2
 import numpy as np
 
 
-def annotate_detections(frame_rgb, people):
-    """Draw numbered green boxes for every detected person."""
+def annotate_detections(frame_rgb, people, selected_ids=()):
+    """Draw numbered boxes; selected players get a thick yellow box, others thin green."""
+    sel = set(int(i) for i in selected_ids)
     img = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
     for i, p in enumerate(people):
         x1, y1, x2, y2 = [int(v) for v in p["bbox"]]
-        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        if i in sel:
+            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 230, 255), 3)   # yellow, thick
+        else:
+            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)     # green, thin
         cv2.putText(img, str(i), (x1, max(y1 - 6, 14)),
                     cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 3)
     return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
 
+def pick_box(people, x, y):
+    """Return the index of the smallest box containing (x, y), or None."""
+    hit, best_area = None, None
+    for i, p in enumerate(people):
+        x1, y1, x2, y2 = p["bbox"]
+        if x1 <= x <= x2 and y1 <= y <= y2:
+            area = (x2 - x1) * (y2 - y1)
+            if best_area is None or area < best_area:
+                best_area, hit = area, i
+    return hit
+
+
 def draw_lines(frame_rgb, pts):
     """Draw the clicked points and the (up to two) line segments.