Update app.py

app.py

@@ -179,6 +179,22 @@ def get_sam3_model_and_processor(model_id="facebook/sam3", hf_token=None):
             "and provided a valid access token in the text box."
         )
 
+def is_different_enough(img1, img2, threshold=0.92):
+    """Detects if two images are visually different enough based on HSV histogram correlation."""
+    # Convert PIL Images to HSV numpy arrays
+    arr1 = cv2.cvtColor(np.array(img1), cv2.COLOR_RGB2HSV)
+    arr2 = cv2.cvtColor(np.array(img2), cv2.COLOR_RGB2HSV)
+    
+    # Calculate histogram
+    hist1 = cv2.calcHist([arr1], [0, 1], None, [30, 32], [0, 180, 0, 256])
+    hist2 = cv2.calcHist([arr2], [0, 1], None, [30, 32], [0, 180, 0, 256])
+    
+    cv2.normalize(hist1, hist1, 0, 1, cv2.NORM_MINMAX)
+    cv2.normalize(hist2, hist2, 0, 1, cv2.NORM_MINMAX)
+    
+    corr = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)
+    return corr < threshold
+
 def analyze_video_sam3(video_path, text_prompts, hf_token, max_duration=8.0, progress=gr.Progress()):
     """
     Core callback function:
@@ -215,14 +231,14 @@ def analyze_video_sam3(video_path, text_prompts, hf_token, max_duration=8.0, pro
     
     orig_fps = orig_fps if orig_fps and orig_fps > 0 else 15.0
     
-    # Subsample FPS to run much faster on CPU (target around 5 FPS)
-    target_fps = 5.0
+    # Target 8 FPS for smooth output rendering
+    target_fps = 8.0
     frame_skip = max(1, int(orig_fps / target_fps))
     fps = orig_fps / frame_skip
     max_frames = int(max_duration * fps)
     
     # Resize resolution to speed up CPU inference dramatically
-    max_dim = 360
+    max_dim = 320
     scale = 1.0
     if max(height, width) > max_dim:
         scale = max_dim / max(height, width)
@@ -246,14 +262,33 @@ def analyze_video_sam3(video_path, text_prompts, hf_token, max_duration=8.0, pro
     if not frames:
         return None, "# ⚠️ Error\nNo frames could be loaded from the video."
         
+    # Extract motion-based keyframes to minimize model run times
+    progress(0.12, desc="Extracting keyframes...")
+    keyframes = [frames[0]]
+    keyframe_indices = [0]
+    max_gap = 6  # force a keyframe at least every 6 frames to keep temporal coherence
+    min_gap = 2  # wait at least 2 frames between keyframes to save time
+    
+    for i in range(1, len(frames)):
+        gap = i - keyframe_indices[-1]
+        if gap >= max_gap:
+            keyframes.append(frames[i])
+            keyframe_indices.append(i)
+        elif gap >= min_gap:
+            if is_different_enough(frames[i], keyframes[-1], threshold=0.94):
+                keyframes.append(frames[i])
+                keyframe_indices.append(i)
+                
+    print(f"Total frames: {len(frames)}, Extracted keyframes: {len(keyframes)}")
+        
     progress(0.15, desc="Initializing SAM 3 video session...")
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float32 if device == "cpu" else torch.bfloat16
     
-    # Initialize video session
+    # Initialize video session with keyframes only
     session = processor.init_video_session(
-        video=frames,
+        video=keyframes,
         inference_device=device,
         inference_state_device=device,
         processing_device=device,
@@ -270,11 +305,11 @@ def analyze_video_sam3(video_path, text_prompts, hf_token, max_duration=8.0, pro
     obj_to_prompt = {}
     track_db = defaultdict(set)
     
-    progress(0.20, desc="Running SAM 3 propagation across frames...")
+    progress(0.20, desc="Running SAM 3 propagation across keyframes...")
     
     session = to_device_recursive(session, device)
     
-    frame_count = len(frames)
+    frame_count = len(keyframes)
     with torch.no_grad():
         for model_outputs in model.propagate_in_video_iterator(
             inference_session=session,
@@ -282,7 +317,7 @@ def analyze_video_sam3(video_path, text_prompts, hf_token, max_duration=8.0, pro
             max_frame_num_to_track=frame_count,
         ):
             processed_outputs = processor.postprocess_outputs(session, model_outputs)
-            frame_idx = model_outputs.frame_idx
+            k_frame_idx = model_outputs.frame_idx  # index in keyframes list
             
             object_ids = processed_outputs["object_ids"]
             masks = processed_outputs["masks"]
@@ -305,9 +340,9 @@ def analyze_video_sam3(video_path, text_prompts, hf_token, max_duration=8.0, pro
                     if mask_2d.ndim == 3:
                         mask_2d = mask_2d.squeeze()
                     mask_2d = (mask_2d > 0.0).astype(np.uint8)
-                    masks_by_frame[frame_idx][current_obj_id] = mask_2d
+                    masks_by_frame[k_frame_idx][current_obj_id] = mask_2d
                     
-            progress(0.20 + 0.60 * ((frame_idx + 1) / frame_count), desc=f"Propagating: Frame {frame_idx + 1}/{frame_count}...")
+            progress(0.20 + 0.60 * ((k_frame_idx + 1) / frame_count), desc=f"Propagating: Keyframe {k_frame_idx + 1}/{frame_count}...")
             
     session = to_device_recursive(session, "cpu")
     
@@ -322,7 +357,9 @@ def analyze_video_sam3(video_path, text_prompts, hf_token, max_duration=8.0, pro
         font = ImageFont.load_default()
         
     for frame_idx, frame in enumerate(frames):
-        frame_masks = masks_by_frame.get(frame_idx, {})
+        # Find the closest keyframe index
+        closest_k = min(range(len(keyframe_indices)), key=lambda k: abs(keyframe_indices[k] - frame_idx))
+        frame_masks = masks_by_frame.get(closest_k, {})
         # Overlay masks on frame
         annotated_img = overlay_masks_on_frame(frame, frame_masks, color_by_obj, alpha=0.55)