Mirko Trasciatti commited on
Commit
da567fd
·
1 Parent(s): bd71d20

Fix SAM2 window: use ALL frames when no kick detected

Browse files

BUG: When kick_frame was None, SAM2 still only analyzed ~110 frames (4s)
instead of the full video.

FIX in _compute_sam_window_from_kick():
- If kick_frame is None → window = (0, total_frames) = ALL frames
- If kick_frame exists → window = 4 seconds around kick (unchanged)

FIX in _apply_selected_ball_to_yolo_state():
- Always call _compute_sam_window_from_kick(), even when kick is None
- Updated status message to indicate 'SAM2 will analyze ALL frames'

Added logging to show which window is being used.

Files changed (1) hide show
  1. app.py +19 -11
app.py CHANGED
@@ -288,16 +288,23 @@ def _compute_sam_window_from_kick(state: AppState, kick_frame: int | None) -> tu
288
  total_frames = state.num_frames
289
  if total_frames == 0:
290
  return 0, 0
291
- fps = state.video_fps if state.video_fps and state.video_fps > 0 else 25.0
292
- target_window_frames = max(1, int(round(fps * 4.0)))
293
- half_window = target_window_frames // 2
294
  if kick_frame is None:
295
  start_idx = 0
 
 
296
  else:
 
 
 
 
297
  start_idx = max(0, int(kick_frame) - half_window)
298
- end_idx = min(total_frames, start_idx + target_window_frames)
299
- if end_idx <= start_idx:
300
- end_idx = min(total_frames, start_idx + 1)
 
 
301
  state.sam_window = (start_idx, end_idx)
302
  return start_idx, end_idx
303
 
@@ -1120,11 +1127,12 @@ def _apply_selected_ball_to_yolo_state(state: AppState) -> None:
1120
  state.yolo_distance_from_start = distance_dict
1121
  state.yolo_kick_distance = [distance_dict.get(f, 0.0) for f in frames_ordered]
1122
 
1123
- # Update kick frame
1124
  kick_frame = tracking.get("kick_frame")
1125
- if kick_frame is not None:
1126
- state.kick_frame = kick_frame
1127
- _compute_sam_window_from_kick(state, kick_frame)
 
1128
 
1129
  # Mark as tracked
1130
  state.is_yolo_tracked = True
@@ -1134,7 +1142,7 @@ def _apply_selected_ball_to_yolo_state(state: AppState) -> None:
1134
  if kick_frame is not None:
1135
  state.yolo_status = f"✅ Ball {idx+1} tracked. Kick @ frame {kick_frame}."
1136
  else:
1137
- state.yolo_status = f"⚠️ Ball {idx+1} tracked ({coverage:.0%} coverage) but no kick detected."
1138
 
1139
 
1140
  def draw_yolo_detections_on_frame(
 
288
  total_frames = state.num_frames
289
  if total_frames == 0:
290
  return 0, 0
291
+
292
+ # If no kick detected, use ALL frames
 
293
  if kick_frame is None:
294
  start_idx = 0
295
+ end_idx = total_frames
296
+ print(f"[_compute_sam_window_from_kick] No kick detected → using ALL {total_frames} frames")
297
  else:
298
+ # If kick detected, use 4-second window around kick
299
+ fps = state.video_fps if state.video_fps and state.video_fps > 0 else 25.0
300
+ target_window_frames = max(1, int(round(fps * 4.0)))
301
+ half_window = target_window_frames // 2
302
  start_idx = max(0, int(kick_frame) - half_window)
303
+ end_idx = min(total_frames, start_idx + target_window_frames)
304
+ if end_idx <= start_idx:
305
+ end_idx = min(total_frames, start_idx + 1)
306
+ print(f"[_compute_sam_window_from_kick] Kick @ {kick_frame} → window [{start_idx}, {end_idx}] ({end_idx - start_idx} frames)")
307
+
308
  state.sam_window = (start_idx, end_idx)
309
  return start_idx, end_idx
310
 
 
1127
  state.yolo_distance_from_start = distance_dict
1128
  state.yolo_kick_distance = [distance_dict.get(f, 0.0) for f in frames_ordered]
1129
 
1130
+ # Update kick frame and SAM window
1131
  kick_frame = tracking.get("kick_frame")
1132
+ state.kick_frame = kick_frame # Can be None
1133
+
1134
+ # Always compute SAM window - if no kick, it will use ALL frames
1135
+ _compute_sam_window_from_kick(state, kick_frame)
1136
 
1137
  # Mark as tracked
1138
  state.is_yolo_tracked = True
 
1142
  if kick_frame is not None:
1143
  state.yolo_status = f"✅ Ball {idx+1} tracked. Kick @ frame {kick_frame}."
1144
  else:
1145
+ state.yolo_status = f"⚠️ Ball {idx+1} tracked ({coverage:.0%} coverage) but no kick detected. SAM2 will analyze ALL frames."
1146
 
1147
 
1148
  def draw_yolo_detections_on_frame(