Refine ghost trail to be anticipatory

app.py

@@ -893,7 +893,8 @@ def compose_frame(state: AppState, frame_idx: int, remove_bg: bool = False) -> I
     frame = state.video_frames[frame_idx]
     masks = state.masks_by_frame.get(frame_idx, {})
     out_img = frame
-    
+    ghost_mask = _build_ball_trail_mask(state, frame_idx)
+
     if len(masks) != 0:
         if remove_bg:
             # Remove background - show only tracked objects
@@ -910,20 +911,19 @@ def compose_frame(state: AppState, frame_idx: int, remove_bg: bool = False) -> I
             out_img = Image.fromarray(result_np)
         else:
             overlay_masks = masks
-            if state.fx_ghost_trail_enabled:
+            if ghost_mask is not None:
                 overlay_masks = {oid: mask for oid, mask in masks.items() if oid != BALL_OBJECT_ID}
             if overlay_masks:
                 out_img = overlay_masks_on_frame(out_img, overlay_masks, state.color_by_obj, alpha=0.65)
-            if state.fx_ghost_trail_enabled:
-                trail_mask = _build_ball_trail_mask(state, frame_idx)
-                if trail_mask is not None:
-                    base_np = np.array(out_img).astype(np.float32) / 255.0
-                    mask_np = np.clip(trail_mask.astype(np.float32), 0.0, 1.0)
-                    mask_np = mask_np[..., None]
-                    base_np = (1.0 - GHOST_TRAIL_ALPHA * mask_np) * base_np + (
-                        GHOST_TRAIL_ALPHA * mask_np
-                    ) * GHOST_TRAIL_COLOR
-                    out_img = Image.fromarray(np.clip(base_np * 255.0, 0, 255).astype(np.uint8))
+
+    if ghost_mask is not None:
+        base_np = np.array(out_img).astype(np.float32) / 255.0
+        mask_np = np.clip(ghost_mask.astype(np.float32), 0.0, 1.0)
+        if mask_np.ndim == 3:
+            mask_np = mask_np.squeeze()
+        mask_np = mask_np[..., None]
+        base_np = (1.0 - GHOST_TRAIL_ALPHA * mask_np) * base_np + (GHOST_TRAIL_ALPHA * mask_np) * GHOST_TRAIL_COLOR
+        out_img = Image.fromarray(np.clip(base_np * 255.0, 0, 255).astype(np.uint8))
 
     # Draw crosses for conditioning frames only (frames with recorded clicks)
     clicks_map = state.clicks_by_frame_obj.get(frame_idx)
@@ -1064,11 +1064,13 @@ def _build_ball_trail_mask(state: AppState, frame_idx: int) -> np.ndarray | None
     ):
         return None
     kick_candidate = state.kick_frame if state.kick_frame is not None else state.kick_debug_kick_frame
-    if kick_candidate is None or frame_idx <= kick_candidate + 1:
+    if kick_candidate is None:
         return None
 
-    start_idx = max(int(kick_candidate) + 1, 0)
-    end_idx = int(frame_idx)
+    start_idx = max(int(kick_candidate) + 1, int(frame_idx) + 1)
+    end_idx = state.num_frames
+    if start_idx >= end_idx:
+        return None
     trail_mask: np.ndarray | None = None
 
     for idx in range(start_idx, end_idx):