app.py

"""
SAM2 Video Segmentation Space
Removes background from videos by tracking specified objects.
Provides both Gradio UI and API endpoints.
"""

import gradio as gr
import torch
import numpy as np
import cv2
import tempfile
import os
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Any
from transformers import Sam2VideoModel, Sam2VideoProcessor
from transformers.video_utils import load_video
from PIL import Image
import json

# Global model variables
MODEL_NAME = "facebook/sam2.1-hiera-tiny"  # Options: tiny, small, base-plus, large
device = None
model = None
processor = None


def initialize_model():
    """Initialize SAM2 model and processor."""
    global device, model, processor
    
    # Determine device
    if torch.cuda.is_available():
        device = torch.device("cuda")
        dtype = torch.float16
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        dtype = torch.float32
    else:
        device = torch.device("cpu")
        dtype = torch.float32
    
    print(f"Loading SAM2 model on {device}...")
    
    # Load model and processor
    model = Sam2VideoModel.from_pretrained(MODEL_NAME).to(device, dtype=dtype)
    processor = Sam2VideoProcessor.from_pretrained(MODEL_NAME)
    
    print("Model loaded successfully!")
    return device, model, processor


def extract_frames_from_video(video_path: str, max_frames: Optional[int] = None) -> Tuple[List[Image.Image], Dict]:
    """Extract frames from video file."""
    video_frames, info = load_video(video_path)
    
    if max_frames and len(video_frames) > max_frames:
        # Sample frames uniformly
        indices = np.linspace(0, len(video_frames) - 1, max_frames, dtype=int)
        video_frames = [video_frames[i] for i in indices]
    
    return video_frames, info


def create_output_video(
    video_frames: List[Image.Image],
    masks: Dict[int, torch.Tensor],
    output_path: str,
    fps: float = 30.0,
    remove_background: bool = True
) -> str:
    """
    Create output video with segmented objects.
    
    Args:
        video_frames: Original video frames
        masks: Dictionary mapping frame_idx to mask tensors
        output_path: Path to save output video
        fps: Frames per second
        remove_background: If True, remove background; if False, highlight objects
    """
    if not masks:
        raise ValueError("No masks provided")
    
    # Get first frame to determine dimensions
    first_frame = np.array(video_frames[0])
    height, width = first_frame.shape[:2]
    
    # Initialize video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    for frame_idx, frame_pil in enumerate(video_frames):
        frame = np.array(frame_pil)
        
        if frame_idx in masks:
            mask = masks[frame_idx].cpu().numpy()
            
            # Handle different mask shapes
            if mask.ndim == 4:  # (batch, num_objects, H, W)
                mask = mask[0]  # Take first batch
            if mask.ndim == 3:  # (num_objects, H, W)
                # Combine all object masks
                mask = mask.max(axis=0)
            
            # Resize mask to frame size if needed
            if mask.shape != (height, width):
                mask = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST)
            
            # Convert to binary mask
            mask_binary = (mask > 0.5).astype(np.uint8)
            
            if remove_background:
                # Keep only the tracked objects (remove background)
                if frame.shape[2] == 3:  # RGB
                    # Create RGBA with alpha channel
                    result = np.zeros((height, width, 4), dtype=np.uint8)
                    result[:, :, :3] = frame
                    result[:, :, 3] = mask_binary * 255
                    
                    # Convert back to RGB with black background
                    background = np.zeros_like(frame)
                    mask_3d = np.repeat(mask_binary[:, :, np.newaxis], 3, axis=2)
                    result_rgb = frame * mask_3d + background * (1 - mask_3d)
                    frame = result_rgb.astype(np.uint8)
            else:
                # Highlight tracked objects (overlay colored mask)
                overlay = frame.copy()
                overlay[mask_binary > 0] = [0, 255, 0]  # Green overlay
                frame = cv2.addWeighted(frame, 0.7, overlay, 0.3, 0)
        
        # Convert RGB to BGR for OpenCV
        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        out.write(frame_bgr)
    
    out.release()
    return output_path


def segment_video(
    video_path: str,
    annotations: List[Dict[str, Any]],
    remove_background: bool = True,
    max_frames: Optional[int] = None
) -> str:
    """
    Main function to segment video based on annotations.
    
    Args:
        video_path: Path to input video
        annotations: List of annotation dictionaries with format:
            [
                {
                    "frame_idx": 0,
                    "object_id": 1,
                    "points": [[x1, y1], [x2, y2], ...],
                    "labels": [1, 1, ...]  # 1 for foreground, 0 for background
                },
                ...
            ]
        remove_background: If True, remove background; if False, highlight objects
        max_frames: Maximum number of frames to process (None = all frames)
    
    Returns:
        Path to output video file
    """
    global device, model, processor
    
    if model is None:
        initialize_model()
    
    # Load video frames
    print("Loading video frames...")
    video_frames, video_info = extract_frames_from_video(video_path, max_frames)
    fps = video_info.get('fps', 30.0)
    
    print(f"Processing {len(video_frames)} frames at {fps} FPS")
    
    # Initialize inference session
    dtype = torch.float16 if device.type == "cuda" else torch.float32
    inference_session = processor.init_video_session(
        video=video_frames,
        inference_device=device,
        dtype=dtype,
    )
    
    # Add annotations to inference session
    print("Adding annotations...")
    for ann in annotations:
        frame_idx = ann["frame_idx"]
        obj_id = ann["object_id"]
        points = ann.get("points", [])
        labels = ann.get("labels", [1] * len(points))
        
        if points:
            # Format points for processor: [[[[x, y], [x, y], ...]]]
            formatted_points = [[points]]
            formatted_labels = [[labels]]
            
            processor.add_inputs_to_inference_session(
                inference_session=inference_session,
                frame_idx=frame_idx,
                obj_ids=obj_id,
                input_points=formatted_points,
                input_labels=formatted_labels,
            )
            
            # Run inference on this frame
            outputs = model(
                inference_session=inference_session,
                frame_idx=frame_idx,
            )
    
    # Propagate through all frames
    print("Propagating masks through video...")
    video_segments = {}
    
    for sam2_output in model.propagate_in_video_iterator(inference_session):
        video_res_masks = processor.post_process_masks(
            [sam2_output.pred_masks],
            original_sizes=[[inference_session.video_height, inference_session.video_width]],
            binarize=False
        )[0]
        video_segments[sam2_output.frame_idx] = video_res_masks
    
    print(f"Generated masks for {len(video_segments)} frames")
    
    # Create output video
    output_path = tempfile.mktemp(suffix=".mp4")
    print("Creating output video...")
    create_output_video(video_frames, video_segments, output_path, fps, remove_background)
    
    print(f"Output video saved to: {output_path}")
    return output_path


# ============================================================================
# GRADIO INTERFACE
# ============================================================================

def gradio_segment_video(
    video_file,
    annotation_json: str,
    remove_bg: bool = True,
    max_frames: Optional[int] = None
):
    """
    Gradio wrapper for video segmentation.
    
    Args:
        video_file: Uploaded video file
        annotation_json: JSON string with annotations
        remove_bg: Whether to remove background
        max_frames: Maximum frames to process
    """
    try:
        # Parse annotations
        annotations = json.loads(annotation_json)
        
        if not isinstance(annotations, list):
            return None, "Error: Annotations must be a list of objects"
        
        # Process video
        output_path = segment_video(
            video_path=video_file,
            annotations=annotations,
            remove_background=remove_bg,
            max_frames=max_frames
        )
        
        return output_path, "✅ Video processed successfully!"
        
    except json.JSONDecodeError as e:
        return None, f"❌ JSON parsing error: {str(e)}"
    except Exception as e:
        return None, f"❌ Error: {str(e)}"


def gradio_simple_segment(
    video_file,
    point_x: int,
    point_y: int,
    frame_idx: int = 0,
    remove_bg: bool = True,
    max_frames: Optional[int] = 300
):
    """
    Simple Gradio interface with single point annotation.
    """
    try:
        # Create simple annotation
        annotations = [{
            "frame_idx": frame_idx,
            "object_id": 1,
            "points": [[point_x, point_y]],
            "labels": [1]
        }]
        
        # Process video
        output_path = segment_video(
            video_path=video_file,
            annotations=annotations,
            remove_background=remove_bg,
            max_frames=max_frames
        )
        
        return output_path, f"✅ Video processed! Tracked from point ({point_x}, {point_y}) on frame {frame_idx}"
        
    except Exception as e:
        return None, f"❌ Error: {str(e)}"


# ============================================================================
# API ENDPOINTS (via Gradio API)
# ============================================================================

def api_segment_video(video_file, annotations_json: str, remove_background: bool = True, max_frames: int = None):
    """
    API endpoint for video segmentation.
    Can be called via gradio_client or direct HTTP requests.
    """
    annotations = json.loads(annotations_json)
    output_path = segment_video(video_file, annotations, remove_background, max_frames)
    return output_path


# ============================================================================
# CREATE GRADIO APP
# ============================================================================

def create_interface():
    """Create the Gradio interface."""
    
    # Initialize model
    initialize_model()
    
    # Create tabs for different interfaces
    with gr.Blocks(title="SAM2 Video Segmentation - Remove Background") as app:
        gr.Markdown("""
        # 🎥 SAM2 Video Background Remover
        
        Remove backgrounds from videos by tracking objects. Uses Meta's Segment Anything Model 2 (SAM2).
        
        **Two ways to use this:**
        1. **Simple Mode**: Click on an object in the first frame
        2. **Advanced Mode**: Provide detailed JSON annotations
        3. **API Mode**: Use the API endpoint programmatically
        """)
        
        with gr.Tabs():
            # ===================== SIMPLE MODE =====================
            with gr.Tab("Simple Mode"):
                gr.Markdown("""
                ### Quick Start
                1. Upload a video
                2. Specify the coordinates of the object you want to track
                3. Click "Process Video"
                
                **Tip**: Open your video in an image viewer to find the x,y coordinates of your target object in the first frame.
                """)
                
                with gr.Row():
                    with gr.Column():
                        simple_video_input = gr.Video(label="Upload Video")
                        
                        with gr.Row():
                            point_x_input = gr.Number(label="Point X", value=320, precision=0)
                            point_y_input = gr.Number(label="Point Y", value=240, precision=0)
                        
                        frame_idx_input = gr.Number(label="Frame Index", value=0, precision=0, 
                                                    info="Which frame to annotate (usually 0 for first frame)")
                        
                        remove_bg_simple = gr.Checkbox(label="Remove Background", value=True,
                                                       info="If checked, removes background. If unchecked, highlights object.")
                        
                        max_frames_simple = gr.Number(label="Max Frames (optional)", value=300, precision=0,
                                                      info="Limit frames for faster processing. Leave at 0 for all frames.")
                        
                        simple_btn = gr.Button("🎬 Process Video", variant="primary")
                    
                    with gr.Column():
                        simple_output_video = gr.Video(label="Output Video")
                        simple_status = gr.Textbox(label="Status", lines=3)
                
                simple_btn.click(
                    fn=gradio_simple_segment,
                    inputs=[simple_video_input, point_x_input, point_y_input, frame_idx_input, 
                           remove_bg_simple, max_frames_simple],
                    outputs=[simple_output_video, simple_status]
                )
                
                gr.Markdown("""
                ### Example:
                For a 640x480 video with a person in the center, try: X=320, Y=240, Frame=0
                """)
            
            # ===================== ADVANCED MODE =====================
            with gr.Tab("Advanced Mode (JSON)"):
                gr.Markdown("""
                ### Advanced Annotations
                Provide detailed JSON annotations for multiple objects and frames.
                
                **JSON Format:**
                ```json
                [
                    {
                        "frame_idx": 0,
                        "object_id": 1,
                        "points": [[x1, y1], [x2, y2]],
                        "labels": [1, 1]
                    }
                ]
                ```
                
                - `frame_idx`: Frame number to annotate
                - `object_id`: Unique ID for each object (1, 2, 3, ...)
                - `points`: List of [x, y] coordinates
                - `labels`: 1 for foreground point, 0 for background point
                """)
                
                with gr.Row():
                    with gr.Column():
                        adv_video_input = gr.Video(label="Upload Video")
                        
                        adv_annotation_input = gr.Textbox(
                            label="Annotations (JSON)",
                            lines=10,
                            value='''[
    {
        "frame_idx": 0,
        "object_id": 1,
        "points": [[320, 240]],
        "labels": [1]
    }
]''',
                            placeholder="Enter JSON annotations here..."
                        )
                        
                        remove_bg_adv = gr.Checkbox(label="Remove Background", value=True)
                        max_frames_adv = gr.Number(label="Max Frames (0 = all)", value=0, precision=0)
                        
                        adv_btn = gr.Button("🎬 Process Video", variant="primary")
                    
                    with gr.Column():
                        adv_output_video = gr.Video(label="Output Video")
                        adv_status = gr.Textbox(label="Status", lines=3)
                
                adv_btn.click(
                    fn=gradio_segment_video,
                    inputs=[adv_video_input, adv_annotation_input, remove_bg_adv, max_frames_adv],
                    outputs=[adv_output_video, adv_status]
                )
            
            # ===================== API INFO =====================
            with gr.Tab("API Documentation"):
                gr.Markdown("""
                ## 📡 API Usage
                
                This Space exposes an API that you can call programmatically.
                
                ### Using Python with `gradio_client`
                
                ```python
                from gradio_client import Client
                import json
                
                # Connect to the Space
                client = Client("YOUR_USERNAME/YOUR_SPACE_NAME")
                
                # Define annotations
                annotations = [
                    {
                        "frame_idx": 0,
                        "object_id": 1,
                        "points": [[320, 240]],
                        "labels": [1]
                    }
                ]
                
                # Call the API
                result = client.predict(
                    video_file="path/to/video.mp4",
                    annotations_json=json.dumps(annotations),
                    remove_background=True,
                    max_frames=300,
                    api_name="/segment_video_api"
                )
                
                print(f"Output video: {result}")
                ```
                
                ### Using cURL
                
                ```bash
                curl -X POST https://YOUR_USERNAME-YOUR_SPACE_NAME.hf.space/api/predict \\
                  -H "Content-Type: application/json" \\
                  -F "data=@video.mp4" \\
                  -F 'annotations=[{"frame_idx":0,"object_id":1,"points":[[320,240]],"labels":[1]}]'
                ```
                
                ### Parameters
                
                - **video_file**: Video file (required)
                - **annotations_json**: JSON string with annotations (required)
                - **remove_background**: Boolean (default: true)
                - **max_frames**: Integer (default: null, processes all frames)
                
                ### Response
                
                Returns the path to the processed video file.
                """)
        
        # Add API endpoint
        api_interface = gr.Interface(
            fn=api_segment_video,
            inputs=[
                gr.Video(label="Video File"),
                gr.Textbox(label="Annotations JSON"),
                gr.Checkbox(label="Remove Background", value=True),
                gr.Number(label="Max Frames", value=None, precision=0)
            ],
            outputs=gr.Video(label="Output Video"),
            api_name="segment_video_api",
            visible=False  # Hidden from UI, only accessible via API
        )
    
    return app


# ============================================================================
# LAUNCH
# ============================================================================

if __name__ == "__main__":
    app = create_interface()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )