Disable propagate buttons until detections

app.py

@@ -47,6 +47,7 @@ YOLO_CONF_THRESHOLD = 0.0
 YOLO_IOU_THRESHOLD = 0.02
 PLAYER_TARGET_NAME = "person"
 PLAYER_OBJECT_ID = 2
+BALL_OBJECT_ID = 1
 
 
 def get_yolo_model(model_filename: str = YOLO_DEFAULT_MODEL) -> YOLO:
@@ -997,6 +998,15 @@ def _format_kick_status(state: AppState) -> str:
     return f"Kick frame ≈ {frame}{time_part}"
 
 
+def _ball_has_masks(state: AppState, target_obj_id: int = BALL_OBJECT_ID) -> bool:
+    if state is None:
+        return False
+    for masks in state.masks_by_frame.values():
+        if target_obj_id in masks:
+            return True
+    return False
+
+
 def _player_has_masks(state: AppState) -> bool:
     if state is None or state.player_obj_id is None:
         return False
@@ -1007,17 +1017,20 @@ def _player_has_masks(state: AppState) -> bool:
     return False
 
 
-def _player_button_updates(state: AppState) -> tuple[Any, Any]:
-    detect_enabled = False
-    propagate_enabled = False
+def _button_updates(state: AppState) -> tuple[Any, Any, Any]:
+    propagate_main_enabled = _ball_has_masks(state)
+    detect_player_enabled = False
+    propagate_player_enabled = False
     if isinstance(state, AppState):
         kick_candidate = state.kick_frame or getattr(state, "kick_debug_kick_frame", None)
         if kick_candidate is not None:
-            detect_enabled = True
-            propagate_enabled = _player_has_masks(state)
+            detect_player_enabled = True
+        if detect_player_enabled and _player_has_masks(state):
+            propagate_player_enabled = True
     return (
-        gr.update(interactive=detect_enabled),
-        gr.update(interactive=propagate_enabled),
+        gr.update(interactive=propagate_main_enabled),
+        gr.update(interactive=detect_player_enabled),
+        gr.update(interactive=propagate_player_enabled),
     )
 
 
@@ -1431,11 +1444,25 @@ def on_image_click(
     return update_frame_display(state, int(frame_idx))
 
 
+def _on_image_click_with_updates(
+    img: Image.Image | np.ndarray,
+    state: AppState,
+    frame_idx: int,
+    obj_id: int,
+    label: str,
+    clear_old: bool,
+    evt: gr.SelectData,
+):
+    preview_img = on_image_click(img, state, frame_idx, obj_id, label, clear_old, evt)
+    propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(state)
+    return preview_img, propagate_main_update, detect_btn_update, propagate_player_update
+
+
 @spaces.GPU()
 def propagate_masks(GLOBAL_STATE: gr.State):
     if GLOBAL_STATE is None or GLOBAL_STATE.inference_session is None:
         # yield GLOBAL_STATE, "Load a video first.", gr.update()
-        detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+        propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
         return (
             GLOBAL_STATE,
             "Load a video first.",
@@ -1443,6 +1470,7 @@ def propagate_masks(GLOBAL_STATE: gr.State):
             _build_kick_plot(GLOBAL_STATE),
             _format_impact_status(GLOBAL_STATE),
             gr.update(value=_format_kick_status(GLOBAL_STATE), visible=True),
+            propagate_main_update,
             detect_btn_update,
             propagate_player_update,
         )
@@ -1459,7 +1487,7 @@ def propagate_masks(GLOBAL_STATE: gr.State):
     processed = 0
 
     # Initial status; no slider change yet
-    detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+    propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
     yield (
         GLOBAL_STATE,
         f"Propagating masks: {processed}/{total}",
@@ -1467,6 +1495,7 @@ def propagate_masks(GLOBAL_STATE: gr.State):
         _build_kick_plot(GLOBAL_STATE),
         _format_impact_status(GLOBAL_STATE),
         gr.update(value=_format_kick_status(GLOBAL_STATE), visible=True),
+        propagate_main_update,
         detect_btn_update,
         propagate_player_update,
     )
@@ -1496,7 +1525,7 @@ def propagate_masks(GLOBAL_STATE: gr.State):
             processed += 1
             # Every 15th frame (or last), move slider to current frame to update preview via slider binding
             if processed % 30 == 0 or processed == total:
-                detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+                propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
                 yield (
                     GLOBAL_STATE,
                     f"Propagating masks: {processed}/{total}",
@@ -1504,6 +1533,7 @@ def propagate_masks(GLOBAL_STATE: gr.State):
                     _build_kick_plot(GLOBAL_STATE),
                     _format_impact_status(GLOBAL_STATE),
                     gr.update(value=_format_kick_status(GLOBAL_STATE), visible=True),
+                    propagate_main_update,
                     detect_btn_update,
                     propagate_player_update,
                 )
@@ -1518,7 +1548,7 @@ def propagate_masks(GLOBAL_STATE: gr.State):
     GLOBAL_STATE.current_frame_idx = target_frame
 
     # Final status; ensure slider points to the target frame (kick frame when detected)
-    detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+    propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
     yield (
         GLOBAL_STATE,
         text,
@@ -1526,16 +1556,17 @@ def propagate_masks(GLOBAL_STATE: gr.State):
         _build_kick_plot(GLOBAL_STATE),
         _format_impact_status(GLOBAL_STATE),
         gr.update(value=_format_kick_status(GLOBAL_STATE), visible=True),
+        propagate_main_update,
         detect_btn_update,
         propagate_player_update,
     )
 
 
-def reset_session(GLOBAL_STATE: gr.State) -> tuple[AppState, Image.Image, int, int, str, any, go.Figure, Any, Any]:
+def reset_session(GLOBAL_STATE: gr.State) -> tuple[AppState, Image.Image, int, int, str, any, go.Figure, Any, Any, Any]:
     # Reset only session-related state, keep uploaded video and model
     if not GLOBAL_STATE.video_frames:
         # Nothing loaded; keep behavior
-        detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+        propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
         return (
             GLOBAL_STATE,
             None,
@@ -1545,6 +1576,7 @@ def reset_session(GLOBAL_STATE: gr.State) -> tuple[AppState, Image.Image, int, i
             gr.update(visible=False, value=""),
             _build_kick_plot(GLOBAL_STATE),
             _format_impact_status(GLOBAL_STATE),
+            propagate_main_update,
             detect_btn_update,
             propagate_player_update,
         )
@@ -1604,7 +1636,7 @@ def reset_session(GLOBAL_STATE: gr.State) -> tuple[AppState, Image.Image, int, i
     slider_minmax = gr.update(minimum=0, maximum=max(GLOBAL_STATE.num_frames - 1, 0), interactive=True)
     slider_value = gr.update(value=current_idx)
     status = "Session reset. Prompts cleared; video preserved."
-    detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+    propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
     # clear and reload model and processor
     return (
         GLOBAL_STATE,
@@ -1615,6 +1647,7 @@ def reset_session(GLOBAL_STATE: gr.State) -> tuple[AppState, Image.Image, int, i
         gr.update(visible=False, value=""),
         _build_kick_plot(GLOBAL_STATE),
         _format_impact_status(GLOBAL_STATE),
+        propagate_main_update,
         detect_btn_update,
         propagate_player_update,
     )
@@ -1918,7 +1951,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
                 )
             with gr.Row():
                 detect_ball_btn = gr.Button("Detect Ball", variant="secondary")
-                propagate_btn = gr.Button("Propagate across video", variant="primary")
+                propagate_btn = gr.Button("Propagate across video", variant="primary", interactive=False)
             detect_player_btn = gr.Button("Detect Player", variant="secondary", interactive=False)
             propagate_player_btn = gr.Button("Propagate Player", variant="primary", interactive=False)
             ball_status = gr.Markdown(visible=False)
@@ -1934,7 +1967,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
     # Wire events
     def _on_video_change(GLOBAL_STATE: gr.State, video):
         GLOBAL_STATE, min_idx, max_idx, first_frame, status = init_video_session(GLOBAL_STATE, video)
-        detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+        propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
         return (
             GLOBAL_STATE,
             gr.update(minimum=min_idx, maximum=max_idx, value=min_idx, interactive=True),
@@ -1943,6 +1976,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
             gr.update(visible=False, value=""),
             _build_kick_plot(GLOBAL_STATE),
             _format_impact_status(GLOBAL_STATE),
+            propagate_main_update,
             detect_btn_update,
             propagate_player_update,
         )
@@ -1950,7 +1984,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
     video_in.change(
         _on_video_change,
         inputs=[GLOBAL_STATE, video_in],
-        outputs=[GLOBAL_STATE, frame_slider, preview, load_status, ball_status, kick_plot, impact_status, detect_player_btn, propagate_player_btn],
+        outputs=[GLOBAL_STATE, frame_slider, preview, load_status, ball_status, kick_plot, impact_status, propagate_btn, detect_player_btn, propagate_player_btn],
         show_progress=True,
     )
 
@@ -1963,7 +1997,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
             examples=examples_list,
             inputs=[GLOBAL_STATE, video_in],
             fn=_on_video_change,
-            outputs=[GLOBAL_STATE, frame_slider, preview, load_status, ball_status, kick_plot, impact_status, detect_player_btn, propagate_player_btn],
+            outputs=[GLOBAL_STATE, frame_slider, preview, load_status, ball_status, kick_plot, impact_status, propagate_btn, detect_player_btn, propagate_player_btn],
             label="Examples",
             cache_examples=False,
             examples_per_page=5,
@@ -2046,11 +2080,12 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
         if s is not None and val is not None:
             s.min_impact_speed_kmh = float(val)
             _recompute_motion_metrics(s)
-        detect_btn_update, propagate_player_update = _player_button_updates(s)
+        propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(s)
         return (
             _build_kick_plot(s),
             _format_impact_status(s),
             gr.update(value=_format_kick_status(s), visible=True),
+            propagate_main_update,
             detect_btn_update,
             propagate_player_update,
         )
@@ -2059,11 +2094,12 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
         if s is not None and val is not None:
             s.goal_distance_m = float(val)
             _recompute_motion_metrics(s)
-        detect_btn_update, propagate_player_update = _player_button_updates(s)
+        propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(s)
         return (
             _build_kick_plot(s),
             _format_impact_status(s),
             gr.update(value=_format_kick_status(s), visible=True),
+            propagate_main_update,
             detect_btn_update,
             propagate_player_update,
         )
@@ -2071,13 +2107,13 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
     min_impact_speed_slider.change(
         _update_min_impact_speed,
         inputs=[GLOBAL_STATE, min_impact_speed_slider],
-        outputs=[kick_plot, impact_status, ball_status, detect_player_btn, propagate_player_btn],
+        outputs=[kick_plot, impact_status, ball_status, propagate_btn, detect_player_btn, propagate_player_btn],
     )
 
     goal_distance_slider.change(
         _update_goal_distance,
         inputs=[GLOBAL_STATE, goal_distance_slider],
-        outputs=[kick_plot, impact_status, ball_status, detect_player_btn, propagate_player_btn],
+        outputs=[kick_plot, impact_status, ball_status, propagate_btn, detect_player_btn, propagate_player_btn],
     )
 
     def _auto_detect_ball(
@@ -2093,7 +2129,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
         frame = state_in.video_frames[frame_idx]
         detection = detect_ball_center(frame)
         if detection is None:
-            detect_btn_update, propagate_player_update = _player_button_updates(state_in)
+            propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(state_in)
             return (
                 update_frame_display(state_in, frame_idx),
                 gr.update(
@@ -2102,6 +2138,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
                 ),
                 gr.update(value=frame_idx),
                 _build_kick_plot(state_in),
+                propagate_main_update,
                 detect_btn_update,
                 propagate_player_update,
             )
@@ -2133,12 +2170,13 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
 
         status_text = f"✅ Auto-detected ball at ({x_center}, {y_center}) (conf={conf:.2f})"
         status_text += f" | {_format_kick_status(state_in)}"
-        detect_btn_update, propagate_player_update = _player_button_updates(state_in)
+        propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(state_in)
         return (
             preview_img,
             gr.update(value=status_text, visible=True),
             gr.update(value=frame_idx),
             _build_kick_plot(state_in),
+            propagate_main_update,
             detect_btn_update,
             propagate_player_update,
         )
@@ -2146,7 +2184,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
     detect_ball_btn.click(
         _auto_detect_ball,
         inputs=[GLOBAL_STATE, obj_id_inp, label_radio, clear_old_chk],
-        outputs=[preview, ball_status, frame_slider, kick_plot, detect_player_btn, propagate_player_btn],
+        outputs=[preview, ball_status, frame_slider, kick_plot, propagate_btn, detect_player_btn, propagate_player_btn],
     )
 
     def _auto_detect_player(state_in: AppState):
@@ -2163,7 +2201,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
         frame = state_in.video_frames[frame_idx]
         detection = detect_person_box(frame)
         if detection is None:
-            detect_btn_update, propagate_player_update = _player_button_updates(state_in)
+            propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(state_in)
             status_text = (
                 f"{_format_kick_status(state_in)} | ⚠️ Unable to auto-detect the player on frame {frame_idx}. "
                 "Please add a box manually."
@@ -2173,6 +2211,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
                 gr.update(value=status_text, visible=True),
                 gr.update(value=frame_idx),
                 _build_kick_plot(state_in),
+                propagate_main_update,
                 detect_btn_update,
                 propagate_player_update,
                 gr.update(),
@@ -2234,7 +2273,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
         state_in.composited_frames.pop(frame_idx, None)
         state_in.current_frame_idx = frame_idx
 
-        detect_btn_update, propagate_player_update = _player_button_updates(state_in)
+        propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(state_in)
         status_text = (
             f"{_format_kick_status(state_in)} | ✅ Player auto-detected on frame {frame_idx} (conf={conf:.2f})"
         )
@@ -2243,6 +2282,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
             gr.update(value=status_text, visible=True),
             gr.update(value=frame_idx),
             _build_kick_plot(state_in),
+            propagate_main_update,
             detect_btn_update,
             propagate_player_update,
             gr.update(value=PLAYER_OBJECT_ID),
@@ -2251,13 +2291,13 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
     detect_player_btn.click(
         _auto_detect_player,
         inputs=[GLOBAL_STATE],
-        outputs=[preview, ball_status, frame_slider, kick_plot, detect_player_btn, propagate_player_btn, obj_id_inp],
+        outputs=[preview, ball_status, frame_slider, kick_plot, propagate_btn, detect_player_btn, propagate_player_btn, obj_id_inp],
     )
 
     @spaces.GPU()
     def propagate_player_masks(GLOBAL_STATE: gr.State):
         if GLOBAL_STATE is None or GLOBAL_STATE.inference_session is None:
-            detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+            propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
             return (
                 GLOBAL_STATE,
                 "Load a video first.",
@@ -2265,11 +2305,12 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
                 _build_kick_plot(GLOBAL_STATE),
                 _format_impact_status(GLOBAL_STATE),
                 gr.update(value=_format_kick_status(GLOBAL_STATE), visible=True),
+                propagate_main_update,
                 detect_btn_update,
                 propagate_player_update,
             )
         if GLOBAL_STATE.player_obj_id is None or not _player_has_masks(GLOBAL_STATE):
-            detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+            propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
             return (
                 GLOBAL_STATE,
                 "Detect the player before propagating.",
@@ -2277,6 +2318,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
                 _build_kick_plot(GLOBAL_STATE),
                 _format_impact_status(GLOBAL_STATE),
                 gr.update(value=_format_kick_status(GLOBAL_STATE), visible=True),
+                propagate_main_update,
                 detect_btn_update,
                 propagate_player_update,
             )
@@ -2291,7 +2333,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
         total = max(1, GLOBAL_STATE.num_frames)
         processed = 0
 
-        detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+        propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
         yield (
             GLOBAL_STATE,
             f"Propagating player: {processed}/{total}",
@@ -2299,6 +2341,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
             _build_kick_plot(GLOBAL_STATE),
             _format_impact_status(GLOBAL_STATE),
             gr.update(value=_format_kick_status(GLOBAL_STATE), visible=True),
+            propagate_main_update,
             detect_btn_update,
             propagate_player_update,
         )
@@ -2333,7 +2376,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
 
                 processed += 1
                 if processed % 30 == 0 or processed == total:
-                    detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+                    propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
                     yield (
                         GLOBAL_STATE,
                         f"Propagating player: {processed}/{total}",
@@ -2341,6 +2384,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
                         _build_kick_plot(GLOBAL_STATE),
                         _format_impact_status(GLOBAL_STATE),
                         gr.update(value=_format_kick_status(GLOBAL_STATE), visible=True),
+                        propagate_main_update,
                         detect_btn_update,
                         propagate_player_update,
                     )
@@ -2354,7 +2398,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
         target_frame = int(np.clip(target_frame, 0, max(0, GLOBAL_STATE.num_frames - 1)))
         GLOBAL_STATE.current_frame_idx = target_frame
 
-        detect_btn_update, propagate_player_update = _player_button_updates(GLOBAL_STATE)
+        propagate_main_update, detect_btn_update, propagate_player_update = _button_updates(GLOBAL_STATE)
         yield (
             GLOBAL_STATE,
             text,
@@ -2362,6 +2406,7 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
             _build_kick_plot(GLOBAL_STATE),
             _format_impact_status(GLOBAL_STATE),
             gr.update(value=_format_kick_status(GLOBAL_STATE), visible=True),
+            propagate_main_update,
             detect_btn_update,
             propagate_player_update,
         )
@@ -2369,12 +2414,14 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
     propagate_player_btn.click(
         propagate_player_masks,
         inputs=[GLOBAL_STATE],
-        outputs=[GLOBAL_STATE, propagate_status, frame_slider, kick_plot, impact_status, ball_status, detect_player_btn, propagate_player_btn],
+        outputs=[GLOBAL_STATE, propagate_status, frame_slider, kick_plot, impact_status, ball_status, propagate_btn, detect_player_btn, propagate_player_btn],
     )
 
     # Image click to add a point and run forward on that frame
     preview.select(
-        on_image_click, [preview, GLOBAL_STATE, frame_slider, obj_id_inp, label_radio, clear_old_chk], preview
+        _on_image_click_with_updates,
+        [preview, GLOBAL_STATE, frame_slider, obj_id_inp, label_radio, clear_old_chk],
+        [preview, propagate_btn, detect_player_btn, propagate_player_btn],
     )
 
     # Playback via MP4 rendering only
@@ -2436,13 +2483,13 @@ with gr.Blocks(title="SAM2 Video (Transformers) - Interactive Segmentation", the
     propagate_btn.click(
         propagate_masks,
         inputs=[GLOBAL_STATE],
-        outputs=[GLOBAL_STATE, propagate_status, frame_slider, kick_plot, impact_status, ball_status, detect_player_btn, propagate_player_btn],
+        outputs=[GLOBAL_STATE, propagate_status, frame_slider, kick_plot, impact_status, ball_status, propagate_btn, detect_player_btn, propagate_player_btn],
     )
 
     reset_btn.click(
         reset_session,
         inputs=GLOBAL_STATE,
-        outputs=[GLOBAL_STATE, preview, frame_slider, frame_slider, load_status, ball_status, kick_plot, impact_status, detect_player_btn, propagate_player_btn],
+        outputs=[GLOBAL_STATE, preview, frame_slider, frame_slider, load_status, ball_status, kick_plot, impact_status, propagate_btn, detect_player_btn, propagate_player_btn],
     )
 
 # ============================================================================