Upload folder using huggingface_hub

.dockerignore

@@ -0,0 +1 @@
+.venv

.gitattributes

@@ -33,3 +33,12 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+sam2-repo/demo/data/gallery/01_dog.mp4 filter=lfs diff=lfs merge=lfs -text
+sam2-repo/demo/data/gallery/02_cups.mp4 filter=lfs diff=lfs merge=lfs -text
+sam2-repo/demo/data/gallery/03_blocks.mp4 filter=lfs diff=lfs merge=lfs -text
+sam2-repo/demo/data/gallery/04_coffee.mp4 filter=lfs diff=lfs merge=lfs -text
+sam2-repo/demo/data/gallery/05_default_juggle.mp4 filter=lfs diff=lfs merge=lfs -text
+sam2-repo/demo/frontend/src/assets/videos/sam2_720px_dark.mp4 filter=lfs diff=lfs merge=lfs -text
+sam2-repo/notebooks/images/cars.jpg filter=lfs diff=lfs merge=lfs -text
+sam2-repo/notebooks/videos/bedroom.mp4 filter=lfs diff=lfs merge=lfs -text
+sam2-repo/sav_dataset/example/sav_000001.mp4 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,11 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+models/

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

.python-version

@@ -0,0 +1 @@
+3.12

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+  "python.languageServer": "None"
+}

Dockerfile

@@ -0,0 +1,22 @@
+# Use an official Python runtime as a parent image
+FROM python:3.9-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy the current directory contents into the container at /app
+COPY . /app
+
+# Install any needed packages specified in requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Install additional dependencies
+RUN apt-get update && apt-get install -y \
+  ffmpeg \
+  && rm -rf /var/lib/apt/lists/*
+
+# Expose port 7860 for Gradio
+EXPOSE 7860
+
+# Run the application
+CMD ["python", "app.py"]

README.md

@@ -1,12 +1,47 @@
 ---
-title: Golf Tracking
-emoji: 🐨
-colorFrom: yellow
-colorTo: green
+title: golf_tracking
+app_file: gradio_app.py
 sdk: gradio
-sdk_version: 5.11.0
-app_file: app.py
-pinned: false
+sdk_version: 5.10.0
 ---
+# Golf Ball Trajectory Tracker
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+This application uses Meta's Segment Anything Model (SAM) 2.1 and physics-based trajectory fitting to track and analyze golf ball trajectories from videos.
+
+## Setup
+
+1. Download the SAM model checkpoint:
+   - Visit: https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
+   - Save the file in the project directory
+
+2. Install dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+## Usage
+
+1. Run the application:
+```bash
+python app.py
+```
+
+2. Using the interface:
+   - Upload a golf swing video
+   - Click "Process Video" to load the video
+   - Click on the golf ball position in key frames (at least 3 points recommended)
+   - The application will fit and display the trajectory
+   - Use "Clear Points" to start over
+
+## Features
+
+- Interactive point selection for trajectory tracking
+- Physics-based trajectory fitting
+- Real-time visualization of predicted path
+- CPU-optimized for Intel UHD 630 graphics
+
+## Notes
+
+- The application uses the SAM base model optimized for CPU usage
+- For best results, provide clear video footage with visible golf ball
+- Select points at different stages of the ball's flight for better trajectory fitting

app.py

@@ -0,0 +1,381 @@
+import gradio as gr
+import numpy as np
+import cv2
+import torch
+import os
+import logging
+import contextlib
+from sam2.build_sam import build_sam2_video_predictor
+
+# Add current directory to path
+import sys
+
+sys.path.append(os.getcwd())
+sys.path.append(os.path.join(os.getcwd(), "sam2"))  # Add sam2 directory to path
+print(f"current dir is {os.getcwd()}")
+
+# Ensure device setup matches the official code
+force_cpu_device = os.environ.get("SAM2_DEMO_FORCE_CPU_DEVICE", "0") == "1"
+if force_cpu_device:
+    logging.info("forcing CPU device for SAM 2 demo")
+if torch.cuda.is_available() and not force_cpu_device:
+    DEVICE = torch.device("cuda")
+elif torch.backends.mps.is_available() and not force_cpu_device:
+    DEVICE = torch.device("mps")
+else:
+    DEVICE = torch.device("cpu")
+logging.info(f"using device: {DEVICE}")
+
+if DEVICE.type == "cuda":
+    if torch.cuda.get_device_properties(0).major >= 8:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+elif DEVICE.type == "mps":
+    logging.warning(
+        "\nSupport for MPS devices is preliminary. SAM 2 is trained with CUDA and might "
+        "give numerically different outputs and sometimes degraded performance on MPS. "
+        "See e.g. https://github.com/pytorch/pytorch/issues/84936 for a discussion."
+    )
+
+
+def load_model_paths(checkpoint_name):
+    """Get model checkpoint and config paths"""
+    if checkpoint_name == "SAM2-T":
+        sam2_checkpoint = "models/sam2.1_hiera_tiny.pt"
+        model_cfg = "configs/sam2.1/sam2.1_hiera_t.yaml"
+    elif checkpoint_name == "SAM2-S":
+        sam2_checkpoint = "models/sam2.1_hiera_small.pt"
+        model_cfg = "configs/sam2.1/sam2.1_hiera_s.yaml"
+    elif checkpoint_name == "SAM2-B_PLUS":
+        sam2_checkpoint = "models/sam2.1_hiera_base_plus.pt"
+        model_cfg = "configs/sam2.1/sam2.1_hiera_b+.yaml"
+    else:
+        raise ValueError(f"Invalid checkpoint name: {checkpoint_name}")
+
+    return sam2_checkpoint, model_cfg
+
+
+# Available checkpoints
+CHECKPOINTS = {
+    "SAM2-B_PLUS": "Base Plus Model",
+    "SAM2-S": "Small Model",
+    "SAM2-T": "Tiny Model",
+}
+
+
+class GolfTracker:
+    def __init__(self, checkpoint="SAM2-T"):
+        """Initialize with specified checkpoint model"""
+        self.current_checkpoint = checkpoint
+        self.predictor = None
+        self.points = []
+        self.frames = []
+        self.current_frame_idx = 0
+        self.video_info = None
+        self.state = None
+        self.obj_id = 1  # Track single object (golf ball)
+        self.device = DEVICE
+        self.out_mask_logits = None
+        self.load_model(checkpoint)
+
+    def load_model(self, checkpoint_name):
+        """Load specified checkpoint model"""
+        if checkpoint_name not in CHECKPOINTS:
+            raise ValueError(f"Invalid checkpoint: {checkpoint_name}")
+
+        print(f"Loading checkpoint: {checkpoint_name}")
+        sam2_checkpoint, model_cfg = load_model_paths(checkpoint_name)
+
+        # Build predictor with model config and checkpoint
+        self.predictor = build_sam2_video_predictor(
+            model_cfg, sam2_checkpoint, self.device
+        )
+        print(f"Model loaded successfully: {CHECKPOINTS[checkpoint_name]}")
+        self.current_checkpoint = checkpoint_name
+
+    def process_video(self, video_path):
+        """Process the video and initialize tracking"""
+        if not os.path.exists(video_path):
+            return None, None, None, "Video file not found"
+
+        # Reset state
+        self.points = []
+        self.frames = []
+        self.current_frame_idx = 0
+        self.state = None
+
+        # Read video frames
+        cap = cv2.VideoCapture(video_path)
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            self.frames.append(frame)
+
+        if not self.frames:
+            return None, None, None, "Failed to read video"
+
+        # Store video info
+        self.video_info = {
+            "path": video_path,
+            "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
+            "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
+            "fps": cap.get(cv2.CAP_PROP_FPS),
+            "total_frames": len(self.frames),
+        }
+
+        cap.release()
+
+        # Initialize SAM2 state
+        with self.autocast_context(), torch.inference_mode():
+            self.state = self.predictor.init_state(video_path)
+
+        return (
+            self.frames[0],  # First frame
+            self.current_checkpoint,
+            gr.Slider(minimum=0, maximum=len(self.frames) - 1, step=1, value=0),
+            "Navigate through frames and click on the golf ball to track",
+        )
+
+    def update_frame(self, frame_idx):
+        """Update displayed frame"""
+        if not self.frames or frame_idx >= len(self.frames):
+            return None
+
+        self.current_frame_idx = int(frame_idx)
+        frame = self.frames[self.current_frame_idx].copy()
+
+        # Draw existing points and trajectory
+        self._draw_tracking(frame)
+        return frame
+
+    def add_point(self, frame, evt: gr.SelectData):
+        """Add a point and get ball prediction with enhanced mask visualization"""
+        if self.state is None:
+            return frame
+
+        x, y = evt.index[0], evt.index[1]
+        self.points.append((self.current_frame_idx, x, y))
+
+        frame_with_points = frame.copy()
+
+        # Get ball prediction using SAM2.1
+        with self.autocast_context(), torch.inference_mode():
+            # Convert points and labels to numpy arrays
+            points = np.array([(x, y)], dtype=np.float32)
+            labels = np.array([1], dtype=np.int32)  # 1 for positive click
+
+            # Add point and get mask
+            _, out_obj_ids, out_mask_logits = self.predictor.add_new_points(
+                inference_state=self.state,
+                frame_idx=self.current_frame_idx,
+                obj_id=self.obj_id,
+                points=points,
+                labels=labels,
+            )
+
+            if out_mask_logits is not None and len(out_mask_logits) > 0:
+                self.out_mask_logits = out_mask_logits
+
+        # Draw tracking visualization
+        self._draw_tracking(frame_with_points)
+        return frame_with_points
+
+    def propagate_masks(self):
+        """Propagate masks to the entire video after user selection"""
+        if self.state is None:
+            return "No state initialized"
+
+        logging.info(f"Propagating masks in video with state: {self.state}")
+
+        # Propagate the masks across the video
+        with self.autocast_context(), torch.inference_mode():
+            frame_idx, obj_ids, video_res_masks = self.predictor.propagate_in_video(
+                inference_state=self.state,
+                start_frame_idx=0,
+                reverse=False,
+            )
+
+            self.out_mask_logits = video_res_masks
+
+        return "Propagation complete"
+
+    def autocast_context(self):
+        if self.device.type == "cuda":
+            return torch.autocast("cuda", dtype=torch.bfloat16)
+        else:
+            return contextlib.nullcontext()
+
+    def _draw_tracking(self, frame):
+        """Draw object mask on frame with enhanced visualization"""
+        # Assuming out_mask_logits is available from propagate_masks
+        if self.current_frame_idx < len(self.frames):
+            mask_np = (self.out_mask_logits[self.current_frame_idx] > 0.0).cpu().numpy()
+            if mask_np.shape[:2] == frame.shape[:2]:
+                overlay = frame.copy()
+                overlay[mask_np > 0] = [0, 0, 255]  # Red color for mask
+                alpha = 0.5  # Transparency factor
+                frame = cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
+        return frame
+
+    def clear_points(self):
+        """Clear all tracked points"""
+        self.points = []
+        if self.frames:
+            return self.frames[self.current_frame_idx].copy()
+        return None
+
+    def change_model(self, checkpoint_name):
+        """Change the current model checkpoint"""
+        if checkpoint_name != self.current_checkpoint:
+            self.load_model(checkpoint_name)
+        return f"Loaded {CHECKPOINTS[checkpoint_name]}"
+
+    def save_output_video(self):
+        """Save the processed video with tracking visualization"""
+        if not self.frames or not self.video_info:
+            return None, "No video loaded"
+
+        output_path = "output_tracked.mp4"
+
+        # Initialize video writer
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        out = cv2.VideoWriter(
+            output_path,
+            fourcc,
+            self.video_info["fps"],
+            (self.video_info["width"], self.video_info["height"]),
+        )
+
+        # Process each frame
+        for frame_idx in range(len(self.frames)):
+            frame = self.frames[frame_idx].copy()
+
+            # Draw tracking for this frame
+            frame_points = [(x, y) for f, x, y in self.points if f == frame_idx]
+            if frame_points:
+                # Draw points
+                for x, y in frame_points:
+                    cv2.circle(frame, (int(x), int(y)), 5, (255, 0, 0), -1)
+
+                # Fit and draw trajectory if enough points
+                if len(frame_points) >= 3:
+                    points_arr = np.array(frame_points)
+                    # fit_results = self.trajectory_fitter.fit_trajectory(points_arr)
+
+                    # if fit_results is not None:
+                    #     trajectory = fit_results["trajectory"]
+                    #     points = trajectory.astype(np.int32)
+                    #     for i in range(len(points) - 1):
+                    #         cv2.line(
+                    #             frame,
+                    #             tuple(points[i]),
+                    #             tuple(points[i + 1]),
+                    #             (0, 255, 0),
+                    #             2,
+                    #         )
+
+                    #     # Calculate and display metrics
+                    #     metrics = self.trajectory_fitter.calculate_metrics(fit_results)
+                    #     cv2.putText(
+                    #         frame,
+                    #         f"Speed: {metrics['initial_velocity_mph']:.1f} mph",
+                    #         (10, 30),
+                    #         cv2.FONT_HERSHEY_SIMPLEX,
+                    #         1,
+                    #         (255, 255, 255),
+                    #         2,
+                    #     )
+                    #     cv2.putText(
+                    #         frame,
+                    #         f"Height: {metrics['max_height']:.1f} m",
+                    #         (10, 70),
+                    #         cv2.FONT_HERSHEY_SIMPLEX,
+                    #         1,
+                    #         (255, 255, 255),
+                    #         2,
+                    #     )
+
+            out.write(frame)
+
+        out.release()
+        return output_path, "Video saved successfully!"
+
+
+def create_ui():
+    tracker = GolfTracker()
+
+    with gr.Blocks() as app:
+        gr.Markdown("# Golf Ball Trajectory Tracker")
+        gr.Markdown(
+            "Upload a video and click on the golf ball positions to track its trajectory"
+        )
+
+        with gr.Row():
+            with gr.Column():
+                video_input = gr.Video(label="Input Video")
+                model_dropdown = gr.Dropdown(
+                    choices=list(CHECKPOINTS.keys()),
+                    value="SAM2-T",
+                    label="Select Model",
+                )
+                upload_button = gr.Button("Process Video")
+                clear_button = gr.Button("Clear Points")
+                save_button = gr.Button("Save Output Video")
+                propagate_button = gr.Button("Propagate Masks")
+
+            with gr.Column():
+                image_output = gr.Image(label="Click on golf ball positions")
+                frame_slider = gr.Slider(
+                    minimum=0,
+                    maximum=0,
+                    step=1,
+                    value=0,
+                    label="Frame",
+                    interactive=True,
+                )
+                current_model = gr.Textbox(label="Current Model", interactive=False)
+                status_text = gr.Textbox(label="Status", interactive=False)
+                output_video = gr.Video(label="Output Video")
+
+        # Event handlers
+        model_dropdown.change(
+            fn=tracker.change_model, inputs=[model_dropdown], outputs=[status_text]
+        )
+
+        video_input.change(
+            fn=tracker.process_video,
+            inputs=[video_input],
+            outputs=[image_output, current_model, frame_slider, status_text],
+        )
+
+        upload_button.click(
+            fn=tracker.process_video,
+            inputs=[video_input],
+            outputs=[image_output, current_model, frame_slider, status_text],
+        )
+
+        clear_button.click(fn=tracker.clear_points, inputs=[], outputs=[image_output])
+
+        frame_slider.change(
+            fn=tracker.update_frame, inputs=[frame_slider], outputs=[image_output]
+        )
+
+        image_output.select(
+            fn=tracker.add_point, inputs=[image_output], outputs=[image_output]
+        )
+
+        save_button.click(
+            fn=tracker.save_output_video, inputs=[], outputs=[output_video, status_text]
+        )
+
+        propagate_button.click(
+            fn=tracker.propagate_masks, inputs=[], outputs=[status_text]
+        )
+
+    return app
+
+
+if __name__ == "__main__":
+    app = create_ui()
+    app.launch()

backup_001.py

@@ -0,0 +1,330 @@
+import gradio as gr
+import torch
+import numpy as np
+import cv2
+from sam2.build_sam import build_sam2_video_predictor
+import tempfile
+import os
+import contextlib
+
+
+class VideoTracker:
+    def __init__(self):
+        self.checkpoint = "./models/sam2.1_hiera_tiny.pt"
+        self.model_cfg = "configs/sam2.1/sam2.1_hiera_t.yaml"
+        self.predictor = build_sam2_video_predictor(
+            self.model_cfg, self.checkpoint, device="cpu", mode="eval"
+        )
+        self.state = None
+        self.video_frames = None
+        self.current_frame_idx = 0
+        self.masks = []
+        self.points = []
+        self.frame_count = 0
+        self.video_info = None
+        self.obj_id = 1
+        self.out_mask_logits = None
+        self.frame_masks = {}  # Store masks for each frame
+
+    def load_video(self, video_path):
+        if video_path is None:
+            return None, gr.Slider(minimum=0, maximum=0, step=1, value=0)
+
+        # Create a temporary file for the video
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
+        temp_file.close()
+
+        # Copy the uploaded video to the temporary file
+        with open(video_path, "rb") as f_src, open(temp_file.name, "wb") as f_dst:
+            f_dst.write(f_src.read())
+
+        # Load video frames using OpenCV
+        cap = cv2.VideoCapture(temp_file.name)
+        frames = []
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frames.append(frame)
+
+        if not frames:
+            cap.release()
+            os.unlink(temp_file.name)
+            return None, gr.Slider(minimum=0, maximum=0, step=1, value=0)
+
+        # Store video info
+        self.video_info = {
+            "path": temp_file.name,
+            "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
+            "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
+            "fps": cap.get(cv2.CAP_PROP_FPS),
+            "total_frames": len(frames),
+        }
+        cap.release()
+
+        self.video_frames = frames
+        self.frame_count = len(frames)
+
+        # Initialize SAM2 state with video path
+        with torch.inference_mode():
+            self.state = self.predictor.init_state(temp_file.name)
+
+        # Now we can remove the temp file
+        os.unlink(temp_file.name)
+
+        return frames[0], gr.Slider(minimum=0, maximum=len(frames) - 1, step=1, value=0)
+
+    def update_frame(self, frame_number):
+        if self.video_frames is None:
+            return None
+
+        self.current_frame_idx = frame_number
+        frame = self.video_frames[frame_number].copy()
+
+        # Apply any existing mask for this frame
+        if frame_number in self.frame_masks:
+            self.out_mask_logits = self.frame_masks[frame_number]
+            frame = self._draw_tracking(frame)
+
+        # Draw points
+        for point in self.points:
+            if point[0] == frame_number:
+                cv2.circle(
+                    frame, (int(point[1]), int(point[2])), 5, (255, 255, 0), -1
+                )  # Yellow dot
+                cv2.circle(
+                    frame, (int(point[1]), int(point[2])), 7, (0, 0, 0), 1
+                )  # Black border
+
+        return frame
+
+    def add_point(self, frame, evt: gr.SelectData):
+        """Add a point and get ball prediction with enhanced mask visualization"""
+        if self.state is None:
+            return frame
+
+        x, y = evt.index[0], evt.index[1]
+        self.points.append((self.current_frame_idx, x, y))
+
+        frame_with_points = frame.copy()
+
+        # Get ball prediction using SAM2.1
+        with torch.inference_mode():
+            # Convert points and labels to numpy arrays
+            points = np.array([(x, y)], dtype=np.float32)
+            labels = np.array([1], dtype=np.int32)  # 1 for positive click
+
+            # Add point and get mask
+            _, out_obj_ids, out_mask_logits = self.predictor.add_new_points(
+                inference_state=self.state,
+                frame_idx=self.current_frame_idx,
+                obj_id=self.obj_id,
+                points=points,
+                labels=labels,
+            )
+
+            if out_mask_logits is not None and len(out_mask_logits) > 0:
+                self.out_mask_logits = (
+                    out_mask_logits[0]
+                    if isinstance(out_mask_logits, list)
+                    else out_mask_logits
+                )
+                # Store mask for this frame
+                self.frame_masks[self.current_frame_idx] = self.out_mask_logits
+
+        # Draw tracking visualization with enhanced mask
+        frame_with_points = self._draw_tracking(frame_with_points)
+
+        # Draw point on top of mask
+        cv2.circle(
+            frame_with_points, (int(x), int(y)), 5, (255, 255, 0), -1
+        )  # Yellow dot
+        cv2.circle(frame_with_points, (int(x), int(y)), 7, (0, 0, 0), 1)  # Black border
+
+        return frame_with_points
+
+    def propagate_video(self):
+        if self.state is None:
+            return None
+
+        output_frames = self.video_frames.copy()
+
+        # Store all masks for smoother visualization
+        all_masks = []
+
+        # First pass: collect all masks
+        with torch.inference_mode():
+            for frame_idx, obj_ids, masks in self.predictor.propagate_in_video(
+                self.state,
+                start_frame_idx=0,
+                reverse=False,
+            ):
+                if masks is not None and len(masks) > 0:
+                    mask = masks[0] if isinstance(masks, list) else masks
+                    all_masks.append((frame_idx, mask))
+                    # Store mask for each frame
+                    self.frame_masks[frame_idx] = mask
+
+        # Second pass: apply visualization with temporal smoothing
+        for i, frame in enumerate(output_frames):
+            frame = frame.copy()
+
+            # Find masks for this frame
+            current_masks = [m[1] for m in all_masks if m[0] == i]
+
+            if current_masks:
+                self.out_mask_logits = current_masks[0]
+
+                # Get binary mask and ensure correct dimensions
+                mask_np = (current_masks[0] > 0.0).cpu().numpy()
+                mask_np = self._handle_mask_dimensions(mask_np)
+
+                # Convert to proper format for OpenCV
+                mask_np = mask_np.astype(np.uint8)
+
+                # Enhanced visualization for video
+                frame = self._draw_tracking(frame, alpha=0.6)
+
+                # Create glowing effect
+                try:
+                    # Create kernel for dilation
+                    kernel = np.ones((5, 5), np.uint8)
+
+                    # Dilate mask for glow effect
+                    dilated_mask = cv2.dilate(mask_np, kernel, iterations=2)
+
+                    # Create glow overlay
+                    glow = frame.copy()
+                    glow[dilated_mask > 0] = [0, 255, 255]  # Yellow glow
+
+                    # Blend glow with frame
+                    frame = cv2.addWeighted(frame, 0.7, glow, 0.3, 0)
+                except cv2.error as e:
+                    print(
+                        f"Warning: Could not apply glow effect. Mask shape: {mask_np.shape}, Frame shape: {frame.shape}"
+                    )
+                    # Continue without glow effect if there's an error
+
+            output_frames[i] = frame
+
+        # Save as video with higher quality
+        temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+        height, width = output_frames[0].shape[:2]
+
+        # Use higher bitrate for better quality
+        writer = cv2.VideoWriter(
+            temp_output, cv2.VideoWriter_fourcc(*"mp4v"), 30, (width, height), True
+        )
+
+        for frame in output_frames:
+            writer.write(frame)
+        writer.release()
+
+        return temp_output
+
+    def _handle_mask_dimensions(self, mask_np):
+        """Helper function to handle various mask dimensions"""
+        # Handle 4D tensor (1, 1, H, W)
+        if len(mask_np.shape) == 4:
+            mask_np = mask_np[0, 0]
+        # Handle 3D tensor (1, H, W) or (H, W, 1)
+        elif len(mask_np.shape) == 3:
+            if mask_np.shape[0] == 1:  # (1, H, W) format
+                mask_np = mask_np[0]
+            elif mask_np.shape[2] == 1:  # (H, W, 1) format
+                mask_np = mask_np[:, :, 0]
+        return mask_np
+
+    def _draw_tracking(self, frame, alpha=0.5):
+        """Draw object mask on frame with enhanced visualization"""
+        if self.out_mask_logits is not None:
+            # Convert logits to binary mask
+            if isinstance(self.out_mask_logits, list):
+                mask = self.out_mask_logits[0]
+            else:
+                mask = self.out_mask_logits
+
+            # Get binary mask and handle dimensions
+            mask_np = (mask > 0.0).cpu().numpy()
+            mask_np = self._handle_mask_dimensions(mask_np)
+
+            if mask_np.shape[:2] == frame.shape[:2]:
+                # Create a red overlay for the mask
+                overlay = frame.copy()
+                overlay[mask_np > 0] = [0, 0, 255]  # BGR format: Red color
+
+                # Add a border around the mask for better visibility
+                contours, _ = cv2.findContours(
+                    mask_np.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+                )
+
+                # Draw thicker contours for better visibility
+                cv2.drawContours(
+                    overlay, contours, -1, (0, 255, 255), 3
+                )  # Thicker yellow border
+
+                # Add a second contour for emphasis
+                cv2.drawContours(
+                    frame, contours, -1, (255, 255, 0), 1
+                )  # Thin bright border
+
+                # Blend the overlay with original frame
+                frame = cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
+
+        return frame
+
+
+def create_interface():
+    tracker = VideoTracker()
+
+    with gr.Blocks() as interface:
+        gr.Markdown("# Object Tracking with SAM2")
+        gr.Markdown("Upload a video and click on objects to track them")
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                video_input = gr.Video(label="Input Video")
+                image_output = gr.Image(label="Current Frame", interactive=True)
+                frame_slider = gr.Slider(
+                    minimum=0,
+                    maximum=0,
+                    step=1,
+                    value=0,
+                    label="Frame Selection",
+                    interactive=True,
+                )
+
+            with gr.Column(scale=1):
+                propagate_btn = gr.Button("Propagate Through Video", variant="primary")
+                video_output = gr.Video(label="Output Video")
+
+        video_input.change(
+            fn=tracker.load_video,
+            inputs=[video_input],
+            outputs=[image_output, frame_slider],
+        )
+
+        frame_slider.change(
+            fn=tracker.update_frame,
+            inputs=[frame_slider],
+            outputs=[image_output],
+        )
+
+        image_output.select(
+            fn=tracker.add_point,
+            inputs=[image_output],
+            outputs=[image_output],
+        )
+
+        propagate_btn.click(
+            fn=tracker.propagate_video,
+            inputs=[],
+            outputs=[video_output],
+        )
+
+    return interface
+
+
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch(share=True)

configs/sam2.1_hiera_b+.yaml

@@ -0,0 +1,116 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 112
+      num_heads: 2
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [896, 448, 224, 112]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

configs/sam2.1_hiera_s.yaml

@@ -0,0 +1,119 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 96
+      num_heads: 1
+      stages: [1, 2, 11, 2]
+      global_att_blocks: [7, 10, 13]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [768, 384, 192, 96]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

configs/sam2.1_hiera_t.yaml

@@ -0,0 +1,121 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 96
+      num_heads: 1
+      stages: [1, 2, 7, 2]
+      global_att_blocks: [5, 7, 9]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [768, 384, 192, 96]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  # SAM decoder
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  # HieraT does not currently support compilation, should always be set to False
+  compile_image_encoder: False

configs/sam2_configs_sam2.1_sam2.1_hiera_l.yaml

@@ -0,0 +1,120 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 144
+      num_heads: 2
+      stages: [2, 6, 36, 4]
+      global_att_blocks: [23, 33, 43]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+      window_spec: [8, 4, 16, 8]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [1152, 576, 288, 144]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

gradio_app.py

@@ -0,0 +1,379 @@
+import gradio as gr
+import torch
+import numpy as np
+import cv2
+from sam2.build_sam import build_sam2_video_predictor
+import tempfile
+import os
+import contextlib
+from trajectory_service import TrajectoryService
+
+
+class VideoTracker:
+    def __init__(self):
+        self.checkpoint = "./models/sam2.1_hiera_tiny.pt"
+        self.model_cfg = "configs/sam2.1/sam2.1_hiera_t.yaml"
+        self.predictor = build_sam2_video_predictor(
+            self.model_cfg, self.checkpoint, device="cpu", mode="eval"
+        )
+        self.state = None
+        self.video_frames = None
+        self.current_frame_idx = 0
+        self.masks = []
+        self.points = []
+        self.frame_count = 0
+        self.video_info = None
+        self.obj_id = 1
+        self.out_mask_logits = None
+        self.frame_masks = {}  # Store masks for each frame
+        self.trajectory_service = None
+
+    def load_video(self, video_path):
+        if video_path is None:
+            return None, gr.Slider(minimum=0, maximum=0, step=1, value=0)
+
+        # Create a temporary file for the video
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
+        temp_file.close()
+
+        # Copy the uploaded video to the temporary file
+        with open(video_path, "rb") as f_src, open(temp_file.name, "wb") as f_dst:
+            f_dst.write(f_src.read())
+
+        # Load video frames using OpenCV
+        cap = cv2.VideoCapture(temp_file.name)
+        frames = []
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frames.append(frame)
+
+        if not frames:
+            cap.release()
+            os.unlink(temp_file.name)
+            return None, gr.Slider(minimum=0, maximum=0, step=1, value=0)
+
+        # Store video info
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        self.video_info = {
+            "path": temp_file.name,
+            "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
+            "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
+            "fps": fps,
+            "total_frames": len(frames),
+        }
+        cap.release()
+
+        self.video_frames = frames
+        self.frame_count = len(frames)
+        self.trajectory_service = TrajectoryService(fps=fps)
+
+        # Initialize SAM2 state with video path
+        with torch.inference_mode():
+            self.state = self.predictor.init_state(temp_file.name)
+
+        # Now we can remove the temp file
+        os.unlink(temp_file.name)
+
+        return frames[0], gr.Slider(minimum=0, maximum=len(frames) - 1, step=1, value=0)
+
+    def update_frame(self, frame_number):
+        if self.video_frames is None:
+            return None
+
+        self.current_frame_idx = frame_number
+        frame = self.video_frames[frame_number].copy()
+
+        # Apply any existing mask for this frame
+        if frame_number in self.frame_masks:
+            self.out_mask_logits = self.frame_masks[frame_number]
+            frame = self._draw_tracking(frame)
+
+        # Draw points (just the points, no trajectory)
+        for point in self.points:
+            if point[0] == frame_number:
+                cv2.circle(
+                    frame, (int(point[1]), int(point[2])), 5, (255, 255, 0), -1
+                )  # Yellow dot
+                cv2.circle(
+                    frame, (int(point[1]), int(point[2])), 7, (0, 0, 0), 1
+                )  # Black border
+
+        return frame
+
+    def add_point(self, frame, evt: gr.SelectData):
+        """Add a point and get ball prediction with enhanced mask visualization"""
+        if self.state is None:
+            return frame
+
+        x, y = evt.index[0], evt.index[1]
+        self.points.append((self.current_frame_idx, x, y))
+
+        # Add point to trajectory service for later use
+        if self.trajectory_service:
+            self.trajectory_service.add_point(self.current_frame_idx, x, y)
+
+        frame_with_points = frame.copy()
+
+        # Get ball prediction using SAM2.1
+        with torch.inference_mode():
+            # Convert points and labels to numpy arrays
+            points = np.array([(x, y)], dtype=np.float32)
+            labels = np.array([1], dtype=np.int32)  # 1 for positive click
+
+            # Add point and get mask
+            _, out_obj_ids, out_mask_logits = self.predictor.add_new_points(
+                inference_state=self.state,
+                frame_idx=self.current_frame_idx,
+                obj_id=self.obj_id,
+                points=points,
+                labels=labels,
+            )
+
+            if out_mask_logits is not None and len(out_mask_logits) > 0:
+                self.out_mask_logits = (
+                    out_mask_logits[0]
+                    if isinstance(out_mask_logits, list)
+                    else out_mask_logits
+                )
+                # Store mask for this frame
+                self.frame_masks[self.current_frame_idx] = self.out_mask_logits
+
+        # Draw tracking visualization with enhanced mask
+        frame_with_points = self._draw_tracking(frame_with_points)
+
+        # Draw point on top of mask (just the point, no trajectory)
+        cv2.circle(
+            frame_with_points, (int(x), int(y)), 5, (255, 255, 0), -1
+        )  # Yellow dot
+        cv2.circle(frame_with_points, (int(x), int(y)), 7, (0, 0, 0), 1)  # Black border
+
+        return frame_with_points
+
+    def propagate_video(self):
+        if self.state is None:
+            return None
+
+        output_frames = self.video_frames.copy()
+
+        # Store all masks and their centers for trajectory calculation
+        all_masks = []
+        mask_centers = []
+
+        # First pass: collect all masks and calculate centers
+        with torch.inference_mode():
+            for frame_idx, obj_ids, masks in self.predictor.propagate_in_video(
+                self.state,
+                start_frame_idx=0,
+                reverse=False,
+            ):
+                if masks is not None and len(masks) > 0:
+                    mask = masks[0] if isinstance(masks, list) else masks
+                    all_masks.append((frame_idx, mask))
+
+                    # Get mask center
+                    mask_np = (mask > 0.0).cpu().numpy()
+                    center = self._get_mask_center(mask_np)
+                    if center is not None:
+                        mask_centers.append((frame_idx, center[0], center[1]))
+
+                    # Store mask for each frame
+                    self.frame_masks[frame_idx] = mask
+
+        # Add detected centers to trajectory service
+        if self.trajectory_service:
+            # Clear existing points and add user-selected points first
+            self.trajectory_service.clear_points()
+            for point in self.points:
+                self.trajectory_service.add_point(point[0], point[1], point[2])
+
+            # Add centers from mask detection
+            for center in mask_centers:
+                if center[0] not in [
+                    p[0] for p in self.points
+                ]:  # Don't duplicate user points
+                    self.trajectory_service.add_point(center[0], center[1], center[2])
+
+            # Calculate trajectory with all points
+            trajectory_points = self.trajectory_service.get_trajectory()
+
+        # Second pass: apply visualization with temporal smoothing and trajectory
+        for i, frame in enumerate(output_frames):
+            frame = frame.copy()
+
+            # Find masks for this frame
+            current_masks = [m[1] for m in all_masks if m[0] == i]
+
+            if current_masks:
+                self.out_mask_logits = current_masks[0]
+                mask_np = (current_masks[0] > 0.0).cpu().numpy()
+                mask_np = self._handle_mask_dimensions(mask_np)
+                mask_np = mask_np.astype(np.uint8)
+                frame = self._draw_tracking(frame, alpha=0.6)
+
+                try:
+                    kernel = np.ones((5, 5), np.uint8)
+                    dilated_mask = cv2.dilate(mask_np, kernel, iterations=2)
+                    glow = frame.copy()
+                    glow[dilated_mask > 0] = [0, 255, 255]  # Yellow glow
+                    frame = cv2.addWeighted(frame, 0.7, glow, 0.3, 0)
+                except cv2.error as e:
+                    print(
+                        f"Warning: Could not apply glow effect. Mask shape: {mask_np.shape}, Frame shape: {frame.shape}"
+                    )
+
+            # Draw trajectory
+            if self.trajectory_service and trajectory_points:
+                frame = self.trajectory_service.draw_trajectory(frame, i)
+
+            output_frames[i] = frame
+
+        # Save as video with higher quality
+        temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+        height, width = output_frames[0].shape[:2]
+        writer = cv2.VideoWriter(
+            temp_output, cv2.VideoWriter_fourcc(*"mp4v"), 30, (width, height), True
+        )
+
+        for frame in output_frames:
+            writer.write(frame)
+        writer.release()
+
+        return temp_output
+
+    def _get_mask_center(self, mask_np):
+        """Get the center point of a mask"""
+        if mask_np is None:
+            return None
+
+        # Ensure mask is 2D
+        mask_np = self._handle_mask_dimensions(mask_np)
+        mask_np = (mask_np > 0.0).astype(np.uint8)
+
+        # Find contours
+        contours, _ = cv2.findContours(
+            mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+        )
+        if not contours:
+            return None
+
+        # Get largest contour
+        largest_contour = max(contours, key=cv2.contourArea)
+
+        # Calculate centroid
+        M = cv2.moments(largest_contour)
+        if M["m00"] == 0:
+            return None
+
+        cx = int(M["m10"] / M["m00"])
+        cy = int(M["m01"] / M["m00"])
+
+        return (cx, cy)
+
+    def _handle_mask_dimensions(self, mask_np):
+        """Helper function to handle various mask dimensions"""
+        # Handle 4D tensor (1, 1, H, W)
+        if len(mask_np.shape) == 4:
+            mask_np = mask_np[0, 0]
+        # Handle 3D tensor (1, H, W) or (H, W, 1)
+        elif len(mask_np.shape) == 3:
+            if mask_np.shape[0] == 1:  # (1, H, W) format
+                mask_np = mask_np[0]
+            elif mask_np.shape[2] == 1:  # (H, W, 1) format
+                mask_np = mask_np[:, :, 0]
+        return mask_np
+
+    def _draw_tracking(self, frame, alpha=0.5):
+        """Draw object mask on frame with enhanced visualization"""
+        if self.out_mask_logits is not None:
+            # Convert logits to binary mask
+            if isinstance(self.out_mask_logits, list):
+                mask = self.out_mask_logits[0]
+            else:
+                mask = self.out_mask_logits
+
+            # Get binary mask and handle dimensions
+            mask_np = (mask > 0.0).cpu().numpy()
+            mask_np = self._handle_mask_dimensions(mask_np)
+
+            if mask_np.shape[:2] == frame.shape[:2]:
+                # Create a red overlay for the mask
+                overlay = frame.copy()
+                overlay[mask_np > 0] = [0, 0, 255]  # BGR format: Red color
+
+                # Add a border around the mask for better visibility
+                contours, _ = cv2.findContours(
+                    mask_np.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+                )
+
+                # Draw thicker contours for better visibility
+                cv2.drawContours(
+                    overlay, contours, -1, (0, 255, 255), 3
+                )  # Thicker yellow border
+
+                # Draw a second contour for emphasis
+                cv2.drawContours(
+                    frame, contours, -1, (255, 255, 0), 1
+                )  # Thin bright border
+
+                # Blend the overlay with original frame
+                frame = cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
+
+        return frame
+
+
+def create_interface():
+    tracker = VideoTracker()
+
+    with gr.Blocks() as interface:
+        gr.Markdown("# Object Tracking with SAM2")
+        gr.Markdown("Upload a video and click on objects to track them")
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                video_input = gr.Video(label="Input Video")
+                image_output = gr.Image(label="Current Frame", interactive=True)
+                frame_slider = gr.Slider(
+                    minimum=0,
+                    maximum=0,
+                    step=1,
+                    value=0,
+                    label="Frame Selection",
+                    interactive=True,
+                )
+
+            with gr.Column(scale=1):
+                propagate_btn = gr.Button("Propagate Through Video", variant="primary")
+                video_output = gr.Video(label="Output Video")
+
+        video_input.change(
+            fn=tracker.load_video,
+            inputs=[video_input],
+            outputs=[image_output, frame_slider],
+        )
+
+        frame_slider.change(
+            fn=tracker.update_frame,
+            inputs=[frame_slider],
+            outputs=[image_output],
+        )
+
+        image_output.select(
+            fn=tracker.add_point,
+            inputs=[image_output],
+            outputs=[image_output],
+        )
+
+        propagate_btn.click(
+            fn=tracker.propagate_video,
+            inputs=[],
+            outputs=[video_output],
+        )
+
+    return interface
+
+
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch(share=True)

hello.py

@@ -0,0 +1,6 @@
+def main():
+    print("Hello from golf-tracking!")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,21 @@
+[project]
+name = "golf-tracking"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "decord>=0.6.0",
+    "gradio>=4.0.0",
+    "hydra-core>=1.3.2",
+    "iopath>=0.1.10",
+    "matplotlib>=3.9.1",
+    "moviepy==1.0.3",
+    "numpy>=1.24.4",
+    "opencv-python>=4.7.0",
+    "pillow>=9.4.0",
+    "sam2>=1.1.0",
+    "scipy>=1.7.0",
+    "torch>=2.0.0",
+    "tqdm>=4.66.1",
+]

ref.py

@@ -0,0 +1,265 @@
+import gradio as gr
+from gradio_image_prompter import ImagePrompter
+import torch
+import numpy as np
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+from uuid import uuid4
+import os
+from huggingface_hub import upload_folder, login
+from PIL import Image as PILImage
+from datasets import Dataset, Features, Array2D, Image
+import shutil
+import random
+from datasets import load_dataset
+
+MODEL = "facebook/sam2-hiera-large"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+PREDICTOR = SAM2ImagePredictor.from_pretrained(MODEL, device=DEVICE)
+
+DESTINATION_DS = "amaye15/object-segmentation"
+
+
+token = os.getenv("TOKEN")
+if token:
+    login(token)
+
+IMAGE = None
+MASKS = None
+MASKED_IMAGES = None
+INDEX = None
+
+
+ds_name = ["amaye15/product_labels"]  #  "amaye15/Products-10k", "amaye15/receipts"
+choices = ["test", "train"]
+max_len = None
+
+ds_stream = load_dataset(random.choice(ds_name), streaming=True)
+
+
+ds_split = ds_stream[random.choice(choices)]
+
+ds_iter = ds_split.iter(batch_size=1)
+
+for idx, val in enumerate(ds_iter):
+    max_len = idx
+
+
+def prompter(prompts):
+    image = np.array(prompts["image"])  # Convert the image to a numpy array
+    points = prompts["points"]  # Get the points from prompts
+
+    # Perform inference with multimask_output=True
+    with torch.inference_mode():
+        PREDICTOR.set_image(image)
+        input_point = [[point[0], point[1]] for point in points]
+        input_label = [1] * len(points)  # Assuming all points are foreground
+        masks, _, _ = PREDICTOR.predict(
+            point_coords=input_point, point_labels=input_label, multimask_output=True
+        )
+
+    # Prepare individual images with separate overlays
+    overlay_images = []
+    for i, mask in enumerate(masks):
+        print(f"Predicted Mask {i+1}:", mask.shape)
+        red_mask = np.zeros_like(image)
+        red_mask[:, :, 0] = mask.astype(np.uint8) * 255  # Apply the red channel
+        red_mask = PILImage.fromarray(red_mask)
+
+        # Convert the original image to a PIL image
+        original_image = PILImage.fromarray(image)
+
+        # Blend the original image with the red mask
+        blended_image = PILImage.blend(original_image, red_mask, alpha=0.5)
+
+        # Add the blended image to the list
+        overlay_images.append(blended_image)
+
+    global IMAGE, MASKS, MASKED_IMAGES
+    IMAGE, MASKS = image, masks
+    MASKED_IMAGES = [np.array(img) for img in overlay_images]
+
+    return overlay_images[0], overlay_images[1], overlay_images[2], masks
+
+
+def select_mask(
+    selected_mask_index,
+    mask1,
+    mask2,
+    mask3,
+):
+    masks = [mask1, mask2, mask3]
+    global INDEX
+    INDEX = selected_mask_index
+    return masks[selected_mask_index]
+
+
+def save_selected_mask(image, mask, output_dir="output"):
+    output_dir = os.path.join(os.getcwd(), output_dir)
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    folder_id = str(uuid4())
+
+    folder_path = os.path.join(output_dir, folder_id)
+
+    os.makedirs(folder_path, exist_ok=True)
+
+    data_path = os.path.join(folder_path, "data.parquet")
+
+    data = {
+        "image": IMAGE,
+        "masked_image": MASKED_IMAGES[INDEX],
+        "mask": MASKS[INDEX],
+    }
+
+    features = Features(
+        {
+            "image": Image(),
+            "masked_image": Image(),
+            "mask": Array2D(
+                dtype="int64", shape=(MASKS[INDEX].shape[0], MASKS[INDEX].shape[1])
+            ),
+        }
+    )
+
+    ds = Dataset.from_list([data], features=features)
+    ds.to_parquet(data_path)
+
+    upload_folder(
+        folder_path=output_dir,
+        repo_id=DESTINATION_DS,
+        repo_type="dataset",
+    )
+
+    shutil.rmtree(folder_path)
+
+    iframe_code = """## Success! 🎉🤖✅
+You've successfully contributed to the dataset.
+Please note that because new data has been added to the dataset, it may take a couple of minutes to render.
+Check it out here:
+[Object Segmentation Dataset](https://huggingface.co/datasets/amaye15/object-segmentation)
+"""
+
+    return iframe_code
+
+
+def get_random_image():
+    """Get a random image from the dataset."""
+    global max_len
+    random_idx = random.choice(range(max_len))
+    image_data = list(ds_split.skip(random_idx).take(1))[0]["pixel_values"]
+    formatted_image = {
+        "image": np.array(image_data),
+        "points": [],
+    }  # Create the correct format
+    return formatted_image
+
+
+# Define the Gradio Blocks app
+with gr.Blocks() as demo:
+    gr.Markdown("# Object Segmentation- Image Point Collector and Mask Overlay Tool")
+    gr.Markdown(
+        """
+        This application utilizes **Segment Anything V2 (SAM2)** to allow you to upload an image or select a random image from a dataset and interactively generate segmentation masks based on multiple points you select on the image.
+        ### How It Works:
+        1. **Upload or Select an Image**: You can either upload your own image or use a random image from the dataset.
+        2. **Point Selection**: Click on the image to indicate points of interest. You can add multiple points, and these will be used collectively to generate segmentation masks using SAM2.
+        3. **Mask Generation**: The app will generate up to three different segmentation masks for the selected points, each displayed separately with a red overlay.
+        4. **Mask Selection**: Carefully review the generated masks and select the one that best fits your needs. **It's important to choose the correct mask, as your selection will be saved and used for further processing.**
+        5. **Save and Contribute**: Save the selected mask along with the image to a dataset, contributing to a shared dataset on Hugging Face.
+        **Disclaimer**: All images and masks you work with will be collected and stored in a public dataset. Please ensure that you are comfortable with your selections and the data you provide before saving.
+
+        This tool is particularly useful for creating precise object segmentation masks for computer vision tasks, such as training models or generating labeled datasets.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.State()
+            # Input: ImagePrompter for uploaded image
+            upload_image_input = ImagePrompter(show_label=False)
+
+            random_image_button = gr.Button("Use Random Image")
+
+            submit_button = gr.Button("Submit")
+
+    with gr.Row():
+        with gr.Column():
+            # Outputs: Up to 3 overlay images
+            image_output_1 = gr.Image(show_label=False)
+        with gr.Column():
+            image_output_2 = gr.Image(show_label=False)
+        with gr.Column():
+            image_output_3 = gr.Image(show_label=False)
+
+    # Dropdown for selecting the correct mask
+    with gr.Row():
+        mask_selector = gr.Radio(
+            label="Select the correct mask",
+            choices=["Mask 1", "Mask 2", "Mask 3"],
+            type="index",
+        )
+        # selected_mask_output = gr.Image(show_label=False)
+
+    save_button = gr.Button("Save Selected Mask and Image")
+    iframe_display = gr.Markdown()
+
+    # Logic for the random image button
+    random_image_button.click(
+        fn=get_random_image,
+        inputs=None,
+        outputs=upload_image_input,  # Pass the formatted random image to ImagePrompter
+    )
+
+    # Logic to use uploaded image
+    upload_image_input.change(
+        fn=lambda img: img, inputs=upload_image_input, outputs=image_input
+    )
+    # Define the action triggered by the submit button
+    submit_button.click(
+        fn=prompter,
+        inputs=upload_image_input,  # The final image input (whether uploaded or random)
+        outputs=[image_output_1, image_output_2, image_output_3, gr.State()],
+        show_progress=True,
+    )
+
+    # Define the action triggered by mask selection
+    mask_selector.change(
+        fn=select_mask,
+        inputs=[mask_selector, image_output_1, image_output_2, image_output_3],
+        outputs=gr.State(),
+    )
+
+    # Define the action triggered by the save button
+    save_button.click(
+        fn=save_selected_mask,
+        inputs=[gr.State(), gr.State()],
+        outputs=iframe_display,
+        show_progress=True,
+    )
+
+# Launch the Gradio app
+demo.launch()
+
+
+# with gr.Column():
+#     source = gr.Textbox(label="Source Dataset")
+#     source_display = gr.Markdown()
+#     iframe_display = gr.HTML()
+
+#     source.change(
+#         save_dataset_name,
+#         inputs=(gr.State("source_dataset"), source),
+#         outputs=(source_display, iframe_display),
+#     )
+
+# with gr.Column():
+
+#     destination = gr.Textbox(label="Destination Dataset")
+#     destination_display = gr.Markdown()
+
+#     destination.change(
+#         save_dataset_name,
+#         inputs=(gr.State("destination_dataset"), destination),
+#         outputs=destination_display,
+#     )

reference_code.py

@@ -0,0 +1,776 @@
+import subprocess
+import re
+from typing import List, Tuple, Optional
+import spaces
+
+# Define the command to be executed
+command = ["python", "setup.py", "build_ext", "--inplace"]
+
+# Execute the command
+result = subprocess.run(command, capture_output=True, text=True)
+
+css = """
+div#component-18, div#component-25, div#component-35, div#component-41{
+    align-items: stretch!important;
+}
+"""
+
+# Print the output and error (if any)
+print("Output:\n", result.stdout)
+print("Errors:\n", result.stderr)
+
+# Check if the command was successful
+if result.returncode == 0:
+    print("Command executed successfully.")
+else:
+    print("Command failed with return code:", result.returncode)
+
+import gradio as gr
+from datetime import datetime
+import os
+
+os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "1"
+import torch
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+from PIL import Image, ImageFilter
+from sam2.build_sam import build_sam2_video_predictor
+
+from moviepy.editor import ImageSequenceClip
+
+
+def sparse_sampling(jpeg_images, original_fps, target_fps=6):
+    # Calculate the frame interval for sampling based on the target fps
+    frame_interval = int(original_fps // target_fps)
+
+    # Sparse sample the jpeg_images by selecting every 'frame_interval' frame
+    sampled_images = [
+        jpeg_images[i] for i in range(0, len(jpeg_images), frame_interval)
+    ]
+
+    return sampled_images
+
+
+def get_video_fps(video_path):
+    # Open the video file
+    cap = cv2.VideoCapture(video_path)
+
+    if not cap.isOpened():
+        print("Error: Could not open video.")
+        return None
+
+    # Get the FPS of the video
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    return fps
+
+
+def clear_points(image):
+    # we clean all
+    return [
+        image,  # first_frame_path
+        [],  # tracking_points
+        [],  # trackings_input_label
+        image,  # points_map
+        # gr.State()     # stored_inference_state
+    ]
+
+
+def preprocess_video_in(video_path):
+    # Generate a unique ID based on the current date and time
+    unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
+
+    # Set directory with this ID to store video frames
+    extracted_frames_output_dir = f"frames_{unique_id}"
+
+    # Create the output directory
+    os.makedirs(extracted_frames_output_dir, exist_ok=True)
+
+    ### Process video frames ###
+    # Open the video file
+    cap = cv2.VideoCapture(video_path)
+
+    if not cap.isOpened():
+        print("Error: Could not open video.")
+        return None
+
+    # Get the frames per second (FPS) of the video
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    # Calculate the number of frames to process (60 seconds of video)
+    max_frames = int(fps * 60)
+
+    frame_number = 0
+    first_frame = None
+
+    while True:
+        ret, frame = cap.read()
+        if not ret or frame_number >= max_frames:
+            break
+        if frame_number % 6 == 0:
+            # Format the frame filename as '00000.jpg'
+            frame_filename = os.path.join(
+                extracted_frames_output_dir, f"{frame_number:05d}.jpg"
+            )
+
+            # Save the frame as a JPEG file
+            cv2.imwrite(frame_filename, frame)
+
+        # Store the first frame
+        if frame_number == 0:
+            first_frame = frame_filename
+
+        frame_number += 1
+
+    # Release the video capture object
+    cap.release()
+
+    # scan all the JPEG frame names in this directory
+    scanned_frames = [
+        p
+        for p in os.listdir(extracted_frames_output_dir)
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+    ]
+    scanned_frames.sort(key=lambda p: int(os.path.splitext(p)[0]))
+    # print(f"SCANNED_FRAMES: {scanned_frames}")
+
+    return [
+        first_frame,  # first_frame_path
+        [],  # tracking_points
+        [],  # trackings_input_label
+        first_frame,  # input_first_frame_image
+        first_frame,  # points_map
+        extracted_frames_output_dir,  # video_frames_dir
+        scanned_frames,  # scanned_frames
+        None,  # stored_inference_state
+        None,  # stored_frame_names
+        gr.update(open=False),  # video_in_drawer
+    ]
+
+
+def get_point(
+    point_type,
+    tracking_points,
+    trackings_input_label,
+    input_first_frame_image,
+    evt: gr.SelectData,
+):
+    print(f"You selected {evt.value} at {evt.index} from {evt.target}")
+
+    tracking_points.append(evt.index)
+    # tracking_points.value.append(evt.index)
+    print(f"TRACKING POINT: {tracking_points}")
+
+    if point_type == "include":
+        trackings_input_label.append(1)
+        # trackings_input_label.value.append(1)
+    elif point_type == "exclude":
+        trackings_input_label.append(0)
+        # trackings_input_label.value.append(0)
+    print(f"TRACKING INPUT LABEL: {trackings_input_label}")
+
+    # Open the image and get its dimensions
+    transparent_background = Image.open(input_first_frame_image).convert("RGBA")
+    w, h = transparent_background.size
+
+    # Define the circle radius as a fraction of the smaller dimension
+    fraction = 0.02  # You can adjust this value as needed
+    radius = int(fraction * min(w, h))
+
+    # Create a transparent layer to draw on
+    transparent_layer = np.zeros((h, w, 4), dtype=np.uint8)
+
+    for index, track in enumerate(tracking_points):
+        if trackings_input_label[index] == 1:
+            cv2.circle(transparent_layer, track, radius, (0, 255, 0, 255), -1)
+        else:
+            cv2.circle(transparent_layer, track, radius, (255, 0, 0, 255), -1)
+
+    # Convert the transparent layer back to an image
+    transparent_layer = Image.fromarray(transparent_layer, "RGBA")
+    selected_point_map = Image.alpha_composite(
+        transparent_background, transparent_layer
+    )
+
+    return tracking_points, trackings_input_label, selected_point_map
+
+
+def show_mask(mask, ax, obj_id=None, random_color=False):
+    if random_color:
+        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+    else:
+        cmap = plt.get_cmap("tab10")
+        cmap_idx = 0 if obj_id is None else obj_id
+        color = np.array([*cmap(cmap_idx)[:3], 0.6])
+    h, w = mask.shape[-2:]
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+    ax.imshow(mask_image)
+
+
+def show_points(coords, labels, ax, marker_size=200):
+    pos_points = coords[labels == 1]
+    neg_points = coords[labels == 0]
+    ax.scatter(
+        pos_points[:, 0],
+        pos_points[:, 1],
+        color="green",
+        marker="*",
+        s=marker_size,
+        edgecolor="white",
+        linewidth=1.25,
+    )
+    ax.scatter(
+        neg_points[:, 0],
+        neg_points[:, 1],
+        color="red",
+        marker="*",
+        s=marker_size,
+        edgecolor="white",
+        linewidth=1.25,
+    )
+
+
+def show_box(box, ax):
+    x0, y0 = box[0], box[1]
+    w, h = box[2] - box[0], box[3] - box[1]
+    ax.add_patch(
+        plt.Rectangle((x0, y0), w, h, edgecolor="green", facecolor=(0, 0, 0, 0), lw=2)
+    )
+
+
+def load_model(checkpoint):
+    # Load model accordingly to user's choice
+    if checkpoint == "tiny":
+        sam2_checkpoint = "./checkpoints/sam2.1_hiera_tiny.pt"
+        model_cfg = "configs/sam2.1/sam2.1_hiera_t.yaml"
+        return [sam2_checkpoint, model_cfg]
+    elif checkpoint == "samll":
+        sam2_checkpoint = "./checkpoints/sam2.1_hiera_small.pt"
+        model_cfg = "configs/sam2.1/sam2.1_hiera_s.yaml"
+        return [sam2_checkpoint, model_cfg]
+    elif checkpoint == "base-plus":
+        sam2_checkpoint = "./checkpoints/sam2.1_hiera_base_plus.pt"
+        model_cfg = "configs/sam2.1/sam2.1_hiera_b+.yaml"
+        return [sam2_checkpoint, model_cfg]
+    # elif checkpoint == "large":
+    #     sam2_checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
+    #     model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
+    #     return [sam2_checkpoint, model_cfg]
+
+
+def get_mask_sam_process(
+    stored_inference_state,
+    input_first_frame_image,
+    checkpoint,
+    tracking_points,
+    trackings_input_label,
+    video_frames_dir,  # extracted_frames_output_dir defined in 'preprocess_video_in' function
+    scanned_frames,
+    working_frame: str = None,  # current frame being added points
+    available_frames_to_check: List[str] = [],
+    # progress=gr.Progress(track_tqdm=True)
+):
+    # get model and model config paths
+    print(f"USER CHOSEN CHECKPOINT: {checkpoint}")
+    sam2_checkpoint, model_cfg = load_model(checkpoint)
+    print("MODEL LOADED")
+
+    # set predictor
+
+    predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cpu")
+
+    print("PREDICTOR READY")
+
+    # `video_dir` a directory of JPEG frames with filenames like `<frame_index>.jpg`
+    # print(f"STATE FRAME OUTPUT DIRECTORY: {video_frames_dir}")
+    video_dir = video_frames_dir
+
+    # scan all the JPEG frame names in this directory
+    frame_names = scanned_frames
+
+    # print(f"STORED INFERENCE STEP: {stored_inference_state}")
+    if stored_inference_state is None:
+        # Init SAM2 inference_state
+        inference_state = predictor.init_state(video_path=video_dir)
+        inference_state["num_pathway"] = 3
+        inference_state["iou_thre"] = 0.3
+        inference_state["uncertainty"] = 2
+        print("NEW INFERENCE_STATE INITIATED")
+    else:
+        inference_state = stored_inference_state
+
+    inference_state["device"] = "cpu"
+
+    # segment and track one object
+    # predictor.reset_state(inference_state) # if any previous tracking, reset
+
+    ### HANDLING WORKING FRAME
+    # new_working_frame = None
+    # Add new point
+    if working_frame is None:
+        ann_frame_idx = (
+            0  # the frame index we interact with, 0 if it is the first frame
+        )
+        working_frame = "00000.jpg"
+    else:
+        # Use a regular expression to find the integer
+        match = re.search(r"frame_(\d+)", working_frame)
+        if match:
+            # Extract the integer from the match
+            frame_number = int(match.group(1))
+            ann_frame_idx = frame_number
+
+    print(f"NEW_WORKING_FRAME PATH: {working_frame}")
+
+    ann_obj_id = (
+        1  # give a unique id to each object we interact with (it can be any integers)
+    )
+
+    # Let's add a positive click at (x, y) = (210, 350) to get started
+    points = np.array(tracking_points, dtype=np.float32)
+    # for labels, `1` means positive click and `0` means negative click
+    labels = np.array(trackings_input_label, np.int32)
+
+    _, out_obj_ids, out_mask_logits = predictor.add_new_points(
+        inference_state=inference_state,
+        frame_idx=ann_frame_idx,
+        obj_id=ann_obj_id,
+        points=points,
+        labels=labels,
+    )
+
+    # Create the plot
+    plt.figure(figsize=(12, 8))
+    plt.title(f"frame {ann_frame_idx}")
+    plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))
+    show_points(points, labels, plt.gca())
+    show_mask(
+        (out_mask_logits[0] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_ids[0]
+    )
+
+    # Save the plot as a JPG file
+    first_frame_output_filename = "output_first_frame.jpg"
+    plt.savefig(first_frame_output_filename, format="jpg")
+    plt.close()
+    # torch.cuda.empty_cache()
+
+    # Assuming available_frames_to_check.value is a list
+    if working_frame not in available_frames_to_check:
+        available_frames_to_check.append(working_frame)
+        print(available_frames_to_check)
+
+    # return gr.update(visible=True), "output_first_frame.jpg", frame_names, predictor, inference_state, gr.update(choices=available_frames_to_check, value=working_frame, visible=True)
+    return (
+        "output_first_frame.jpg",
+        frame_names,
+        predictor,
+        inference_state,
+        gr.update(
+            choices=available_frames_to_check, value=working_frame, visible=False
+        ),
+    )
+
+
+@spaces.GPU
+def propagate_to_all(
+    video_in,
+    checkpoint,
+    stored_inference_state,
+    stored_frame_names,
+    video_frames_dir,
+    vis_frame_type,
+    available_frames_to_check,
+    working_frame,
+    progress=gr.Progress(track_tqdm=True),
+):
+    # use bfloat16 for the entire notebook
+    torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
+
+    if torch.cuda.get_device_properties(0).major >= 8:
+        # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+    #### PROPAGATION ####
+    sam2_checkpoint, model_cfg = load_model(checkpoint)
+    # set predictor
+
+    inference_state = stored_inference_state
+
+    if torch.cuda.is_available():
+        inference_state["device"] = "cuda"
+        predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
+    else:
+        inference_state["device"] = "cpu"
+        predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cpu")
+
+    frame_names = stored_frame_names
+    video_dir = video_frames_dir
+
+    # Define a directory to save the JPEG images
+    frames_output_dir = "frames_output_images"
+    os.makedirs(frames_output_dir, exist_ok=True)
+
+    # Initialize a list to store file paths of saved images
+    jpeg_images = []
+
+    # run propagation throughout the video and collect the results in a dict
+    video_segments = {}  # video_segments contains the per-frame segmentation results
+    # for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
+    #     video_segments[out_frame_idx] = {
+    #         out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
+    #         for i, out_obj_id in enumerate(out_obj_ids)
+    #     }
+
+    out_obj_ids, out_mask_logits = predictor.propagate_in_video(
+        inference_state,
+        start_frame_idx=0,
+        reverse=False,
+    )
+    print(out_obj_ids)
+    for frame_idx in range(0, inference_state["num_frames"]):
+        video_segments[frame_idx] = {
+            out_obj_ids[0]: (out_mask_logits[frame_idx] > 0.0).cpu().numpy()
+        }
+        # output_scores_per_object[object_id][frame_idx] = out_mask_logits[frame_idx].cpu().numpy()
+
+    # render the segmentation results every few frames
+    if vis_frame_type == "check":
+        vis_frame_stride = 15
+    elif vis_frame_type == "render":
+        vis_frame_stride = 1
+
+    plt.close("all")
+    for out_frame_idx in range(0, len(frame_names), vis_frame_stride):
+        plt.figure(figsize=(6, 4))
+        plt.title(f"frame {out_frame_idx}")
+        plt.imshow(Image.open(os.path.join(video_dir, frame_names[out_frame_idx])))
+        for out_obj_id, out_mask in video_segments[out_frame_idx].items():
+            show_mask(out_mask, plt.gca(), obj_id=out_obj_id)
+
+        # Define the output filename and save the figure as a JPEG file
+        output_filename = os.path.join(frames_output_dir, f"frame_{out_frame_idx}.jpg")
+        plt.savefig(output_filename, format="jpg")
+
+        # Close the plot
+        plt.close()
+
+        # Append the file path to the list
+        jpeg_images.append(output_filename)
+
+        if f"frame_{out_frame_idx}.jpg" not in available_frames_to_check:
+            available_frames_to_check.append(f"frame_{out_frame_idx}.jpg")
+
+    torch.cuda.empty_cache()
+    print(f"JPEG_IMAGES: {jpeg_images}")
+
+    if vis_frame_type == "check":
+        return (
+            gr.update(value=jpeg_images),
+            gr.update(value=None),
+            gr.update(
+                choices=available_frames_to_check, value=working_frame, visible=True
+            ),
+            available_frames_to_check,
+            gr.update(visible=True),
+        )
+    elif vis_frame_type == "render":
+        # Create a video clip from the image sequence
+        original_fps = get_video_fps(video_in)
+        # sampled_images = sparse_sampling(jpeg_images, original_fps, target_fps=6)
+        clip = ImageSequenceClip(jpeg_images, fps=original_fps // 6)
+        # clip = ImageSequenceClip(jpeg_images, fps=fps)
+        # Write the result to a file
+        final_vid_output_path = "output_video.mp4"
+
+        # Write the result to a file
+        clip.write_videofile(final_vid_output_path, codec="libx264")
+
+        return (
+            gr.update(value=None),
+            gr.update(value=final_vid_output_path),
+            working_frame,
+            available_frames_to_check,
+            gr.update(visible=True),
+        )
+
+
+def update_ui(vis_frame_type):
+    if vis_frame_type == "check":
+        return gr.update(visible=True), gr.update(visible=False)
+    elif vis_frame_type == "render":
+        return gr.update(visible=False), gr.update(visible=True)
+
+
+def switch_working_frame(working_frame, scanned_frames, video_frames_dir):
+    new_working_frame = None
+    if working_frame == None:
+        new_working_frame = os.path.join(video_frames_dir, scanned_frames[0])
+
+    else:
+        # Use a regular expression to find the integer
+        match = re.search(r"frame_(\d+)", working_frame)
+        if match:
+            # Extract the integer from the match
+            frame_number = int(match.group(1))
+            ann_frame_idx = frame_number
+            new_working_frame = os.path.join(
+                video_frames_dir, scanned_frames[ann_frame_idx]
+            )
+    return gr.State([]), gr.State([]), new_working_frame, new_working_frame
+
+
+def reset_propagation(first_frame_path, predictor, stored_inference_state):
+    predictor.reset_state(stored_inference_state)
+    # print(f"RESET State: {stored_inference_state} ")
+    return (
+        first_frame_path,
+        [],
+        [],
+        gr.update(value=None, visible=False),
+        stored_inference_state,
+        None,
+        ["frame_0.jpg"],
+        first_frame_path,
+        "frame_0.jpg",
+        gr.update(visible=False),
+    )
+
+
+with gr.Blocks(css=css) as demo:
+    first_frame_path = gr.State()
+    tracking_points = gr.State([])
+    trackings_input_label = gr.State([])
+    video_frames_dir = gr.State()
+    scanned_frames = gr.State()
+    loaded_predictor = gr.State()
+    stored_inference_state = gr.State()
+    stored_frame_names = gr.State()
+    available_frames_to_check = gr.State([])
+    with gr.Column():
+        gr.Markdown(
+            """
+            <h1 style="text-align: center;">🔥 SAM2Long Demo 🔥</h1>
+            """
+        )
+        gr.Markdown(
+            """
+            This is a simple demo for video segmentation with [SAM2Long](https://github.com/Mark12Ding/SAM2Long).
+            """
+        )
+        gr.Markdown(
+            """
+            ### 📋 Instructions:
+            It is largely built on the [SAM2-Video-Predictor](https://huggingface.co/spaces/fffiloni/SAM2-Video-Predictor).
+            1. **Upload your video** [MP4-24fps]
+            2. With **'include' point type** selected, click on the object to mask on the first frame
+            3. Switch to **'exclude' point type** if you want to specify an area to avoid
+            4. **Get Mask!**
+            5. **Check Propagation** every 15 frames
+            6. **Propagate with "render"** to render the final masked video
+            7. **Hit Reset** button if you want to refresh and start again
+
+            *Note: Input video will be processed for up to 60 seconds only for demo purposes.*
+            """
+        )
+        with gr.Row():
+            with gr.Column():
+                with gr.Group():
+                    with gr.Group():
+                        with gr.Row():
+                            point_type = gr.Radio(
+                                label="point type",
+                                choices=["include", "exclude"],
+                                value="include",
+                                scale=2,
+                            )
+                            clear_points_btn = gr.Button("Clear Points", scale=1)
+
+                    input_first_frame_image = gr.Image(
+                        label="input image",
+                        interactive=False,
+                        type="filepath",
+                        visible=False,
+                    )
+
+                    points_map = gr.Image(
+                        label="Point n Click map", type="filepath", interactive=False
+                    )
+
+                    with gr.Group():
+                        with gr.Row():
+                            checkpoint = gr.Dropdown(
+                                label="Checkpoint",
+                                choices=["tiny", "small", "base-plus"],
+                                value="tiny",
+                            )
+                            submit_btn = gr.Button("Get Mask", size="lg")
+
+                with gr.Accordion("Your video IN", open=True) as video_in_drawer:
+                    video_in = gr.Video(label="Video IN", format="mp4")
+
+                gr.HTML("""
+
+                <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}?duplicate=true">
+                    <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg-dark.svg" alt="Duplicate this Space" />
+                </a> to skip queue and avoid OOM errors from heavy public load
+                """)
+
+            with gr.Column():
+                with gr.Group():
+                    # with gr.Group():
+                    # with gr.Row():
+                    working_frame = gr.Dropdown(
+                        label="working frame ID",
+                        choices=["frame_0.jpg"],
+                        value="frame_0.jpg",
+                        visible=False,
+                        allow_custom_value=False,
+                        interactive=True,
+                    )
+                    #     change_current = gr.Button("change current", visible=False)
+                    # working_frame = []
+                    output_result = gr.Image(label="current working mask ref")
+                with gr.Group():
+                    with gr.Row():
+                        vis_frame_type = gr.Radio(
+                            label="Propagation level",
+                            choices=["check", "render"],
+                            value="check",
+                            scale=2,
+                        )
+                        propagate_btn = gr.Button("Propagate", scale=2)
+
+                reset_prpgt_brn = gr.Button("Reset", visible=False)
+                output_propagated = gr.Gallery(
+                    label="Propagated Mask samples gallery", columns=4, visible=False
+                )
+                output_video = gr.Video(visible=False)
+                # output_result_mask = gr.Image()
+
+    # When new video is uploaded
+    video_in.upload(
+        fn=preprocess_video_in,
+        inputs=[video_in],
+        outputs=[
+            first_frame_path,
+            tracking_points,  # update Tracking Points in the gr.State([]) object
+            trackings_input_label,  # update Tracking Labels in the gr.State([]) object
+            input_first_frame_image,  # hidden component used as ref when clearing points
+            points_map,  # Image component where we add new tracking points
+            video_frames_dir,  # Array where frames from video_in are deep stored
+            scanned_frames,  # Scanned frames by SAM2
+            stored_inference_state,  # Sam2 inference state
+            stored_frame_names,  #
+            video_in_drawer,  # Accordion to hide uploaded video player
+        ],
+        queue=False,
+    )
+
+    # triggered when we click on image to add new points
+    points_map.select(
+        fn=get_point,
+        inputs=[
+            point_type,  # "include" or "exclude"
+            tracking_points,  # get tracking_points values
+            trackings_input_label,  # get tracking label values
+            input_first_frame_image,  # gr.State() first frame path
+        ],
+        outputs=[
+            tracking_points,  # updated with new points
+            trackings_input_label,  # updated with corresponding labels
+            points_map,  # updated image with points
+        ],
+        queue=False,
+    )
+
+    # Clear every points clicked and added to the map
+    clear_points_btn.click(
+        fn=clear_points,
+        inputs=input_first_frame_image,  # we get the untouched hidden image
+        outputs=[
+            first_frame_path,
+            tracking_points,
+            trackings_input_label,
+            points_map,
+            # stored_inference_state,
+        ],
+        queue=False,
+    )
+
+    # change_current.click(
+    #     fn = switch_working_frame,
+    #     inputs = [working_frame, scanned_frames, video_frames_dir],
+    #     outputs = [tracking_points, trackings_input_label, input_first_frame_image, points_map],
+    #     queue=False
+    # )
+
+    submit_btn.click(
+        fn=get_mask_sam_process,
+        inputs=[
+            stored_inference_state,
+            input_first_frame_image,
+            checkpoint,
+            tracking_points,
+            trackings_input_label,
+            video_frames_dir,
+            scanned_frames,
+            working_frame,
+            available_frames_to_check,
+        ],
+        outputs=[
+            output_result,
+            stored_frame_names,
+            loaded_predictor,
+            stored_inference_state,
+            working_frame,
+        ],
+        queue=False,
+    )
+
+    reset_prpgt_brn.click(
+        fn=reset_propagation,
+        inputs=[first_frame_path, loaded_predictor, stored_inference_state],
+        outputs=[
+            points_map,
+            tracking_points,
+            trackings_input_label,
+            output_propagated,
+            stored_inference_state,
+            output_result,
+            available_frames_to_check,
+            input_first_frame_image,
+            working_frame,
+            reset_prpgt_brn,
+        ],
+        queue=False,
+    )
+
+    propagate_btn.click(
+        fn=update_ui,
+        inputs=[vis_frame_type],
+        outputs=[output_propagated, output_video],
+        queue=False,
+    ).then(
+        fn=propagate_to_all,
+        inputs=[
+            video_in,
+            checkpoint,
+            stored_inference_state,
+            stored_frame_names,
+            video_frames_dir,
+            vis_frame_type,
+            available_frames_to_check,
+            working_frame,
+        ],
+        outputs=[
+            output_propagated,
+            output_video,
+            working_frame,
+            available_frames_to_check,
+            reset_prpgt_brn,
+        ],
+    )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+gradio>=5.0.0
+torch>=2.0.0
+numpy>=1.24.0
+opencv-python>=4.8.0
+decord>=0.6.0

sam2-repo/.clang-format

@@ -0,0 +1,85 @@
+AccessModifierOffset: -1
+AlignAfterOpenBracket: AlwaysBreak
+AlignConsecutiveAssignments: false
+AlignConsecutiveDeclarations: false
+AlignEscapedNewlinesLeft: true
+AlignOperands:   false
+AlignTrailingComments: false
+AllowAllParametersOfDeclarationOnNextLine: false
+AllowShortBlocksOnASingleLine: false
+AllowShortCaseLabelsOnASingleLine: false
+AllowShortFunctionsOnASingleLine: Empty
+AllowShortIfStatementsOnASingleLine: false
+AllowShortLoopsOnASingleLine: false
+AlwaysBreakAfterReturnType: None
+AlwaysBreakBeforeMultilineStrings: true
+AlwaysBreakTemplateDeclarations: true
+BinPackArguments: false
+BinPackParameters: false
+BraceWrapping:
+  AfterClass:      false
+  AfterControlStatement: false
+  AfterEnum:       false
+  AfterFunction:   false
+  AfterNamespace:  false
+  AfterObjCDeclaration: false
+  AfterStruct:     false
+  AfterUnion:      false
+  BeforeCatch:     false
+  BeforeElse:      false
+  IndentBraces:    false
+BreakBeforeBinaryOperators: None
+BreakBeforeBraces: Attach
+BreakBeforeTernaryOperators: true
+BreakConstructorInitializersBeforeComma: false
+BreakAfterJavaFieldAnnotations: false
+BreakStringLiterals: false
+ColumnLimit:     80
+CommentPragmas:  '^ IWYU pragma:'
+ConstructorInitializerAllOnOneLineOrOnePerLine: true
+ConstructorInitializerIndentWidth: 4
+ContinuationIndentWidth: 4
+Cpp11BracedListStyle: true
+DerivePointerAlignment: false
+DisableFormat:   false
+ForEachMacros:   [ FOR_EACH, FOR_EACH_R, FOR_EACH_RANGE, ]
+IncludeCategories:
+  - Regex:           '^<.*\.h(pp)?>'
+    Priority:        1
+  - Regex:           '^<.*'
+    Priority:        2
+  - Regex:           '.*'
+    Priority:        3
+IndentCaseLabels: true
+IndentWidth:     2
+IndentWrappedFunctionNames: false
+KeepEmptyLinesAtTheStartOfBlocks: false
+MacroBlockBegin: ''
+MacroBlockEnd:   ''
+MaxEmptyLinesToKeep: 1
+NamespaceIndentation: None
+ObjCBlockIndentWidth: 2
+ObjCSpaceAfterProperty: false
+ObjCSpaceBeforeProtocolList: false
+PenaltyBreakBeforeFirstCallParameter: 1
+PenaltyBreakComment: 300
+PenaltyBreakFirstLessLess: 120
+PenaltyBreakString: 1000
+PenaltyExcessCharacter: 1000000
+PenaltyReturnTypeOnItsOwnLine: 200
+PointerAlignment: Left
+ReflowComments:  true
+SortIncludes:    true
+SpaceAfterCStyleCast: false
+SpaceBeforeAssignmentOperators: true
+SpaceBeforeParens: ControlStatements
+SpaceInEmptyParentheses: false
+SpacesBeforeTrailingComments: 1
+SpacesInAngles:  false
+SpacesInContainerLiterals: true
+SpacesInCStyleCastParentheses: false
+SpacesInParentheses: false
+SpacesInSquareBrackets: false
+Standard:        Cpp11
+TabWidth:        8
+UseTab:          Never

sam2-repo/.github/workflows/check_fmt.yml

@@ -0,0 +1,17 @@
+name: SAM2/fmt
+on:
+  pull_request:
+    branches:
+    - main
+jobs:
+  ufmt_check:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check formatting
+        uses: omnilib/ufmt@action-v1
+        with:
+          path: sam2 tools
+          version: "2.0.0b2"
+          python-version: "3.10"
+          black-version: "24.2.0"
+          usort-version: "1.0.2"

sam2-repo/.gitignore

@@ -0,0 +1,11 @@
+.vscode/
+.DS_Store
+__pycache__/
+*-checkpoint.ipynb
+.venv
+*.egg*
+build/*
+_C.*
+outputs/*
+checkpoints/*.pt
+demo/backend/checkpoints/*.pt

sam2-repo/.watchmanconfig

@@ -0,0 +1 @@
+{}

sam2-repo/CODE_OF_CONDUCT.md

@@ -0,0 +1,80 @@
+# Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to make participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, sex characteristics, gender identity and expression,
+level of experience, education, socio-economic status, nationality, personal
+appearance, race, religion, or sexual identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+  advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies within all project spaces, and it also applies when
+an individual is representing the project or its community in public spaces.
+Examples of representing a project or community include using an official
+project e-mail address, posting via an official social media account, or acting
+as an appointed representative at an online or offline event. Representation of
+a project may be further defined and clarified by project maintainers.
+
+This Code of Conduct also applies outside the project spaces when there is a
+reasonable belief that an individual's behavior may have a negative impact on
+the project or its community.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at <[email protected]>. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see
+https://www.contributor-covenant.org/faq

sam2-repo/CONTRIBUTING.md

@@ -0,0 +1,31 @@
+# Contributing to segment-anything
+We want to make contributing to this project as easy and transparent as
+possible.
+
+## Pull Requests
+We actively welcome your pull requests.
+
+1. Fork the repo and create your branch from `main`.
+2. If you've added code that should be tested, add tests.
+3. If you've changed APIs, update the documentation.
+4. Ensure the test suite passes.
+5. Make sure your code lints, using the `ufmt format` command. Linting requires `black==24.2.0`, `usort==1.0.2`, and `ufmt==2.0.0b2`, which can be installed via `pip install -e ".[dev]"`.
+6. If you haven't already, complete the Contributor License Agreement ("CLA").
+
+## Contributor License Agreement ("CLA")
+In order to accept your pull request, we need you to submit a CLA. You only need
+to do this once to work on any of Facebook's open source projects.
+
+Complete your CLA here: <https://code.facebook.com/cla>
+
+## Issues
+We use GitHub issues to track public bugs. Please ensure your description is
+clear and has sufficient instructions to be able to reproduce the issue.
+
+Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
+disclosure of security bugs. In those cases, please go through the process
+outlined on that page and do not file a public issue.
+
+## License
+By contributing to segment-anything, you agree that your contributions will be licensed
+under the LICENSE file in the root directory of this source tree.

sam2-repo/INSTALL.md

@@ -0,0 +1,189 @@
+## Installation
+
+### Requirements
+
+- Linux with Python ≥ 3.10, PyTorch ≥ 2.5.1 and [torchvision](https://github.com/pytorch/vision/) that matches the PyTorch installation. Install them together at https://pytorch.org to ensure this.
+  * Note older versions of Python or PyTorch may also work. However, the versions above are strongly recommended to provide all features such as `torch.compile`.
+- [CUDA toolkits](https://developer.nvidia.com/cuda-toolkit-archive) that match the CUDA version for your PyTorch installation. This should typically be CUDA 12.1 if you follow the default installation command.
+- If you are installing on Windows, it's strongly recommended to use [Windows Subsystem for Linux (WSL)](https://learn.microsoft.com/en-us/windows/wsl/install) with Ubuntu.
+
+Then, install SAM 2 from the root of this repository via
+```bash
+pip install -e ".[notebooks]"
+```
+
+Note that you may skip building the SAM 2 CUDA extension during installation via environment variable `SAM2_BUILD_CUDA=0`, as follows:
+```bash
+# skip the SAM 2 CUDA extension
+SAM2_BUILD_CUDA=0 pip install -e ".[notebooks]"
+```
+This would also skip the post-processing step at runtime (removing small holes and sprinkles in the output masks, which requires the CUDA extension), but shouldn't affect the results in most cases.
+
+### Building the SAM 2 CUDA extension
+
+By default, we allow the installation to proceed even if the SAM 2 CUDA extension fails to build. (In this case, the build errors are hidden unless using `-v` for verbose output in `pip install`.)
+
+If you see a message like `Skipping the post-processing step due to the error above` at runtime or `Failed to build the SAM 2 CUDA extension due to the error above` during installation, it indicates that the SAM 2 CUDA extension failed to build in your environment. In this case, **you can still use SAM 2 for both image and video applications**. The post-processing step (removing small holes and sprinkles in the output masks) will be skipped, but this shouldn't affect the results in most cases.
+
+If you would like to enable this post-processing step, you can reinstall SAM 2 on a GPU machine with environment variable `SAM2_BUILD_ALLOW_ERRORS=0` to force building the CUDA extension (and raise errors if it fails to build), as follows
+```bash
+pip uninstall -y SAM-2 && \
+rm -f ./sam2/*.so && \
+SAM2_BUILD_ALLOW_ERRORS=0 pip install -v -e ".[notebooks]"
+```
+
+Note that PyTorch needs to be installed first before building the SAM 2 CUDA extension. It's also necessary to install [CUDA toolkits](https://developer.nvidia.com/cuda-toolkit-archive) that match the CUDA version for your PyTorch installation. (This should typically be CUDA 12.1 if you follow the default installation command.) After installing the CUDA toolkits, you can check its version via `nvcc --version`.
+
+Please check the section below on common installation issues if the CUDA extension fails to build during installation or load at runtime.
+
+### Common Installation Issues
+
+Click each issue for its solutions:
+
+<details>
+<summary>
+I got `ImportError: cannot import name '_C' from 'sam2'`
+</summary>
+<br/>
+
+This is usually because you haven't run the `pip install -e ".[notebooks]"` step above or the installation failed. Please install SAM 2 first, and see the other issues if your installation fails.
+
+In some systems, you may need to run `python setup.py build_ext --inplace` in the SAM 2 repo root as suggested in https://github.com/facebookresearch/sam2/issues/77.
+</details>
+
+<details>
+<summary>
+I got `MissingConfigException: Cannot find primary config 'configs/sam2.1/sam2.1_hiera_l.yaml'`
+</summary>
+<br/>
+
+This is usually because you haven't run the `pip install -e .` step above, so `sam2` isn't in your Python's `sys.path`. Please run this installation step. In case it still fails after the installation step, you may try manually adding the root of this repo to `PYTHONPATH` via
+```bash
+export SAM2_REPO_ROOT=/path/to/sam2  # path to this repo
+export PYTHONPATH="${SAM2_REPO_ROOT}:${PYTHONPATH}"
+```
+to manually add `sam2_configs` into your Python's `sys.path`.
+
+</details>
+
+<details>
+<summary>
+I got `RuntimeError: Error(s) in loading state_dict for SAM2Base` when loading the new SAM 2.1 checkpoints
+</summary>
+<br/>
+
+This is likely because you have installed a previous version of this repo, which doesn't have the new modules to support the SAM 2.1 checkpoints yet. Please try the following steps:
+
+1. pull the latest code from the `main` branch of this repo
+2. run `pip uninstall -y SAM-2` to uninstall any previous installations
+3. then install the latest repo again using `pip install -e ".[notebooks]"`
+
+In case the steps above still don't resolve the error, please try running in your Python environment the following
+```python
+from sam2.modeling import sam2_base
+
+print(sam2_base.__file__)
+```
+and check whether the content in the printed local path of `sam2/modeling/sam2_base.py` matches the latest one in https://github.com/facebookresearch/sam2/blob/main/sam2/modeling/sam2_base.py (e.g. whether your local file has `no_obj_embed_spatial`) to indentify if you're still using a previous installation.
+
+</details>
+
+<details>
+<summary>
+My installation failed with `CUDA_HOME environment variable is not set`
+</summary>
+<br/>
+
+This usually happens because the installation step cannot find the CUDA toolkits (that contain the NVCC compiler) to build a custom CUDA kernel in SAM 2. Please install [CUDA toolkits](https://developer.nvidia.com/cuda-toolkit-archive) or the version that matches the CUDA version for your PyTorch installation. If the error persists after installing CUDA toolkits, you may explicitly specify `CUDA_HOME` via
+```
+export CUDA_HOME=/usr/local/cuda  # change to your CUDA toolkit path
+```
+and rerun the installation.
+
+Also, you should make sure
+```
+python -c 'import torch; from torch.utils.cpp_extension import CUDA_HOME; print(torch.cuda.is_available(), CUDA_HOME)'
+```
+print `(True, a directory with cuda)` to verify that the CUDA toolkits are correctly set up.
+
+If you are still having problems after verifying that the CUDA toolkit is installed and the `CUDA_HOME` environment variable is set properly, you may have to add the `--no-build-isolation` flag to the pip command:
+```
+pip install --no-build-isolation -e .
+```
+
+</details>
+
+<details>
+<summary>
+I got `undefined symbol: _ZN3c1015SmallVectorBaseIjE8grow_podEPKvmm` (or similar errors)
+</summary>
+<br/>
+
+This usually happens because you have multiple versions of dependencies (PyTorch or CUDA) in your environment. During installation, the SAM 2 library is compiled against one version library while at run time it links against another version. This might be due to that you have different versions of PyTorch or CUDA installed separately via `pip` or `conda`. You may delete one of the duplicates to only keep a single PyTorch and CUDA version.
+
+In particular, if you have a lower PyTorch version than 2.5.1, it's recommended to upgrade to PyTorch 2.5.1 or higher first. Otherwise, the installation script will try to upgrade to the latest PyTorch using `pip`, which could sometimes lead to duplicated PyTorch installation if you have previously installed another PyTorch version using `conda`.
+
+We have been building SAM 2 against PyTorch 2.5.1 internally. However, a few user comments (e.g. https://github.com/facebookresearch/sam2/issues/22, https://github.com/facebookresearch/sam2/issues/14) suggested that downgrading to PyTorch 2.1.0 might resolve this problem. In case the error persists, you may try changing the restriction from `torch>=2.5.1` to `torch==2.1.0` in both [`pyproject.toml`](pyproject.toml) and [`setup.py`](setup.py) to allow PyTorch 2.1.0.
+</details>
+
+<details>
+<summary>
+I got `CUDA error: no kernel image is available for execution on the device`
+</summary>
+<br/>
+
+A possible cause could be that the CUDA kernel is somehow not compiled towards your GPU's CUDA [capability](https://developer.nvidia.com/cuda-gpus). This could happen if the installation is done in an environment different from the runtime (e.g. in a slurm system).
+
+You can try pulling the latest code from the SAM 2 repo and running the following
+```
+export TORCH_CUDA_ARCH_LIST=9.0 8.0 8.6 8.9 7.0 7.2 7.5 6.0`
+```
+to manually specify the CUDA capability in the compilation target that matches your GPU.
+</details>
+
+<details>
+<summary>
+I got `RuntimeError: No available kernel. Aborting execution.` (or similar errors)
+</summary>
+<br/>
+
+This is probably because your machine doesn't have a GPU or a compatible PyTorch version for Flash Attention (see also https://discuss.pytorch.org/t/using-f-scaled-dot-product-attention-gives-the-error-runtimeerror-no-available-kernel-aborting-execution/180900 for a discussion in PyTorch forum). You may be able to resolve this error by replacing the line
+```python
+OLD_GPU, USE_FLASH_ATTN, MATH_KERNEL_ON = get_sdpa_settings()
+```
+in [`sam2/modeling/sam/transformer.py`](sam2/modeling/sam/transformer.py) with
+```python
+OLD_GPU, USE_FLASH_ATTN, MATH_KERNEL_ON = True, True, True
+```
+to relax the attention kernel setting and use other kernels than Flash Attention.
+</details>
+
+<details>
+<summary>
+I got `Error compiling objects for extension`
+</summary>
+<br/>
+
+You may see error log of:
+> unsupported Microsoft Visual Studio version! Only the versions between 2017 and 2022 (inclusive) are supported! The nvcc flag '-allow-unsupported-compiler' can be used to override this version check; however, using an unsupported host compiler may cause compilation failure or incorrect run time execution. Use at your own risk.
+
+This is probably because your versions of CUDA and Visual Studio are incompatible. (see also https://stackoverflow.com/questions/78515942/cuda-compatibility-with-visual-studio-2022-version-17-10 for a discussion in stackoverflow).<br> 
+You may be able to fix this by adding the `-allow-unsupported-compiler` argument to `nvcc` after L48 in the [setup.py](https://github.com/facebookresearch/sam2/blob/main/setup.py). <br>
+After adding the argument, `get_extension()` will look like this:
+```python
+def get_extensions():
+    srcs = ["sam2/csrc/connected_components.cu"]
+    compile_args = {
+        "cxx": [],
+        "nvcc": [
+            "-DCUDA_HAS_FP16=1",
+            "-D__CUDA_NO_HALF_OPERATORS__",
+            "-D__CUDA_NO_HALF_CONVERSIONS__",
+            "-D__CUDA_NO_HALF2_OPERATORS__",
+            "-allow-unsupported-compiler"  # Add this argument
+        ],
+    }
+    ext_modules = [CUDAExtension("sam2._C", srcs, extra_compile_args=compile_args)]
+    return ext_modules
+```
+</details>

sam2-repo/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

sam2-repo/LICENSE_cctorch

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2020, the respective contributors, as shown by the AUTHORS file.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

sam2-repo/MANIFEST.in

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+recursive-include sam2 *.yaml #include all config files

sam2-repo/README.md

@@ -0,0 +1,224 @@
+# SAM 2: Segment Anything in Images and Videos
+
+**[AI at Meta, FAIR](https://ai.meta.com/research/)**
+
+[Nikhila Ravi](https://nikhilaravi.com/), [Valentin Gabeur](https://gabeur.github.io/), [Yuan-Ting Hu](https://scholar.google.com/citations?user=E8DVVYQAAAAJ&hl=en), [Ronghang Hu](https://ronghanghu.com/), [Chaitanya Ryali](https://scholar.google.com/citations?user=4LWx24UAAAAJ&hl=en), [Tengyu Ma](https://scholar.google.com/citations?user=VeTSl0wAAAAJ&hl=en), [Haitham Khedr](https://hkhedr.com/), [Roman Rädle](https://scholar.google.de/citations?user=Tpt57v0AAAAJ&hl=en), [Chloe Rolland](https://scholar.google.com/citations?hl=fr&user=n-SnMhoAAAAJ), [Laura Gustafson](https://scholar.google.com/citations?user=c8IpF9gAAAAJ&hl=en), [Eric Mintun](https://ericmintun.github.io/), [Junting Pan](https://junting.github.io/), [Kalyan Vasudev Alwala](https://scholar.google.co.in/citations?user=m34oaWEAAAAJ&hl=en), [Nicolas Carion](https://www.nicolascarion.com/), [Chao-Yuan Wu](https://chaoyuan.org/), [Ross Girshick](https://www.rossgirshick.info/), [Piotr Dollár](https://pdollar.github.io/), [Christoph Feichtenhofer](https://feichtenhofer.github.io/)
+
+[[`Paper`](https://ai.meta.com/research/publications/sam-2-segment-anything-in-images-and-videos/)] [[`Project`](https://ai.meta.com/sam2)] [[`Demo`](https://sam2.metademolab.com/)] [[`Dataset`](https://ai.meta.com/datasets/segment-anything-video)] [[`Blog`](https://ai.meta.com/blog/segment-anything-2)] [[`BibTeX`](#citing-sam-2)]
+
+![SAM 2 architecture](assets/model_diagram.png?raw=true)
+
+**Segment Anything Model 2 (SAM 2)** is a foundation model towards solving promptable visual segmentation in images and videos. We extend SAM to video by considering images as a video with a single frame. The model design is a simple transformer architecture with streaming memory for real-time video processing. We build a model-in-the-loop data engine, which improves model and data via user interaction, to collect [**our SA-V dataset**](https://ai.meta.com/datasets/segment-anything-video), the largest video segmentation dataset to date. SAM 2 trained on our data provides strong performance across a wide range of tasks and visual domains.
+
+![SA-V dataset](assets/sa_v_dataset.jpg?raw=true)
+
+## Latest updates
+
+**12/11/2024 -- full model compilation for a major VOS speedup and a new `SAM2VideoPredictor` to better handle multi-object tracking**
+
+- We now support `torch.compile` of the entire SAM 2 model on videos, which can be turned on by setting `vos_optimized=True` in `build_sam2_video_predictor`, leading to a major speedup for VOS inference.
+- We update the implementation of `SAM2VideoPredictor` to support independent per-object inference, allowing us to relax the assumption of prompting for multi-object tracking and adding new objects after tracking starts.
+- See [`RELEASE_NOTES.md`](RELEASE_NOTES.md) for full details.
+
+**09/30/2024 -- SAM 2.1 Developer Suite (new checkpoints, training code, web demo) is released**
+
+- A new suite of improved model checkpoints (denoted as **SAM 2.1**) are released. See [Model Description](#model-description) for details.
+  * To use the new SAM 2.1 checkpoints, you need the latest model code from this repo. If you have installed an earlier version of this repo, please first uninstall the previous version via `pip uninstall SAM-2`, pull the latest code from this repo (with `git pull`), and then reinstall the repo following [Installation](#installation) below.
+- The training (and fine-tuning) code has been released. See [`training/README.md`](training/README.md) on how to get started.
+- The frontend + backend code for the SAM 2 web demo has been released. See [`demo/README.md`](demo/README.md) for details.
+
+## Installation
+
+SAM 2 needs to be installed first before use. The code requires `python>=3.10`, as well as `torch>=2.5.1` and `torchvision>=0.20.1`. Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install both PyTorch and TorchVision dependencies. You can install SAM 2 on a GPU machine using:
+
+```bash
+git clone https://github.com/facebookresearch/sam2.git && cd sam2
+
+pip install -e .
+```
+If you are installing on Windows, it's strongly recommended to use [Windows Subsystem for Linux (WSL)](https://learn.microsoft.com/en-us/windows/wsl/install) with Ubuntu.
+
+To use the SAM 2 predictor and run the example notebooks, `jupyter` and `matplotlib` are required and can be installed by:
+
+```bash
+pip install -e ".[notebooks]"
+```
+
+Note:
+1. It's recommended to create a new Python environment via [Anaconda](https://www.anaconda.com/) for this installation and install PyTorch 2.5.1 (or higher) via `pip` following https://pytorch.org/. If you have a PyTorch version lower than 2.5.1 in your current environment, the installation command above will try to upgrade it to the latest PyTorch version using `pip`.
+2. The step above requires compiling a custom CUDA kernel with the `nvcc` compiler. If it isn't already available on your machine, please install the [CUDA toolkits](https://developer.nvidia.com/cuda-toolkit-archive) with a version that matches your PyTorch CUDA version.
+3. If you see a message like `Failed to build the SAM 2 CUDA extension` during installation, you can ignore it and still use SAM 2 (some post-processing functionality may be limited, but it doesn't affect the results in most cases).
+
+Please see [`INSTALL.md`](./INSTALL.md) for FAQs on potential issues and solutions.
+
+## Getting Started
+
+### Download Checkpoints
+
+First, we need to download a model checkpoint. All the model checkpoints can be downloaded by running:
+
+```bash
+cd checkpoints && \
+./download_ckpts.sh && \
+cd ..
+```
+
+or individually from:
+
+- [sam2.1_hiera_tiny.pt](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt)
+- [sam2.1_hiera_small.pt](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt)
+- [sam2.1_hiera_base_plus.pt](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_base_plus.pt)
+- [sam2.1_hiera_large.pt](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt)
+
+(note that these are the improved checkpoints denoted as SAM 2.1; see [Model Description](#model-description) for details.)
+
+Then SAM 2 can be used in a few lines as follows for image and video prediction.
+
+### Image prediction
+
+SAM 2 has all the capabilities of [SAM](https://github.com/facebookresearch/segment-anything) on static images, and we provide image prediction APIs that closely resemble SAM for image use cases. The `SAM2ImagePredictor` class has an easy interface for image prompting.
+
+```python
+import torch
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
+model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
+predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint))
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    predictor.set_image(<your_image>)
+    masks, _, _ = predictor.predict(<input_prompts>)
+```
+
+Please refer to the examples in [image_predictor_example.ipynb](./notebooks/image_predictor_example.ipynb) (also in Colab [here](https://colab.research.google.com/github/facebookresearch/sam2/blob/main/notebooks/image_predictor_example.ipynb)) for static image use cases.
+
+SAM 2 also supports automatic mask generation on images just like SAM. Please see [automatic_mask_generator_example.ipynb](./notebooks/automatic_mask_generator_example.ipynb) (also in Colab [here](https://colab.research.google.com/github/facebookresearch/sam2/blob/main/notebooks/automatic_mask_generator_example.ipynb)) for automatic mask generation in images.
+
+### Video prediction
+
+For promptable segmentation and tracking in videos, we provide a video predictor with APIs for example to add prompts and propagate masklets throughout a video. SAM 2 supports video inference on multiple objects and uses an inference state to keep track of the interactions in each video.
+
+```python
+import torch
+from sam2.build_sam import build_sam2_video_predictor
+
+checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
+model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
+predictor = build_sam2_video_predictor(model_cfg, checkpoint)
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    state = predictor.init_state(<your_video>)
+
+    # add new prompts and instantly get the output on the same frame
+    frame_idx, object_ids, masks = predictor.add_new_points_or_box(state, <your_prompts>):
+
+    # propagate the prompts to get masklets throughout the video
+    for frame_idx, object_ids, masks in predictor.propagate_in_video(state):
+        ...
+```
+
+Please refer to the examples in [video_predictor_example.ipynb](./notebooks/video_predictor_example.ipynb) (also in Colab [here](https://colab.research.google.com/github/facebookresearch/sam2/blob/main/notebooks/video_predictor_example.ipynb)) for details on how to add click or box prompts, make refinements, and track multiple objects in videos.
+
+## Load from 🤗 Hugging Face
+
+Alternatively, models can also be loaded from [Hugging Face](https://huggingface.co/models?search=facebook/sam2) (requires `pip install huggingface_hub`).
+
+For image prediction:
+
+```python
+import torch
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    predictor.set_image(<your_image>)
+    masks, _, _ = predictor.predict(<input_prompts>)
+```
+
+For video prediction:
+
+```python
+import torch
+from sam2.sam2_video_predictor import SAM2VideoPredictor
+
+predictor = SAM2VideoPredictor.from_pretrained("facebook/sam2-hiera-large")
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    state = predictor.init_state(<your_video>)
+
+    # add new prompts and instantly get the output on the same frame
+    frame_idx, object_ids, masks = predictor.add_new_points_or_box(state, <your_prompts>):
+
+    # propagate the prompts to get masklets throughout the video
+    for frame_idx, object_ids, masks in predictor.propagate_in_video(state):
+        ...
+```
+
+## Model Description
+
+### SAM 2.1 checkpoints
+
+The table below shows the improved SAM 2.1 checkpoints released on September 29, 2024.
+|      **Model**       | **Size (M)** |    **Speed (FPS)**     | **SA-V test (J&F)** | **MOSE val (J&F)** | **LVOS v2 (J&F)** |
+| :------------------: | :----------: | :--------------------: | :-----------------: | :----------------: | :---------------: |
+|   sam2.1_hiera_tiny <br /> ([config](sam2/configs/sam2.1/sam2.1_hiera_t.yaml), [checkpoint](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt))    |     38.9     |          91.2          |        76.5         |        71.8        |       77.3        |
+|   sam2.1_hiera_small <br /> ([config](sam2/configs/sam2.1/sam2.1_hiera_s.yaml), [checkpoint](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt))   |      46      |          84.8          |        76.6         |        73.5        |       78.3        |
+| sam2.1_hiera_base_plus <br /> ([config](sam2/configs/sam2.1/sam2.1_hiera_b+.yaml), [checkpoint](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_base_plus.pt)) |     80.8     |        64.1          |        78.2         |        73.7        |       78.2        |
+|   sam2.1_hiera_large <br /> ([config](sam2/configs/sam2.1/sam2.1_hiera_l.yaml), [checkpoint](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt))   |    224.4     |          39.5          |        79.5         |        74.6        |       80.6        |
+
+### SAM 2 checkpoints
+
+The previous SAM 2 checkpoints released on July 29, 2024 can be found as follows:
+
+|      **Model**       | **Size (M)** |    **Speed (FPS)**     | **SA-V test (J&F)** | **MOSE val (J&F)** | **LVOS v2 (J&F)** |
+| :------------------: | :----------: | :--------------------: | :-----------------: | :----------------: | :---------------: |
+|   sam2_hiera_tiny <br /> ([config](sam2/configs/sam2/sam2_hiera_t.yaml), [checkpoint](https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt))   |     38.9     |          91.5          |        75.0         |        70.9        |       75.3        |
+|   sam2_hiera_small <br /> ([config](sam2/configs/sam2/sam2_hiera_s.yaml), [checkpoint](https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt))   |      46      |          85.6          |        74.9         |        71.5        |       76.4        |
+| sam2_hiera_base_plus <br /> ([config](sam2/configs/sam2/sam2_hiera_b+.yaml), [checkpoint](https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt)) |     80.8     |     64.8    |        74.7         |        72.8        |       75.8        |
+|   sam2_hiera_large <br /> ([config](sam2/configs/sam2/sam2_hiera_l.yaml), [checkpoint](https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt))   |    224.4     | 39.7 |        76.0         |        74.6        |       79.8        |
+
+Speed measured on an A100 with `torch 2.5.1, cuda 12.4`. See `benchmark.py` for an example on benchmarking (compiling all the model components). Compiling only the image encoder can be more flexible and also provide (a smaller) speed-up (set `compile_image_encoder: True` in the config).
+## Segment Anything Video Dataset
+
+See [sav_dataset/README.md](sav_dataset/README.md) for details.
+
+## Training SAM 2
+
+You can train or fine-tune SAM 2 on custom datasets of images, videos, or both. Please check the training [README](training/README.md) on how to get started.
+
+## Web demo for SAM 2
+
+We have released the frontend + backend code for the SAM 2 web demo (a locally deployable version similar to https://sam2.metademolab.com/demo). Please see the web demo [README](demo/README.md) for details.
+
+## License
+
+The SAM 2 model checkpoints, SAM 2 demo code (front-end and back-end), and SAM 2 training code are licensed under [Apache 2.0](./LICENSE), however the [Inter Font](https://github.com/rsms/inter?tab=OFL-1.1-1-ov-file) and [Noto Color Emoji](https://github.com/googlefonts/noto-emoji) used in the SAM 2 demo code are made available under the [SIL Open Font License, version 1.1](https://openfontlicense.org/open-font-license-official-text/).
+
+## Contributing
+
+See [contributing](CONTRIBUTING.md) and the [code of conduct](CODE_OF_CONDUCT.md).
+
+## Contributors
+
+The SAM 2 project was made possible with the help of many contributors (alphabetical):
+
+Karen Bergan, Daniel Bolya, Alex Bosenberg, Kai Brown, Vispi Cassod, Christopher Chedeau, Ida Cheng, Luc Dahlin, Shoubhik Debnath, Rene Martinez Doehner, Grant Gardner, Sahir Gomez, Rishi Godugu, Baishan Guo, Caleb Ho, Andrew Huang, Somya Jain, Bob Kamma, Amanda Kallet, Jake Kinney, Alexander Kirillov, Shiva Koduvayur, Devansh Kukreja, Robert Kuo, Aohan Lin, Parth Malani, Jitendra Malik, Mallika Malhotra, Miguel Martin, Alexander Miller, Sasha Mitts, William Ngan, George Orlin, Joelle Pineau, Kate Saenko, Rodrick Shepard, Azita Shokrpour, David Soofian, Jonathan Torres, Jenny Truong, Sagar Vaze, Meng Wang, Claudette Ward, Pengchuan Zhang.
+
+Third-party code: we use a GPU-based connected component algorithm adapted from [`cc_torch`](https://github.com/zsef123/Connected_components_PyTorch) (with its license in [`LICENSE_cctorch`](./LICENSE_cctorch)) as an optional post-processing step for the mask predictions.
+
+## Citing SAM 2
+
+If you use SAM 2 or the SA-V dataset in your research, please use the following BibTeX entry.
+
+```bibtex
+@article{ravi2024sam2,
+  title={SAM 2: Segment Anything in Images and Videos},
+  author={Ravi, Nikhila and Gabeur, Valentin and Hu, Yuan-Ting and Hu, Ronghang and Ryali, Chaitanya and Ma, Tengyu and Khedr, Haitham and R{\"a}dle, Roman and Rolland, Chloe and Gustafson, Laura and Mintun, Eric and Pan, Junting and Alwala, Kalyan Vasudev and Carion, Nicolas and Wu, Chao-Yuan and Girshick, Ross and Doll{\'a}r, Piotr and Feichtenhofer, Christoph},
+  journal={arXiv preprint arXiv:2408.00714},
+  url={https://arxiv.org/abs/2408.00714},
+  year={2024}
+}
+```

sam2-repo/RELEASE_NOTES.md

@@ -0,0 +1,27 @@
+## SAM 2 release notes
+
+### 12/11/2024 -- full model compilation for a major VOS speedup and a new `SAM2VideoPredictor` to better handle multi-object tracking
+
+- We now support `torch.compile` of the entire SAM 2 model on videos, which can be turned on by setting `vos_optimized=True` in `build_sam2_video_predictor` (it uses the new `SAM2VideoPredictorVOS` predictor class in `sam2/sam2_video_predictor.py`).
+  * Compared to the previous setting (which only compiles the image encoder backbone), the new full model compilation gives a major speedup in inference FPS.
+  * In the VOS prediction script `tools/vos_inference.py`, you can specify this option in `tools/vos_inference.py` via the `--use_vos_optimized_video_predictor` flag.
+  * Note that turning on this flag might introduce a small variance in the predictions due to numerical differences caused by `torch.compile` of the full model.
+  * **PyTorch 2.5.1 is the minimum version for full support of this feature**. (Earlier PyTorch versions might run into compilation errors in some cases.) Therefore, we have updated the minimum PyTorch version to 2.5.1 accordingly in the installation scripts.
+- We also update the implementation of the `SAM2VideoPredictor` class for the SAM 2 video prediction in `sam2/sam2_video_predictor.py`, which allows for independent per-object inference. Specifically, in the new `SAM2VideoPredictor`:
+  * Now **we handle the inference of each object independently** (as if we are opening a separate session for each object) while sharing their backbone features.
+  * This change allows us to relax the assumption of prompting for multi-object tracking. Previously (due to the batching behavior in inference), if a video frame receives clicks for only a subset of objects, the rest of the (non-prompted) objects are assumed to be non-existent in this frame (i.e., in such frames, the user is telling SAM 2 that the rest of the objects don't appear). Now, if a frame receives clicks for only a subset of objects, we do not make any assumptions about the remaining (non-prompted) objects (i.e., now each object is handled independently and is not affected by how other objects are prompted). As a result, **we allow adding new objects after tracking starts** after this change (which was previously a restriction on usage).
+  * We believe that the new version is a more natural inference behavior and therefore switched to it as the default behavior. The previous implementation of `SAM2VideoPredictor` is backed up to in `sam2/sam2_video_predictor_legacy.py`. All the VOS inference results using `tools/vos_inference.py` should remain the same after this change to the `SAM2VideoPredictor` class.
+
+### 09/30/2024 -- SAM 2.1 Developer Suite (new checkpoints, training code, web demo) is released
+
+- A new suite of improved model checkpoints (denoted as **SAM 2.1**) are released. See [Model Description](#model-description) for details.
+  * To use the new SAM 2.1 checkpoints, you need the latest model code from this repo. If you have installed an earlier version of this repo, please first uninstall the previous version via `pip uninstall SAM-2`, pull the latest code from this repo (with `git pull`), and then reinstall the repo following [Installation](#installation) below.
+- The training (and fine-tuning) code has been released. See [`training/README.md`](training/README.md) on how to get started.
+- The frontend + backend code for the SAM 2 web demo has been released. See [`demo/README.md`](demo/README.md) for details.
+
+### 07/29/2024 -- SAM 2 is released
+
+- We release Segment Anything Model 2 (SAM 2), a foundation model towards solving promptable visual segmentation in images and videos.
+  * SAM 2 code: https://github.com/facebookresearch/sam2
+  * SAM 2 demo: https://sam2.metademolab.com/
+  * SAM 2 paper: https://arxiv.org/abs/2408.00714

sam2-repo/backend.Dockerfile

@@ -0,0 +1,64 @@
+ARG BASE_IMAGE=pytorch/pytorch:2.5.1-cuda12.1-cudnn9-runtime
+ARG MODEL_SIZE=base_plus
+
+FROM ${BASE_IMAGE}
+
+# Gunicorn environment variables
+ENV GUNICORN_WORKERS=1
+ENV GUNICORN_THREADS=2
+ENV GUNICORN_PORT=5000
+
+# SAM 2 environment variables
+ENV APP_ROOT=/opt/sam2
+ENV PYTHONUNBUFFERED=1
+ENV SAM2_BUILD_CUDA=0
+ENV MODEL_SIZE=${MODEL_SIZE}
+
+# Install system requirements
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    ffmpeg \
+    libavutil-dev \
+    libavcodec-dev \
+    libavformat-dev \
+    libswscale-dev \
+    pkg-config \
+    build-essential \
+    libffi-dev
+
+COPY setup.py .
+COPY README.md .
+
+RUN pip install --upgrade pip setuptools
+RUN pip install -e ".[interactive-demo]"
+
+# https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite/issues/69#issuecomment-1826764707
+RUN rm /opt/conda/bin/ffmpeg && ln -s /bin/ffmpeg /opt/conda/bin/ffmpeg
+
+# Make app directory. This directory will host all files required for the
+# backend and SAM 2 inference files.
+RUN mkdir ${APP_ROOT}
+
+# Copy backend server files
+COPY demo/backend/server ${APP_ROOT}/server
+
+# Copy SAM 2 inference files
+COPY sam2 ${APP_ROOT}/server/sam2
+
+# Download SAM 2.1 checkpoints
+ADD https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt ${APP_ROOT}/checkpoints/sam2.1_hiera_tiny.pt
+ADD https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt ${APP_ROOT}/checkpoints/sam2.1_hiera_small.pt
+ADD https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_base_plus.pt ${APP_ROOT}/checkpoints/sam2.1_hiera_base_plus.pt
+ADD https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt ${APP_ROOT}/checkpoints/sam2.1_hiera_large.pt
+
+WORKDIR ${APP_ROOT}/server
+
+# https://pythonspeed.com/articles/gunicorn-in-docker/
+CMD gunicorn --worker-tmp-dir /dev/shm \
+    --worker-class gthread app:app \
+    --log-level info \
+    --access-logfile /dev/stdout \
+    --log-file /dev/stderr \
+    --workers ${GUNICORN_WORKERS} \
+    --threads ${GUNICORN_THREADS} \
+    --bind 0.0.0.0:${GUNICORN_PORT} \
+    --timeout 60

sam2-repo/checkpoints/download_ckpts.sh

@@ -0,0 +1,59 @@
+#!/bin/bash
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Use either wget or curl to download the checkpoints
+if command -v wget &> /dev/null; then
+    CMD="wget"
+elif command -v curl &> /dev/null; then
+    CMD="curl -L -O"
+else
+    echo "Please install wget or curl to download the checkpoints."
+    exit 1
+fi
+
+# Define the URLs for SAM 2 checkpoints
+# SAM2_BASE_URL="https://dl.fbaipublicfiles.com/segment_anything_2/072824"
+# sam2_hiera_t_url="${SAM2_BASE_URL}/sam2_hiera_tiny.pt"
+# sam2_hiera_s_url="${SAM2_BASE_URL}/sam2_hiera_small.pt"
+# sam2_hiera_b_plus_url="${SAM2_BASE_URL}/sam2_hiera_base_plus.pt"
+# sam2_hiera_l_url="${SAM2_BASE_URL}/sam2_hiera_large.pt"
+
+# Download each of the four checkpoints using wget
+# echo "Downloading sam2_hiera_tiny.pt checkpoint..."
+# $CMD $sam2_hiera_t_url || { echo "Failed to download checkpoint from $sam2_hiera_t_url"; exit 1; }
+
+# echo "Downloading sam2_hiera_small.pt checkpoint..."
+# $CMD $sam2_hiera_s_url || { echo "Failed to download checkpoint from $sam2_hiera_s_url"; exit 1; }
+
+# echo "Downloading sam2_hiera_base_plus.pt checkpoint..."
+# $CMD $sam2_hiera_b_plus_url || { echo "Failed to download checkpoint from $sam2_hiera_b_plus_url"; exit 1; }
+
+# echo "Downloading sam2_hiera_large.pt checkpoint..."
+# $CMD $sam2_hiera_l_url || { echo "Failed to download checkpoint from $sam2_hiera_l_url"; exit 1; }
+
+# Define the URLs for SAM 2.1 checkpoints
+SAM2p1_BASE_URL="https://dl.fbaipublicfiles.com/segment_anything_2/092824"
+sam2p1_hiera_t_url="${SAM2p1_BASE_URL}/sam2.1_hiera_tiny.pt"
+sam2p1_hiera_s_url="${SAM2p1_BASE_URL}/sam2.1_hiera_small.pt"
+sam2p1_hiera_b_plus_url="${SAM2p1_BASE_URL}/sam2.1_hiera_base_plus.pt"
+sam2p1_hiera_l_url="${SAM2p1_BASE_URL}/sam2.1_hiera_large.pt"
+
+# SAM 2.1 checkpoints
+echo "Downloading sam2.1_hiera_tiny.pt checkpoint..."
+$CMD $sam2p1_hiera_t_url || { echo "Failed to download checkpoint from $sam2p1_hiera_t_url"; exit 1; }
+
+echo "Downloading sam2.1_hiera_small.pt checkpoint..."
+$CMD $sam2p1_hiera_s_url || { echo "Failed to download checkpoint from $sam2p1_hiera_s_url"; exit 1; }
+
+echo "Downloading sam2.1_hiera_base_plus.pt checkpoint..."
+$CMD $sam2p1_hiera_b_plus_url || { echo "Failed to download checkpoint from $sam2p1_hiera_b_plus_url"; exit 1; }
+
+echo "Downloading sam2.1_hiera_large.pt checkpoint..."
+$CMD $sam2p1_hiera_l_url || { echo "Failed to download checkpoint from $sam2p1_hiera_l_url"; exit 1; }
+
+echo "All checkpoints are downloaded successfully."

sam2-repo/demo/.gitignore

@@ -0,0 +1,2 @@
+data/uploads
+data/posters

sam2-repo/demo/README.md

@@ -0,0 +1,173 @@
+# SAM 2 Demo
+
+Welcome to the SAM 2 Demo! This project consists of a frontend built with React TypeScript and Vite and a backend service using Python Flask and Strawberry GraphQL. Both components can be run in Docker containers or locally on MPS (Metal Performance Shaders) or CPU. However, running the backend service on MPS or CPU devices may result in significantly slower performance (FPS).
+
+## Prerequisites
+
+Before you begin, ensure you have the following installed on your system:
+
+- Docker and Docker Compose
+- [OPTIONAL] Node.js and Yarn for running frontend locally
+- [OPTIONAL] Anaconda for running backend locally
+
+### Installing Docker
+
+To install Docker, follow these steps:
+
+1. Go to the [Docker website](https://www.docker.com/get-started)
+2. Follow the installation instructions for your operating system.
+
+### [OPTIONAL] Installing Node.js and Yarn
+
+To install Node.js and Yarn, follow these steps:
+
+1. Go to the [Node.js website](https://nodejs.org/en/download/).
+2. Follow the installation instructions for your operating system.
+3. Once Node.js is installed, open a terminal or command prompt and run the following command to install Yarn:
+
+```
+npm install -g yarn
+```
+
+### [OPTIONAL] Installing Anaconda
+
+To install Anaconda, follow these steps:
+
+1. Go to the [Anaconda website](https://www.anaconda.com/products/distribution).
+2. Follow the installation instructions for your operating system.
+
+## Quick Start
+
+To get both the frontend and backend running quickly using Docker, you can use the following command:
+
+```bash
+docker compose up --build
+```
+
+> [!WARNING]
+> On macOS, Docker containers only support running on CPU. MPS is not supported through Docker. If you want to run the demo backend service on MPS, you will need to run it locally (see "Running the Backend Locally" below).
+
+This will build and start both services. You can access them at:
+
+- **Frontend:** [http://localhost:7262](http://localhost:7262)
+- **Backend:** [http://localhost:7263/graphql](http://localhost:7263/graphql)
+
+## Running Backend with MPS Support
+
+MPS (Metal Performance Shaders) is not supported with Docker. To use MPS, you need to run the backend on your local machine.
+
+### Setting Up Your Environment
+
+1. **Create Conda environment**
+
+   Create a new Conda environment for this project by running the following command or use your existing conda environment for SAM 2:
+
+   ```
+   conda create --name sam2-demo python=3.10 --yes
+   ```
+
+   This will create a new environment named `sam2-demo` with Python 3.10 as the interpreter.
+
+2. **Activate the Conda environment:**
+
+   ```bash
+   conda activate sam2-demo
+   ```
+
+3. **Install ffmpeg**
+
+   ```bash
+   conda install -c conda-forge ffmpeg
+   ```
+
+4. **Install SAM 2 demo dependencies:**
+
+Install project dependencies by running the following command in the SAM 2 checkout root directory:
+
+```bash
+pip install -e '.[interactive-demo]'
+```
+
+### Running the Backend Locally
+
+Download the SAM 2 checkpoints:
+
+```bash
+(cd ./checkpoints && ./download_ckpts.sh)
+```
+
+Use the following command to start the backend with MPS support:
+
+```bash
+cd demo/backend/server/
+```
+
+```bash
+PYTORCH_ENABLE_MPS_FALLBACK=1 \
+APP_ROOT="$(pwd)/../../../" \
+API_URL=http://localhost:7263 \
+MODEL_SIZE=base_plus \
+DATA_PATH="$(pwd)/../../data" \
+DEFAULT_VIDEO_PATH=gallery/05_default_juggle.mp4 \
+gunicorn \
+    --worker-class gthread app:app \
+    --workers 1 \
+    --threads 2 \
+    --bind 0.0.0.0:7263 \
+    --timeout 60
+```
+
+Options for the `MODEL_SIZE` argument are "tiny", "small", "base_plus" (default), and "large".
+
+> [!WARNING]
+> Running the backend service on MPS devices can cause fatal crashes with the Gunicorn worker due to insufficient MPS memory. Try switching to CPU devices by setting the `SAM2_DEMO_FORCE_CPU_DEVICE=1` environment variable.
+
+### Starting the Frontend
+
+If you wish to run the frontend separately (useful for development), follow these steps:
+
+1. **Navigate to demo frontend directory:**
+
+   ```bash
+   cd demo/frontend
+   ```
+
+2. **Install dependencies:**
+
+   ```bash
+   yarn install
+   ```
+
+3. **Start the development server:**
+
+   ```bash
+   yarn dev --port 7262
+   ```
+
+This will start the frontend development server on [http://localhost:7262](http://localhost:7262).
+
+## Docker Tips
+
+- To rebuild the Docker containers (useful if you've made changes to the Dockerfile or dependencies):
+
+  ```bash
+  docker compose up --build
+  ```
+
+- To stop the Docker containers:
+
+  ```bash
+  docker compose down
+  ```
+
+## Contributing
+
+Contributions are welcome! Please read our contributing guidelines to get started.
+
+## License
+
+See the LICENSE file for details.
+
+---
+
+By following these instructions, you should have a fully functional development environment for both the frontend and backend of the SAM 2 Demo. Happy coding!

sam2-repo/demo/backend/server/app.py

@@ -0,0 +1,140 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from typing import Any, Generator
+
+from app_conf import (
+    GALLERY_PATH,
+    GALLERY_PREFIX,
+    POSTERS_PATH,
+    POSTERS_PREFIX,
+    UPLOADS_PATH,
+    UPLOADS_PREFIX,
+)
+from data.loader import preload_data
+from data.schema import schema
+from data.store import set_videos
+from flask import Flask, make_response, Request, request, Response, send_from_directory
+from flask_cors import CORS
+from inference.data_types import PropagateDataResponse, PropagateInVideoRequest
+from inference.multipart import MultipartResponseBuilder
+from inference.predictor import InferenceAPI
+from strawberry.flask.views import GraphQLView
+
+logger = logging.getLogger(__name__)
+
+app = Flask(__name__)
+cors = CORS(app, supports_credentials=True)
+
+videos = preload_data()
+set_videos(videos)
+
+inference_api = InferenceAPI()
+
+
+@app.route("/healthy")
+def healthy() -> Response:
+    return make_response("OK", 200)
+
+
+@app.route(f"/{GALLERY_PREFIX}/<path:path>", methods=["GET"])
+def send_gallery_video(path: str) -> Response:
+    try:
+        return send_from_directory(
+            GALLERY_PATH,
+            path,
+        )
+    except:
+        raise ValueError("resource not found")
+
+
+@app.route(f"/{POSTERS_PREFIX}/<path:path>", methods=["GET"])
+def send_poster_image(path: str) -> Response:
+    try:
+        return send_from_directory(
+            POSTERS_PATH,
+            path,
+        )
+    except:
+        raise ValueError("resource not found")
+
+
+@app.route(f"/{UPLOADS_PREFIX}/<path:path>", methods=["GET"])
+def send_uploaded_video(path: str):
+    try:
+        return send_from_directory(
+            UPLOADS_PATH,
+            path,
+        )
+    except:
+        raise ValueError("resource not found")
+
+
+# TOOD: Protect route with ToS permission check
+@app.route("/propagate_in_video", methods=["POST"])
+def propagate_in_video() -> Response:
+    data = request.json
+    args = {
+        "session_id": data["session_id"],
+        "start_frame_index": data.get("start_frame_index", 0),
+    }
+
+    boundary = "frame"
+    frame = gen_track_with_mask_stream(boundary, **args)
+    return Response(frame, mimetype="multipart/x-savi-stream; boundary=" + boundary)
+
+
+def gen_track_with_mask_stream(
+    boundary: str,
+    session_id: str,
+    start_frame_index: int,
+) -> Generator[bytes, None, None]:
+    with inference_api.autocast_context():
+        request = PropagateInVideoRequest(
+            type="propagate_in_video",
+            session_id=session_id,
+            start_frame_index=start_frame_index,
+        )
+
+        for chunk in inference_api.propagate_in_video(request=request):
+            yield MultipartResponseBuilder.build(
+                boundary=boundary,
+                headers={
+                    "Content-Type": "application/json; charset=utf-8",
+                    "Frame-Current": "-1",
+                    # Total frames minus the reference frame
+                    "Frame-Total": "-1",
+                    "Mask-Type": "RLE[]",
+                },
+                body=chunk.to_json().encode("UTF-8"),
+            ).get_message()
+
+
+class MyGraphQLView(GraphQLView):
+    def get_context(self, request: Request, response: Response) -> Any:
+        return {"inference_api": inference_api}
+
+
+# Add GraphQL route to Flask app.
+app.add_url_rule(
+    "/graphql",
+    view_func=MyGraphQLView.as_view(
+        "graphql_view",
+        schema=schema,
+        # Disable GET queries
+        # https://strawberry.rocks/docs/operations/deployment
+        # https://strawberry.rocks/docs/integrations/flask
+        allow_queries_via_get=False,
+        # Strawberry recently changed multipart request handling, which now
+        # requires enabling support explicitly for views.
+        # https://github.com/strawberry-graphql/strawberry/issues/3655
+        multipart_uploads_enabled=True,
+    ),
+)
+
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=5000)

sam2-repo/demo/backend/server/app_conf.py

@@ -0,0 +1,55 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+APP_ROOT = os.getenv("APP_ROOT", "/opt/sam2")
+
+API_URL = os.getenv("API_URL", "http://localhost:7263")
+
+MODEL_SIZE = os.getenv("MODEL_SIZE", "base_plus")
+
+logger.info(f"using model size {MODEL_SIZE}")
+
+FFMPEG_NUM_THREADS = int(os.getenv("FFMPEG_NUM_THREADS", "1"))
+
+# Path for all data used in API
+DATA_PATH = Path(os.getenv("DATA_PATH", "/data"))
+
+# Max duration an uploaded video can have in seconds. The default is 10
+# seconds.
+MAX_UPLOAD_VIDEO_DURATION = float(os.environ.get("MAX_UPLOAD_VIDEO_DURATION", "10"))
+
+# If set, it will define which video is returned by the default video query for
+# desktop
+DEFAULT_VIDEO_PATH = os.getenv("DEFAULT_VIDEO_PATH")
+
+# Prefix for gallery videos
+GALLERY_PREFIX = "gallery"
+
+# Path where all gallery videos are stored
+GALLERY_PATH = DATA_PATH / GALLERY_PREFIX
+
+# Prefix for uploaded videos
+UPLOADS_PREFIX = "uploads"
+
+# Path where all uploaded videos are stored
+UPLOADS_PATH = DATA_PATH / UPLOADS_PREFIX
+
+# Prefix for video posters (1st frame of video)
+POSTERS_PREFIX = "posters"
+
+# Path where all posters are stored
+POSTERS_PATH = DATA_PATH / POSTERS_PREFIX
+
+# Make sure any of those paths exist
+os.makedirs(DATA_PATH, exist_ok=True)
+os.makedirs(GALLERY_PATH, exist_ok=True)
+os.makedirs(UPLOADS_PATH, exist_ok=True)
+os.makedirs(POSTERS_PATH, exist_ok=True)

sam2-repo/demo/backend/server/data/data_types.py

@@ -0,0 +1,154 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass
+from typing import Iterable, List, Optional
+
+import strawberry
+from app_conf import API_URL
+from data.resolver import resolve_videos
+from dataclasses_json import dataclass_json
+from strawberry import relay
+
+
+@strawberry.type
+class Video(relay.Node):
+    """Core type for video."""
+
+    code: relay.NodeID[str]
+    path: str
+    poster_path: Optional[str]
+    width: int
+    height: int
+
+    @strawberry.field
+    def url(self) -> str:
+        return f"{API_URL}/{self.path}"
+
+    @strawberry.field
+    def poster_url(self) -> str:
+        return f"{API_URL}/{self.poster_path}"
+
+    @classmethod
+    def resolve_nodes(
+        cls,
+        *,
+        info: relay.PageInfo,
+        node_ids: Iterable[str],
+        required: bool = False,
+    ):
+        return resolve_videos(node_ids, required)
+
+
+@strawberry.type
+class RLEMask:
+    """Core type for Onevision GraphQL RLE mask."""
+
+    size: List[int]
+    counts: str
+    order: str
+
+
+@strawberry.type
+class RLEMaskForObject:
+    """Type for RLE mask associated with a specific object id."""
+
+    object_id: int
+    rle_mask: RLEMask
+
+
+@strawberry.type
+class RLEMaskListOnFrame:
+    """Type for a list of object-associated RLE masks on a specific video frame."""
+
+    frame_index: int
+    rle_mask_list: List[RLEMaskForObject]
+
+
+@strawberry.input
+class StartSessionInput:
+    path: str
+
+
+@strawberry.type
+class StartSession:
+    session_id: str
+
+
+@strawberry.input
+class PingInput:
+    session_id: str
+
+
+@strawberry.type
+class Pong:
+    success: bool
+
+
+@strawberry.input
+class CloseSessionInput:
+    session_id: str
+
+
+@strawberry.type
+class CloseSession:
+    success: bool
+
+
+@strawberry.input
+class AddPointsInput:
+    session_id: str
+    frame_index: int
+    clear_old_points: bool
+    object_id: int
+    labels: List[int]
+    points: List[List[float]]
+
+
+@strawberry.input
+class ClearPointsInFrameInput:
+    session_id: str
+    frame_index: int
+    object_id: int
+
+
+@strawberry.input
+class ClearPointsInVideoInput:
+    session_id: str
+
+
+@strawberry.type
+class ClearPointsInVideo:
+    success: bool
+
+
+@strawberry.input
+class RemoveObjectInput:
+    session_id: str
+    object_id: int
+
+
+@strawberry.input
+class PropagateInVideoInput:
+    session_id: str
+    start_frame_index: int
+
+
+@strawberry.input
+class CancelPropagateInVideoInput:
+    session_id: str
+
+
+@strawberry.type
+class CancelPropagateInVideo:
+    success: bool
+
+
+@strawberry.type
+class SessionExpiration:
+    session_id: str
+    expiration_time: int
+    max_expiration_time: int
+    ttl: int

sam2-repo/demo/backend/server/data/loader.py

@@ -0,0 +1,92 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import shutil
+import subprocess
+from glob import glob
+from pathlib import Path
+from typing import Dict, Optional
+
+import imagesize
+from app_conf import GALLERY_PATH, POSTERS_PATH, POSTERS_PREFIX
+from data.data_types import Video
+from tqdm import tqdm
+
+
+def preload_data() -> Dict[str, Video]:
+    """
+    Preload data including gallery videos and their posters.
+    """
+    # Dictionaries for videos and datasets on the backend.
+    # Note that since Python 3.7, dictionaries preserve their insert order, so
+    # when looping over its `.values()`, elements inserted first also appear first.
+    # https://stackoverflow.com/questions/39980323/are-dictionaries-ordered-in-python-3-6
+    all_videos = {}
+
+    video_path_pattern = os.path.join(GALLERY_PATH, "**/*.mp4")
+    video_paths = glob(video_path_pattern, recursive=True)
+
+    for p in tqdm(video_paths):
+        video = get_video(p, GALLERY_PATH)
+        all_videos[video.code] = video
+
+    return all_videos
+
+
+def get_video(
+    filepath: os.PathLike,
+    absolute_path: Path,
+    file_key: Optional[str] = None,
+    generate_poster: bool = True,
+    width: Optional[int] = None,
+    height: Optional[int] = None,
+    verbose: Optional[bool] = False,
+) -> Video:
+    """
+    Get video object given
+    """
+    # Use absolute_path to include the parent directory in the video
+    video_path = os.path.relpath(filepath, absolute_path.parent)
+    poster_path = None
+    if generate_poster:
+        poster_id = os.path.splitext(os.path.basename(filepath))[0]
+        poster_filename = f"{str(poster_id)}.jpg"
+        poster_path = f"{POSTERS_PREFIX}/{poster_filename}"
+
+        # Extract the first frame from video
+        poster_output_path = os.path.join(POSTERS_PATH, poster_filename)
+        ffmpeg = shutil.which("ffmpeg")
+        subprocess.call(
+            [
+                ffmpeg,
+                "-y",
+                "-i",
+                str(filepath),
+                "-pix_fmt",
+                "yuv420p",
+                "-frames:v",
+                "1",
+                "-update",
+                "1",
+                "-strict",
+                "unofficial",
+                str(poster_output_path),
+            ],
+            stdout=None if verbose else subprocess.DEVNULL,
+            stderr=None if verbose else subprocess.DEVNULL,
+        )
+
+        # Extract video width and height from poster. This is important to optimize
+        # rendering previews in the mosaic video preview.
+        width, height = imagesize.get(poster_output_path)
+
+    return Video(
+        code=video_path,
+        path=video_path if file_key is None else file_key,
+        poster_path=poster_path,
+        width=width,
+        height=height,
+    )

sam2-repo/demo/backend/server/data/resolver.py

@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Iterable
+
+
+def resolve_videos(node_ids: Iterable[str], required: bool = False):
+    """
+    Resolve videos given node ids.
+    """
+    from data.store import get_videos
+
+    all_videos = get_videos()
+    return [
+        all_videos[nid] if required else all_videos.get(nid, None) for nid in node_ids
+    ]

sam2-repo/demo/backend/server/data/schema.py

@@ -0,0 +1,357 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import hashlib
+import os
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Iterable, List, Optional, Tuple, Union
+
+import av
+import strawberry
+from app_conf import (
+    DATA_PATH,
+    DEFAULT_VIDEO_PATH,
+    MAX_UPLOAD_VIDEO_DURATION,
+    UPLOADS_PATH,
+    UPLOADS_PREFIX,
+)
+from data.data_types import (
+    AddPointsInput,
+    CancelPropagateInVideo,
+    CancelPropagateInVideoInput,
+    ClearPointsInFrameInput,
+    ClearPointsInVideo,
+    ClearPointsInVideoInput,
+    CloseSession,
+    CloseSessionInput,
+    RemoveObjectInput,
+    RLEMask,
+    RLEMaskForObject,
+    RLEMaskListOnFrame,
+    StartSession,
+    StartSessionInput,
+    Video,
+)
+from data.loader import get_video
+from data.store import get_videos
+from data.transcoder import get_video_metadata, transcode, VideoMetadata
+from inference.data_types import (
+    AddPointsRequest,
+    CancelPropagateInVideoRequest,
+    CancelPropagateInVideoRequest,
+    ClearPointsInFrameRequest,
+    ClearPointsInVideoRequest,
+    CloseSessionRequest,
+    RemoveObjectRequest,
+    StartSessionRequest,
+)
+from inference.predictor import InferenceAPI
+from strawberry import relay
+from strawberry.file_uploads import Upload
+
+
+@strawberry.type
+class Query:
+
+    @strawberry.field
+    def default_video(self) -> Video:
+        """
+        Return the default video.
+
+        The default video can be set with the DEFAULT_VIDEO_PATH environment
+        variable. It will return the video that matches this path. If no video
+        is found, it will return the first video.
+        """
+        all_videos = get_videos()
+
+        # Find the video that matches the default path and return that as
+        # default video.
+        for _, v in all_videos.items():
+            if v.path == DEFAULT_VIDEO_PATH:
+                return v
+
+        # Fallback is returning the first video
+        return next(iter(all_videos.values()))
+
+    @relay.connection(relay.ListConnection[Video])
+    def videos(
+        self,
+    ) -> Iterable[Video]:
+        """
+        Return all available videos.
+        """
+        all_videos = get_videos()
+        return all_videos.values()
+
+
+@strawberry.type
+class Mutation:
+
+    @strawberry.mutation
+    def upload_video(
+        self,
+        file: Upload,
+        start_time_sec: Optional[float] = None,
+        duration_time_sec: Optional[float] = None,
+    ) -> Video:
+        """
+        Receive a video file and store it in the configured S3 bucket.
+        """
+        max_time = MAX_UPLOAD_VIDEO_DURATION
+        filepath, file_key, vm = process_video(
+            file,
+            max_time=max_time,
+            start_time_sec=start_time_sec,
+            duration_time_sec=duration_time_sec,
+        )
+
+        video = get_video(
+            filepath,
+            UPLOADS_PATH,
+            file_key=file_key,
+            width=vm.width,
+            height=vm.height,
+            generate_poster=False,
+        )
+        return video
+
+    @strawberry.mutation
+    def start_session(
+        self, input: StartSessionInput, info: strawberry.Info
+    ) -> StartSession:
+        inference_api: InferenceAPI = info.context["inference_api"]
+
+        request = StartSessionRequest(
+            type="start_session",
+            path=f"{DATA_PATH}/{input.path}",
+        )
+
+        response = inference_api.start_session(request=request)
+
+        return StartSession(session_id=response.session_id)
+
+    @strawberry.mutation
+    def close_session(
+        self, input: CloseSessionInput, info: strawberry.Info
+    ) -> CloseSession:
+        inference_api: InferenceAPI = info.context["inference_api"]
+
+        request = CloseSessionRequest(
+            type="close_session",
+            session_id=input.session_id,
+        )
+        response = inference_api.close_session(request)
+        return CloseSession(success=response.success)
+
+    @strawberry.mutation
+    def add_points(
+        self, input: AddPointsInput, info: strawberry.Info
+    ) -> RLEMaskListOnFrame:
+        inference_api: InferenceAPI = info.context["inference_api"]
+
+        request = AddPointsRequest(
+            type="add_points",
+            session_id=input.session_id,
+            frame_index=input.frame_index,
+            object_id=input.object_id,
+            points=input.points,
+            labels=input.labels,
+            clear_old_points=input.clear_old_points,
+        )
+        reponse = inference_api.add_points(request)
+
+        return RLEMaskListOnFrame(
+            frame_index=reponse.frame_index,
+            rle_mask_list=[
+                RLEMaskForObject(
+                    object_id=r.object_id,
+                    rle_mask=RLEMask(counts=r.mask.counts, size=r.mask.size, order="F"),
+                )
+                for r in reponse.results
+            ],
+        )
+
+    @strawberry.mutation
+    def remove_object(
+        self, input: RemoveObjectInput, info: strawberry.Info
+    ) -> List[RLEMaskListOnFrame]:
+        inference_api: InferenceAPI = info.context["inference_api"]
+
+        request = RemoveObjectRequest(
+            type="remove_object", session_id=input.session_id, object_id=input.object_id
+        )
+
+        response = inference_api.remove_object(request)
+
+        return [
+            RLEMaskListOnFrame(
+                frame_index=res.frame_index,
+                rle_mask_list=[
+                    RLEMaskForObject(
+                        object_id=r.object_id,
+                        rle_mask=RLEMask(
+                            counts=r.mask.counts, size=r.mask.size, order="F"
+                        ),
+                    )
+                    for r in res.results
+                ],
+            )
+            for res in response.results
+        ]
+
+    @strawberry.mutation
+    def clear_points_in_frame(
+        self, input: ClearPointsInFrameInput, info: strawberry.Info
+    ) -> RLEMaskListOnFrame:
+        inference_api: InferenceAPI = info.context["inference_api"]
+
+        request = ClearPointsInFrameRequest(
+            type="clear_points_in_frame",
+            session_id=input.session_id,
+            frame_index=input.frame_index,
+            object_id=input.object_id,
+        )
+
+        response = inference_api.clear_points_in_frame(request)
+
+        return RLEMaskListOnFrame(
+            frame_index=response.frame_index,
+            rle_mask_list=[
+                RLEMaskForObject(
+                    object_id=r.object_id,
+                    rle_mask=RLEMask(counts=r.mask.counts, size=r.mask.size, order="F"),
+                )
+                for r in response.results
+            ],
+        )
+
+    @strawberry.mutation
+    def clear_points_in_video(
+        self, input: ClearPointsInVideoInput, info: strawberry.Info
+    ) -> ClearPointsInVideo:
+        inference_api: InferenceAPI = info.context["inference_api"]
+
+        request = ClearPointsInVideoRequest(
+            type="clear_points_in_video",
+            session_id=input.session_id,
+        )
+        response = inference_api.clear_points_in_video(request)
+        return ClearPointsInVideo(success=response.success)
+
+    @strawberry.mutation
+    def cancel_propagate_in_video(
+        self, input: CancelPropagateInVideoInput, info: strawberry.Info
+    ) -> CancelPropagateInVideo:
+        inference_api: InferenceAPI = info.context["inference_api"]
+
+        request = CancelPropagateInVideoRequest(
+            type="cancel_propagate_in_video",
+            session_id=input.session_id,
+        )
+        response = inference_api.cancel_propagate_in_video(request)
+        return CancelPropagateInVideo(success=response.success)
+
+
+def get_file_hash(video_path_or_file) -> str:
+    if isinstance(video_path_or_file, str):
+        with open(video_path_or_file, "rb") as in_f:
+            result = hashlib.sha256(in_f.read()).hexdigest()
+    else:
+        video_path_or_file.seek(0)
+        result = hashlib.sha256(video_path_or_file.read()).hexdigest()
+    return result
+
+
+def _get_start_sec_duration_sec(
+    start_time_sec: Union[float, None],
+    duration_time_sec: Union[float, None],
+    max_time: float,
+) -> Tuple[float, float]:
+    default_seek_t = int(os.environ.get("VIDEO_ENCODE_SEEK_TIME", "0"))
+    if start_time_sec is None:
+        start_time_sec = default_seek_t
+
+    if duration_time_sec is not None:
+        duration_time_sec = min(duration_time_sec, max_time)
+    else:
+        duration_time_sec = max_time
+    return start_time_sec, duration_time_sec
+
+
+def process_video(
+    file: Upload,
+    max_time: float,
+    start_time_sec: Optional[float] = None,
+    duration_time_sec: Optional[float] = None,
+) -> Tuple[Optional[str], str, str, VideoMetadata]:
+    """
+    Process file upload including video trimming and content moderation checks.
+
+    Returns the filepath, s3_file_key, hash & video metaedata as a tuple.
+    """
+    with tempfile.TemporaryDirectory() as tempdir:
+        in_path = f"{tempdir}/in.mp4"
+        out_path = f"{tempdir}/out.mp4"
+        with open(in_path, "wb") as in_f:
+            in_f.write(file.read())
+
+        try:
+            video_metadata = get_video_metadata(in_path)
+        except av.InvalidDataError:
+            raise Exception("not valid video file")
+
+        if video_metadata.num_video_streams == 0:
+            raise Exception("video container does not contain a video stream")
+        if video_metadata.width is None or video_metadata.height is None:
+            raise Exception("video container does not contain width or height metadata")
+
+        if video_metadata.duration_sec in (None, 0):
+            raise Exception("video container does time duration metadata")
+
+        start_time_sec, duration_time_sec = _get_start_sec_duration_sec(
+            max_time=max_time,
+            start_time_sec=start_time_sec,
+            duration_time_sec=duration_time_sec,
+        )
+
+        # Transcode video to make sure videos returned to the app are all in
+        # the same format, duration, resolution, fps.
+        transcode(
+            in_path,
+            out_path,
+            video_metadata,
+            seek_t=start_time_sec,
+            duration_time_sec=duration_time_sec,
+        )
+
+        os.remove(in_path)  # don't need original video now
+
+        out_video_metadata = get_video_metadata(out_path)
+        if out_video_metadata.num_video_frames == 0:
+            raise Exception(
+                "transcode produced empty video; check seek time or your input video"
+            )
+
+        filepath = None
+        file_key = None
+        with open(out_path, "rb") as file_data:
+            file_hash = get_file_hash(file_data)
+            file_data.seek(0)
+
+            file_key = UPLOADS_PREFIX + "/" + f"{file_hash}.mp4"
+            filepath = os.path.join(UPLOADS_PATH, f"{file_hash}.mp4")
+
+        assert filepath is not None and file_key is not None
+        shutil.move(out_path, filepath)
+
+        return filepath, file_key, out_video_metadata
+
+
+schema = strawberry.Schema(
+    query=Query,
+    mutation=Mutation,
+)

sam2-repo/demo/backend/server/data/store.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Dict
+
+from data.data_types import Video
+
+ALL_VIDEOS: Dict[str, Video] = []
+
+
+def set_videos(videos: Dict[str, Video]) -> None:
+    """
+    Set the videos available in the backend. The data is kept in-memory, but a future change could replace the
+    in-memory storage with a database backend. This would also be more efficient when querying videos given a
+    dataset name etc.
+    """
+    global ALL_VIDEOS
+    ALL_VIDEOS = videos
+
+
+def get_videos() -> Dict[str, Video]:
+    """
+    Return the videos available in the backend.
+    """
+    global ALL_VIDEOS
+    return ALL_VIDEOS

sam2-repo/demo/backend/server/data/transcoder.py

@@ -0,0 +1,186 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import ast
+import math
+import os
+import shutil
+import subprocess
+from dataclasses import dataclass
+from typing import Optional
+
+import av
+from app_conf import FFMPEG_NUM_THREADS
+from dataclasses_json import dataclass_json
+
+TRANSCODE_VERSION = 1
+
+
+@dataclass_json
+@dataclass
+class VideoMetadata:
+    duration_sec: Optional[float]
+    video_duration_sec: Optional[float]
+    container_duration_sec: Optional[float]
+    fps: Optional[float]
+    width: Optional[int]
+    height: Optional[int]
+    num_video_frames: int
+    num_video_streams: int
+    video_start_time: float
+
+
+def transcode(
+    in_path: str,
+    out_path: str,
+    in_metadata: Optional[VideoMetadata],
+    seek_t: float,
+    duration_time_sec: float,
+):
+    codec = os.environ.get("VIDEO_ENCODE_CODEC", "libx264")
+    crf = int(os.environ.get("VIDEO_ENCODE_CRF", "23"))
+    fps = int(os.environ.get("VIDEO_ENCODE_FPS", "24"))
+    max_w = int(os.environ.get("VIDEO_ENCODE_MAX_WIDTH", "1280"))
+    max_h = int(os.environ.get("VIDEO_ENCODE_MAX_HEIGHT", "720"))
+    verbose = ast.literal_eval(os.environ.get("VIDEO_ENCODE_VERBOSE", "False"))
+
+    normalize_video(
+        in_path=in_path,
+        out_path=out_path,
+        max_w=max_w,
+        max_h=max_h,
+        seek_t=seek_t,
+        max_time=duration_time_sec,
+        in_metadata=in_metadata,
+        codec=codec,
+        crf=crf,
+        fps=fps,
+        verbose=verbose,
+    )
+
+
+def get_video_metadata(path: str) -> VideoMetadata:
+    with av.open(path) as cont:
+        num_video_streams = len(cont.streams.video)
+        width, height, fps = None, None, None
+        video_duration_sec = 0
+        container_duration_sec = float((cont.duration or 0) / av.time_base)
+        video_start_time = 0.0
+        rotation_deg = 0
+        num_video_frames = 0
+        if num_video_streams > 0:
+            video_stream = cont.streams.video[0]
+            assert video_stream.time_base is not None
+
+            # for rotation, see: https://github.com/PyAV-Org/PyAV/pull/1249
+            rotation_deg = video_stream.side_data.get("DISPLAYMATRIX", 0)
+            num_video_frames = video_stream.frames
+            video_start_time = float(video_stream.start_time * video_stream.time_base)
+            width, height = video_stream.width, video_stream.height
+            fps = float(video_stream.guessed_rate)
+            fps_avg = video_stream.average_rate
+            if video_stream.duration is not None:
+                video_duration_sec = float(
+                    video_stream.duration * video_stream.time_base
+                )
+            if fps is None:
+                fps = float(fps_avg)
+
+            if not math.isnan(rotation_deg) and int(rotation_deg) in (
+                90,
+                -90,
+                270,
+                -270,
+            ):
+                width, height = height, width
+
+        duration_sec = max(container_duration_sec, video_duration_sec)
+
+        return VideoMetadata(
+            duration_sec=duration_sec,
+            container_duration_sec=container_duration_sec,
+            video_duration_sec=video_duration_sec,
+            video_start_time=video_start_time,
+            fps=fps,
+            width=width,
+            height=height,
+            num_video_streams=num_video_streams,
+            num_video_frames=num_video_frames,
+        )
+
+
+def normalize_video(
+    in_path: str,
+    out_path: str,
+    max_w: int,
+    max_h: int,
+    seek_t: float,
+    max_time: float,
+    in_metadata: Optional[VideoMetadata],
+    codec: str = "libx264",
+    crf: int = 23,
+    fps: int = 24,
+    verbose: bool = False,
+):
+    if in_metadata is None:
+        in_metadata = get_video_metadata(in_path)
+
+    assert in_metadata.num_video_streams > 0, "no video stream present"
+
+    w, h = in_metadata.width, in_metadata.height
+    assert w is not None, "width not available"
+    assert h is not None, "height not available"
+
+    # rescale to max_w:max_h if needed & preserve aspect ratio
+    r = w / h
+    if r < 1:
+        h = min(720, h)
+        w = h * r
+    else:
+        w = min(1280, w)
+        h = w / r
+
+    # h264 cannot encode w/ odd dimensions
+    w = int(w)
+    h = int(h)
+    if w % 2 != 0:
+        w += 1
+    if h % 2 != 0:
+        h += 1
+
+    ffmpeg = shutil.which("ffmpeg")
+    cmd = [
+        ffmpeg,
+        "-threads",
+        f"{FFMPEG_NUM_THREADS}",  # global threads
+        "-ss",
+        f"{seek_t:.2f}",
+        "-t",
+        f"{max_time:.2f}",
+        "-i",
+        in_path,
+        "-threads",
+        f"{FFMPEG_NUM_THREADS}",  # decode (or filter..?) threads
+        "-vf",
+        f"fps={fps},scale={w}:{h},setsar=1:1",
+        "-c:v",
+        codec,
+        "-crf",
+        f"{crf}",
+        "-pix_fmt",
+        "yuv420p",
+        "-threads",
+        f"{FFMPEG_NUM_THREADS}",  # encode threads
+        out_path,
+        "-y",
+    ]
+    if verbose:
+        print(" ".join(cmd))
+
+    subprocess.call(
+        cmd,
+        stdout=None if verbose else subprocess.DEVNULL,
+        stderr=None if verbose else subprocess.DEVNULL,
+    )

sam2-repo/demo/backend/server/inference/data_types.py

@@ -0,0 +1,191 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+from dataclasses_json import dataclass_json
+from torch import Tensor
+
+
+@dataclass_json
+@dataclass
+class Mask:
+    size: List[int]
+    counts: str
+
+
+@dataclass_json
+@dataclass
+class BaseRequest:
+    type: str
+
+
+@dataclass_json
+@dataclass
+class StartSessionRequest(BaseRequest):
+    type: str
+    path: str
+    session_id: Optional[str] = None
+
+
+@dataclass_json
+@dataclass
+class SaveSessionRequest(BaseRequest):
+    type: str
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class LoadSessionRequest(BaseRequest):
+    type: str
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class RenewSessionRequest(BaseRequest):
+    type: str
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class CloseSessionRequest(BaseRequest):
+    type: str
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class AddPointsRequest(BaseRequest):
+    type: str
+    session_id: str
+    frame_index: int
+    clear_old_points: bool
+    object_id: int
+    labels: List[int]
+    points: List[List[float]]
+
+
+@dataclass_json
+@dataclass
+class AddMaskRequest(BaseRequest):
+    type: str
+    session_id: str
+    frame_index: int
+    object_id: int
+    mask: Mask
+
+
+@dataclass_json
+@dataclass
+class ClearPointsInFrameRequest(BaseRequest):
+    type: str
+    session_id: str
+    frame_index: int
+    object_id: int
+
+
+@dataclass_json
+@dataclass
+class ClearPointsInVideoRequest(BaseRequest):
+    type: str
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class RemoveObjectRequest(BaseRequest):
+    type: str
+    session_id: str
+    object_id: int
+
+
+@dataclass_json
+@dataclass
+class PropagateInVideoRequest(BaseRequest):
+    type: str
+    session_id: str
+    start_frame_index: int
+
+
+@dataclass_json
+@dataclass
+class CancelPropagateInVideoRequest(BaseRequest):
+    type: str
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class StartSessionResponse:
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class SaveSessionResponse:
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class LoadSessionResponse:
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class RenewSessionResponse:
+    session_id: str
+
+
+@dataclass_json
+@dataclass
+class CloseSessionResponse:
+    success: bool
+
+
+@dataclass_json
+@dataclass
+class ClearPointsInVideoResponse:
+    success: bool
+
+
+@dataclass_json
+@dataclass
+class PropagateDataValue:
+    object_id: int
+    mask: Mask
+
+
+@dataclass_json
+@dataclass
+class PropagateDataResponse:
+    frame_index: int
+    results: List[PropagateDataValue]
+
+
+@dataclass_json
+@dataclass
+class RemoveObjectResponse:
+    results: List[PropagateDataResponse]
+
+
+@dataclass_json
+@dataclass
+class CancelPorpagateResponse:
+    success: bool
+
+
+@dataclass_json
+@dataclass
+class InferenceSession:
+    start_time: float
+    last_use_time: float
+    session_id: str
+    state: Dict[str, Dict[str, Union[Tensor, Dict[int, Tensor]]]]

sam2-repo/demo/backend/server/inference/multipart.py

@@ -0,0 +1,48 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Dict, Union
+
+
+class MultipartResponseBuilder:
+    message: bytes
+
+    def __init__(self, boundary: str) -> None:
+        self.message = b"--" + boundary.encode("utf-8") + b"\r\n"
+
+    @classmethod
+    def build(
+        cls, boundary: str, headers: Dict[str, str], body: Union[str, bytes]
+    ) -> "MultipartResponseBuilder":
+        builder = cls(boundary=boundary)
+        for k, v in headers.items():
+            builder.__append_header(key=k, value=v)
+        if isinstance(body, bytes):
+            builder.__append_body(body)
+        elif isinstance(body, str):
+            builder.__append_body(body.encode("utf-8"))
+        else:
+            raise ValueError(
+                f"body needs to be of type bytes or str but got {type(body)}"
+            )
+
+        return builder
+
+    def get_message(self) -> bytes:
+        return self.message
+
+    def __append_header(self, key: str, value: str) -> "MultipartResponseBuilder":
+        self.message += key.encode("utf-8") + b": " + value.encode("utf-8") + b"\r\n"
+        return self
+
+    def __close_header(self) -> "MultipartResponseBuilder":
+        self.message += b"\r\n"
+        return self
+
+    def __append_body(self, body: bytes) -> "MultipartResponseBuilder":
+        self.__append_header(key="Content-Length", value=str(len(body)))
+        self.__close_header()
+        self.message += body
+        return self

sam2-repo/demo/backend/server/inference/predictor.py

@@ -0,0 +1,427 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import contextlib
+import logging
+import os
+import uuid
+from pathlib import Path
+from threading import Lock
+from typing import Any, Dict, Generator, List
+
+import numpy as np
+import torch
+from app_conf import APP_ROOT, MODEL_SIZE
+from inference.data_types import (
+    AddMaskRequest,
+    AddPointsRequest,
+    CancelPorpagateResponse,
+    CancelPropagateInVideoRequest,
+    ClearPointsInFrameRequest,
+    ClearPointsInVideoRequest,
+    ClearPointsInVideoResponse,
+    CloseSessionRequest,
+    CloseSessionResponse,
+    Mask,
+    PropagateDataResponse,
+    PropagateDataValue,
+    PropagateInVideoRequest,
+    RemoveObjectRequest,
+    RemoveObjectResponse,
+    StartSessionRequest,
+    StartSessionResponse,
+)
+from pycocotools.mask import decode as decode_masks, encode as encode_masks
+from sam2.build_sam import build_sam2_video_predictor
+
+
+logger = logging.getLogger(__name__)
+
+
+class InferenceAPI:
+
+    def __init__(self) -> None:
+        super(InferenceAPI, self).__init__()
+
+        self.session_states: Dict[str, Any] = {}
+        self.score_thresh = 0
+
+        if MODEL_SIZE == "tiny":
+            checkpoint = Path(APP_ROOT) / "checkpoints/sam2.1_hiera_tiny.pt"
+            model_cfg = "configs/sam2.1/sam2.1_hiera_t.yaml"
+        elif MODEL_SIZE == "small":
+            checkpoint = Path(APP_ROOT) / "checkpoints/sam2.1_hiera_small.pt"
+            model_cfg = "configs/sam2.1/sam2.1_hiera_s.yaml"
+        elif MODEL_SIZE == "large":
+            checkpoint = Path(APP_ROOT) / "checkpoints/sam2.1_hiera_large.pt"
+            model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
+        else:  # base_plus (default)
+            checkpoint = Path(APP_ROOT) / "checkpoints/sam2.1_hiera_base_plus.pt"
+            model_cfg = "configs/sam2.1/sam2.1_hiera_b+.yaml"
+
+        # select the device for computation
+        force_cpu_device = os.environ.get("SAM2_DEMO_FORCE_CPU_DEVICE", "0") == "1"
+        if force_cpu_device:
+            logger.info("forcing CPU device for SAM 2 demo")
+        if torch.cuda.is_available() and not force_cpu_device:
+            device = torch.device("cuda")
+        elif torch.backends.mps.is_available() and not force_cpu_device:
+            device = torch.device("mps")
+        else:
+            device = torch.device("cpu")
+        logger.info(f"using device: {device}")
+
+        if device.type == "cuda":
+            # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
+            if torch.cuda.get_device_properties(0).major >= 8:
+                torch.backends.cuda.matmul.allow_tf32 = True
+                torch.backends.cudnn.allow_tf32 = True
+        elif device.type == "mps":
+            logging.warning(
+                "\nSupport for MPS devices is preliminary. SAM 2 is trained with CUDA and might "
+                "give numerically different outputs and sometimes degraded performance on MPS. "
+                "See e.g. https://github.com/pytorch/pytorch/issues/84936 for a discussion."
+            )
+
+        self.device = device
+        self.predictor = build_sam2_video_predictor(
+            model_cfg, checkpoint, device=device
+        )
+        self.inference_lock = Lock()
+
+    def autocast_context(self):
+        if self.device.type == "cuda":
+            return torch.autocast("cuda", dtype=torch.bfloat16)
+        else:
+            return contextlib.nullcontext()
+
+    def start_session(self, request: StartSessionRequest) -> StartSessionResponse:
+        with self.autocast_context(), self.inference_lock:
+            session_id = str(uuid.uuid4())
+            # for MPS devices, we offload the video frames to CPU by default to avoid
+            # memory fragmentation in MPS (which sometimes crashes the entire process)
+            offload_video_to_cpu = self.device.type == "mps"
+            inference_state = self.predictor.init_state(
+                request.path,
+                offload_video_to_cpu=offload_video_to_cpu,
+            )
+            self.session_states[session_id] = {
+                "canceled": False,
+                "state": inference_state,
+            }
+            return StartSessionResponse(session_id=session_id)
+
+    def close_session(self, request: CloseSessionRequest) -> CloseSessionResponse:
+        is_successful = self.__clear_session_state(request.session_id)
+        return CloseSessionResponse(success=is_successful)
+
+    def add_points(
+        self, request: AddPointsRequest, test: str = ""
+    ) -> PropagateDataResponse:
+        with self.autocast_context(), self.inference_lock:
+            session = self.__get_session(request.session_id)
+            inference_state = session["state"]
+
+            frame_idx = request.frame_index
+            obj_id = request.object_id
+            points = request.points
+            labels = request.labels
+            clear_old_points = request.clear_old_points
+
+            # add new prompts and instantly get the output on the same frame
+            frame_idx, object_ids, masks = self.predictor.add_new_points_or_box(
+                inference_state=inference_state,
+                frame_idx=frame_idx,
+                obj_id=obj_id,
+                points=points,
+                labels=labels,
+                clear_old_points=clear_old_points,
+                normalize_coords=False,
+            )
+
+            masks_binary = (masks > self.score_thresh)[:, 0].cpu().numpy()
+
+            rle_mask_list = self.__get_rle_mask_list(
+                object_ids=object_ids, masks=masks_binary
+            )
+
+            return PropagateDataResponse(
+                frame_index=frame_idx,
+                results=rle_mask_list,
+            )
+
+    def add_mask(self, request: AddMaskRequest) -> PropagateDataResponse:
+        """
+        Add new points on a specific video frame.
+        - mask is a numpy array of shape [H_im, W_im] (containing 1 for foreground and 0 for background).
+        Note: providing an input mask would overwrite any previous input points on this frame.
+        """
+        with self.autocast_context(), self.inference_lock:
+            session_id = request.session_id
+            frame_idx = request.frame_index
+            obj_id = request.object_id
+            rle_mask = {
+                "counts": request.mask.counts,
+                "size": request.mask.size,
+            }
+
+            mask = decode_masks(rle_mask)
+
+            logger.info(
+                f"add mask on frame {frame_idx} in session {session_id}: {obj_id=}, {mask.shape=}"
+            )
+            session = self.__get_session(session_id)
+            inference_state = session["state"]
+
+            frame_idx, obj_ids, video_res_masks = self.model.add_new_mask(
+                inference_state=inference_state,
+                frame_idx=frame_idx,
+                obj_id=obj_id,
+                mask=torch.tensor(mask > 0),
+            )
+            masks_binary = (video_res_masks > self.score_thresh)[:, 0].cpu().numpy()
+
+            rle_mask_list = self.__get_rle_mask_list(
+                object_ids=obj_ids, masks=masks_binary
+            )
+
+            return PropagateDataResponse(
+                frame_index=frame_idx,
+                results=rle_mask_list,
+            )
+
+    def clear_points_in_frame(
+        self, request: ClearPointsInFrameRequest
+    ) -> PropagateDataResponse:
+        """
+        Remove all input points in a specific frame.
+        """
+        with self.autocast_context(), self.inference_lock:
+            session_id = request.session_id
+            frame_idx = request.frame_index
+            obj_id = request.object_id
+
+            logger.info(
+                f"clear inputs on frame {frame_idx} in session {session_id}: {obj_id=}"
+            )
+            session = self.__get_session(session_id)
+            inference_state = session["state"]
+            frame_idx, obj_ids, video_res_masks = (
+                self.predictor.clear_all_prompts_in_frame(
+                    inference_state, frame_idx, obj_id
+                )
+            )
+            masks_binary = (video_res_masks > self.score_thresh)[:, 0].cpu().numpy()
+
+            rle_mask_list = self.__get_rle_mask_list(
+                object_ids=obj_ids, masks=masks_binary
+            )
+
+            return PropagateDataResponse(
+                frame_index=frame_idx,
+                results=rle_mask_list,
+            )
+
+    def clear_points_in_video(
+        self, request: ClearPointsInVideoRequest
+    ) -> ClearPointsInVideoResponse:
+        """
+        Remove all input points in all frames throughout the video.
+        """
+        with self.autocast_context(), self.inference_lock:
+            session_id = request.session_id
+            logger.info(f"clear all inputs across the video in session {session_id}")
+            session = self.__get_session(session_id)
+            inference_state = session["state"]
+            self.predictor.reset_state(inference_state)
+            return ClearPointsInVideoResponse(success=True)
+
+    def remove_object(self, request: RemoveObjectRequest) -> RemoveObjectResponse:
+        """
+        Remove an object id from the tracking state.
+        """
+        with self.autocast_context(), self.inference_lock:
+            session_id = request.session_id
+            obj_id = request.object_id
+            logger.info(f"remove object in session {session_id}: {obj_id=}")
+            session = self.__get_session(session_id)
+            inference_state = session["state"]
+            new_obj_ids, updated_frames = self.predictor.remove_object(
+                inference_state, obj_id
+            )
+
+            results = []
+            for frame_index, video_res_masks in updated_frames:
+                masks = (video_res_masks > self.score_thresh)[:, 0].cpu().numpy()
+                rle_mask_list = self.__get_rle_mask_list(
+                    object_ids=new_obj_ids, masks=masks
+                )
+                results.append(
+                    PropagateDataResponse(
+                        frame_index=frame_index,
+                        results=rle_mask_list,
+                    )
+                )
+
+            return RemoveObjectResponse(results=results)
+
+    def propagate_in_video(
+        self, request: PropagateInVideoRequest
+    ) -> Generator[PropagateDataResponse, None, None]:
+        session_id = request.session_id
+        start_frame_idx = request.start_frame_index
+        propagation_direction = "both"
+        max_frame_num_to_track = None
+
+        """
+        Propagate existing input points in all frames to track the object across video.
+        """
+
+        # Note that as this method is a generator, we also need to use autocast_context
+        # in caller to this method to ensure that it's called under the correct context
+        # (we've added `autocast_context` to `gen_track_with_mask_stream` in app.py).
+        with self.autocast_context(), self.inference_lock:
+            logger.info(
+                f"propagate in video in session {session_id}: "
+                f"{propagation_direction=}, {start_frame_idx=}, {max_frame_num_to_track=}"
+            )
+
+            try:
+                session = self.__get_session(session_id)
+                session["canceled"] = False
+
+                inference_state = session["state"]
+                if propagation_direction not in ["both", "forward", "backward"]:
+                    raise ValueError(
+                        f"invalid propagation direction: {propagation_direction}"
+                    )
+
+                # First doing the forward propagation
+                if propagation_direction in ["both", "forward"]:
+                    for outputs in self.predictor.propagate_in_video(
+                        inference_state=inference_state,
+                        start_frame_idx=start_frame_idx,
+                        max_frame_num_to_track=max_frame_num_to_track,
+                        reverse=False,
+                    ):
+                        if session["canceled"]:
+                            return None
+
+                        frame_idx, obj_ids, video_res_masks = outputs
+                        masks_binary = (
+                            (video_res_masks > self.score_thresh)[:, 0].cpu().numpy()
+                        )
+
+                        rle_mask_list = self.__get_rle_mask_list(
+                            object_ids=obj_ids, masks=masks_binary
+                        )
+
+                        yield PropagateDataResponse(
+                            frame_index=frame_idx,
+                            results=rle_mask_list,
+                        )
+
+                # Then doing the backward propagation (reverse in time)
+                if propagation_direction in ["both", "backward"]:
+                    for outputs in self.predictor.propagate_in_video(
+                        inference_state=inference_state,
+                        start_frame_idx=start_frame_idx,
+                        max_frame_num_to_track=max_frame_num_to_track,
+                        reverse=True,
+                    ):
+                        if session["canceled"]:
+                            return None
+
+                        frame_idx, obj_ids, video_res_masks = outputs
+                        masks_binary = (
+                            (video_res_masks > self.score_thresh)[:, 0].cpu().numpy()
+                        )
+
+                        rle_mask_list = self.__get_rle_mask_list(
+                            object_ids=obj_ids, masks=masks_binary
+                        )
+
+                        yield PropagateDataResponse(
+                            frame_index=frame_idx,
+                            results=rle_mask_list,
+                        )
+            finally:
+                # Log upon completion (so that e.g. we can see if two propagations happen in parallel).
+                # Using `finally` here to log even when the tracking is aborted with GeneratorExit.
+                logger.info(
+                    f"propagation ended in session {session_id}; {self.__get_session_stats()}"
+                )
+
+    def cancel_propagate_in_video(
+        self, request: CancelPropagateInVideoRequest
+    ) -> CancelPorpagateResponse:
+        session = self.__get_session(request.session_id)
+        session["canceled"] = True
+        return CancelPorpagateResponse(success=True)
+
+    def __get_rle_mask_list(
+        self, object_ids: List[int], masks: np.ndarray
+    ) -> List[PropagateDataValue]:
+        """
+        Return a list of data values, i.e. list of object/mask combos.
+        """
+        return [
+            self.__get_mask_for_object(object_id=object_id, mask=mask)
+            for object_id, mask in zip(object_ids, masks)
+        ]
+
+    def __get_mask_for_object(
+        self, object_id: int, mask: np.ndarray
+    ) -> PropagateDataValue:
+        """
+        Create a data value for an object/mask combo.
+        """
+        mask_rle = encode_masks(np.array(mask, dtype=np.uint8, order="F"))
+        mask_rle["counts"] = mask_rle["counts"].decode()
+        return PropagateDataValue(
+            object_id=object_id,
+            mask=Mask(
+                size=mask_rle["size"],
+                counts=mask_rle["counts"],
+            ),
+        )
+
+    def __get_session(self, session_id: str):
+        session = self.session_states.get(session_id, None)
+        if session is None:
+            raise RuntimeError(
+                f"Cannot find session {session_id}; it might have expired"
+            )
+        return session
+
+    def __get_session_stats(self):
+        """Get a statistics string for live sessions and their GPU usage."""
+        # print both the session ids and their video frame numbers
+        live_session_strs = [
+            f"'{session_id}' ({session['state']['num_frames']} frames, "
+            f"{len(session['state']['obj_ids'])} objects)"
+            for session_id, session in self.session_states.items()
+        ]
+        session_stats_str = (
+            "Test String Here - -"
+            f"live sessions: [{', '.join(live_session_strs)}], GPU memory: "
+            f"{torch.cuda.memory_allocated() // 1024**2} MiB used and "
+            f"{torch.cuda.memory_reserved() // 1024**2} MiB reserved"
+            f" (max over time: {torch.cuda.max_memory_allocated() // 1024**2} MiB used "
+            f"and {torch.cuda.max_memory_reserved() // 1024**2} MiB reserved)"
+        )
+        return session_stats_str
+
+    def __clear_session_state(self, session_id: str) -> bool:
+        session = self.session_states.pop(session_id, None)
+        if session is None:
+            logger.warning(
+                f"cannot close session {session_id} as it does not exist (it might have expired); "
+                f"{self.__get_session_stats()}"
+            )
+            return False
+        else:
+            logger.info(f"removed session {session_id}; {self.__get_session_stats()}")
+            return True