app update

app.py

@@ -68,34 +68,53 @@ def initialize_models():
     print("All models initialized successfully!")
 
 
-def extract_frames_from_video(video_path, max_frames=50):
+def extract_frames_from_video(video_path, max_frames=24):
     """
     Extract frames from video file
     
     Args:
         video_path: Path to video file
-        max_frames: Maximum number of frames to extract
+        max_frames: Maximum number of frames to extract (default: 24)
         
     Returns:
         frames: List of numpy arrays (H,W,3), uint8 RGB
-        fps: Original FPS of video
+        adjusted_fps: Adjusted FPS for output video to maintain normal playback speed
     """
     cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS)
+    original_fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
-    frames = []
-    while cap.isOpened() and len(frames) < max_frames:
+    # Read all frames first
+    all_frames = []
+    while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
         # Convert BGR to RGB
         frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        frames.append(frame_rgb)
+        all_frames.append(frame_rgb)
     
     cap.release()
-    print(f"Extracted {len(frames)} frames from video (FPS: {fps})")
     
-    return frames, fps
+    # If video has more frames than max_frames, randomly sample
+    if len(all_frames) > max_frames:
+        print(f"Video has {len(all_frames)} frames, randomly sampling {max_frames} frames...")
+        # Sort indices to maintain temporal order
+        sampled_indices = sorted(np.random.choice(len(all_frames), max_frames, replace=False))
+        frames = [all_frames[i] for i in sampled_indices]
+        print(f"Sampled frame indices: {sampled_indices}")
+        
+        # Adjust FPS to maintain normal playback speed
+        # If we sampled N frames from M total frames, adjust FPS proportionally
+        adjusted_fps = original_fps * (len(frames) / len(all_frames))
+    else:
+        frames = all_frames
+        adjusted_fps = original_fps
+        print(f"Video has {len(frames)} frames (≤ {max_frames}), using all frames")
+    
+    print(f"Using {len(frames)} frames from video (Original FPS: {original_fps:.2f}, Adjusted FPS: {adjusted_fps:.2f})")
+    
+    return frames, adjusted_fps
 
 
 def get_prompt(click_state, click_input):
@@ -123,17 +142,21 @@ def get_prompt(click_state, click_input):
     return click_state
 
 
-def load_video(video_input, video_state):
+def load_video(video_input, video_state, num_frames):
     """
     Load video and extract first frame for mask generation
     """
+    # Clean up old output files if they exist
+    if video_state is not None and "output_paths" in video_state:
+        cleanup_old_videos(video_state["output_paths"])
+    
     if video_input is None:
         return video_state, None, \
                gr.update(visible=False), gr.update(visible=False), \
                gr.update(visible=False), gr.update(visible=False)
     
-    # Extract frames
-    frames, fps = extract_frames_from_video(video_input, max_frames=50)
+    # Extract frames with user-specified number
+    frames, fps = extract_frames_from_video(video_input, max_frames=num_frames)
     
     if len(frames) == 0:
         return video_state, None, \
@@ -359,6 +382,9 @@ def run_videomama_with_sam2(video_state, click_state):
     
     status_msg = f"✓ Complete! Generated {len(output_frames)} frames."
     
+    # Store paths for cleanup later
+    video_state["output_paths"] = [str(output_video_path), str(mask_video_path), str(greenscreen_path)]
+    
     return video_state, str(output_video_path), str(mask_video_path), str(greenscreen_path), status_msg
 
 
@@ -382,6 +408,42 @@ def save_video(frames, output_path, fps):
     print(f"Saved video to {output_path}")
 
 
+def cleanup_old_videos(video_paths):
+    """Remove old output videos to save storage space"""
+    if video_paths is None:
+        return
+    
+    for path in video_paths:
+        try:
+            if os.path.exists(path):
+                os.remove(path)
+                print(f"Cleaned up: {path}")
+        except Exception as e:
+            print(f"Failed to remove {path}: {e}")
+
+
+def cleanup_old_outputs(max_age_minutes=30):
+    """
+    Remove output files older than max_age_minutes to prevent storage overflow
+    This runs periodically to clean up abandoned files
+    """
+    output_dir = Path("outputs")
+    if not output_dir.exists():
+        return
+    
+    current_time = time.time()
+    max_age_seconds = max_age_minutes * 60
+    
+    for file_path in output_dir.glob("*.mp4"):
+        try:
+            file_age = current_time - file_path.stat().st_mtime
+            if file_age > max_age_seconds:
+                file_path.unlink()
+                print(f"Cleaned up old file: {file_path} (age: {file_age/60:.1f} minutes)")
+        except Exception as e:
+            print(f"Failed to clean up {file_path}: {e}")
+
+
 def restart():
     """Reset all states"""
     return None, [[], []], None, \
@@ -431,6 +493,14 @@ with gr.Blocks(title="VideoMaMa Demo") as demo:
         with gr.Column(scale=1):
             gr.Markdown("### Step 1: Upload Video")
             video_input = gr.Video(label="Input Video")
+            num_frames_slider = gr.Slider(
+                minimum=1,
+                maximum=50,
+                value=24,
+                step=1,
+                label="Number of Frames",
+                info="More frames = better quality but slower processing"
+            )
             load_button = gr.Button("📁 Load Video", variant="primary")
             
             gr.Markdown("### Step 2: Mark Object")
@@ -466,7 +536,7 @@ with gr.Blocks(title="VideoMaMa Demo") as demo:
     # Event handlers
     load_button.click(
         fn=load_video,
-        inputs=[video_input, video_state],
+        inputs=[video_input, video_state, num_frames_slider],
         outputs=[video_state, first_frame_display, 
                 point_prompt, clear_button, run_button, status_text]
     )
@@ -511,6 +581,9 @@ if __name__ == "__main__":
     print("VideoMaMa Interactive Demo")
     print("=" * 60)
     
+    # Clean up old output files on startup
+    cleanup_old_outputs(max_age_minutes=30)
+    
     # Models will be initialized on first use (lazy loading for ZeroGPU)
     # initialize_models()