DOCSEARCHbyAlgolia
Submit Docs

EdgeTAM

repository·main·Indexed 21 days ago

https://github.com/facebookresearch/edgetam

An on-device executable variant of SAM 2 optimized for high-speed promptable segmentation and tracking in videos on mobile hardware. It includes tools for exporting to CoreML format for iOS and macOS deployment, a Gradio demo, and support for semi-supervised video object segmentation (VOS) inference on datasets such as DAVIS, MOSE, and SA-V.

Tokens
3.3K
Snippets
12
Records
13
Agent score
74%

What's inside EdgeTAM

  1. Install EdgeTAM

    main

    EdgeTAM requires python>=3.10, torch>=2.3.1, and torchvision>=0.18.1. It is recommended to use an Anaconda environment.

    To install the core package on a GPU machine:

    git clone https://github.com/facebookresearch/EdgeTAM.git && cd EdgeTAM

pip install -e .

    Notes:

    • If you have a PyTorch version lower than 2.3.1, the installation will attempt to upgrade it.
    • The installation requires nvcc to compile a custom CUDA kernel. Ensure your CUDA toolkit version matches your PyTorch CUDA version.
    • If you see Failed to build the SAM 2 CUDA extension, you can ignore it; the model will still work, though some post-processing might be limited.

  3. Export EdgeTAM to CoreML format

    main

    To prepare EdgeTAM for iOS or macOS deployment, export the model to CoreML format using the export_to_coreml.py script. This process generates three separate .mlpackage files located in the ./coreml_models/ directory:

    • edgetam_image_encoder.mlpackage (9.6MB)
    • edgetam_prompt_encoder.mlpackage (2.0MB)
    • edgetam_mask_decoder.mlpackage (9.8MB)

    Ensure you have the EdgeTAM checkpoint and the appropriate SAM2 configuration file available.

    python coreml/export_to_coreml.py \
  --sam2_cfg sam2/configs/edgetam.yaml \
  --sam2_checkpoint checkpoints/edgetam.pt

  4. Run the EdgeTAM Gradio demo

    main

    To run a local on-device demo via Gradio:

    1. Install Gradio dependencies: pip install -e ".[gradio]".
    2. Run the application: python gradio_app.py.

    The demo is available by default at http://127.0.0.1:7860/. You can specify a custom port using the --port argument.

    python gradio_app.py

  5. Install EdgeTAM optional dependencies

    main

    Depending on your use case, you may need additional dependencies:

    • For Jupyter Notebooks and Matplotlib:
      pip install -e ".[notebooks]"
    • For CoreML export (iOS/macOS deployment):
      pip install -e ".[coreml]"
    • For the Gradio demo:
      pip install -e ".[gradio]"
    pip install -e ".[notebooks]"
pip install -e ".[coreml]"
pip install -e ".[gradio]"

  6. Export EdgeTAM to CoreML for iOS/macOS deployment

    main

    Use the export_to_coreml.py script to convert EdgeTAM models into CoreML format for on-device inference on Apple platforms. This process generates three optimized .mlpackage models:

    1. edgetam_image_encoder.mlpackage: Processes input images into feature embeddings.
    2. edgetam_prompt_encoder.mlpackage: Handles user prompts (points, boxes, masks).
    3. edgetam_mask_decoder.mlpackage: Generates segmentation masks from features and prompts.

    Setup and Execution

    First, install the necessary CoreML dependencies:

    pip install -e ".[coreml]"

    Then, run the export script. Note: You must use the EdgeTAM config (sam2/configs/edgetam.yaml) rather than SAM 2.1 configs, as SAM 2.1 configs fail due to unsupported bicubic interpolation.

    Arguments

    • --sam2_cfg: Path to the EdgeTAM config (e.g., sam2/configs/edgetam.yaml).
    • --sam2_checkpoint: Path to the EdgeTAM checkpoint.
    • --output_dir: Directory where the .mlpackage files will be saved.
    • --validate: Enables validation during export.
    • --compute_units {ALL,CPU_ONLY,CPU_AND_GPU,CPU_AND_NE}: Target compute units (default: ALL).
    • --deployment_target {iOS16,iOS17,macOS13,macOS14}: Minimum deployment target (default: iOS16).
    • --test_image path/to/image.jpg: Path to an image used for validation.
    # Install dependencies
pip install -e ".[coreml]"

# Export models
python ./tools/export_to_coreml.py \
  --sam2_cfg sam2/configs/edgetam.yaml \
  --sam2_checkpoint ./checkpoints/edgetam.pt \
  --output_dir ./coreml_models \
  --validate

  7. Perform semi-supervised VOS inference with vos_inference.py

    main

    The vos_inference.py script generates predictions for semi-supervised video object segmentation (VOS) evaluation on datasets like DAVIS, MOSE, or SA-V. It saves prediction PNG files to the directory specified by --output_mask_dir.

    Key Flags

    • --sam2_cfg: Path to the SAM 2 configuration file.
    • --sam2_checkpoint: Path to the SAM 2 checkpoint file.
    • --base_video_dir: Directory containing the video frames.
    • --input_mask_dir: Directory containing the input object masks.
    • --video_list_file: File containing the list of videos to process.
    • --output_mask_dir: Directory where prediction PNGs will be saved.
    • --per_obj_png_file: Required for SA-V dataset. When this flag is present, the script saves per-object PNG files for both input and output masks.
    • --track_object_appearing_later_in_video: Required for LVOS or YouTube-VOS. Use this flag if objects to be tracked do not all appear in the first frame (frame 0) of the video.
    # Example for DAVIS 2017
python ./tools/vos_inference.py \
  --sam2_cfg configs/sam2.1/sam2.1_hiera_b+.yaml \
  --sam2_checkpoint ./checkpoints/sam2.1_hiera_base_plus.pt \
  --base_video_dir /path-to-davis-2017/JPEGImages/480p \
  --input_mask_dir /path-to-davis-2017/Annotations/480p \
  --video_list_file /path-to-davis-2017/ImageSets/2017/val.txt \
  --output_mask_dir ./outputs/davis_2017_pred_pngs

# Example for SA-V (requires --per_obj_png_file)
python ./tools/vos_inference.py \
  --sam2_cfg configs/sam2.1/sam2.1_hiera_b+.yaml \
  --sam2_checkpoint ./checkpoints/sam2.1_hiera_base_plus.pt \
  --base_video_dir /path-to-sav-val/JPEGImages_24fps \
  --input_mask_dir /path-to-sav-val/Annotations_6fps \
  --video_list_file /path-to-sav-val/sav_val.txt \
  --per_obj_png_file \
  --output_mask_dir ./outputs/sav_val_pred_pngs

  8. Export EdgeTAM to CoreML for iOS/macOS

    main

    You can export EdgeTAM to CoreML format to enable hardware-accelerated on-device inference on Apple devices. This process generates three optimized models:

    • Image Encoder: Processes input images to feature embeddings (~9.6MB)
    • Prompt Encoder: Handles user prompts (points, boxes, masks) (~2MB)
    • Mask Decoder: Generates segmentation masks from features (~8MB)

    Run the export script with the following arguments:

    python ./coreml/export_to_coreml.py \
  --sam2_cfg ./sam2/configs/edgetam.yaml \
  --sam2_checkpoint ./checkpoints/edgetam.pt \
  --output_dir ./coreml_models

  10. Use exported CoreML models for image segmentation

    main

    Once exported, you can load the three CoreML models using coremltools. To perform segmentation with a point prompt, first pass a resized image through the image_encoder to obtain the encoder output.

    import coremltools as ct
from PIL import Image

# Load models
image_encoder = ct.models.MLModel("coreml_models/edgetam_image_encoder.mlpackage")
prompt_encoder = ct.models.MLModel("coreml_models/edgetam_prompt_encoder.mlpackage")
mask_decoder = ct.models.MLModel("coreml_models/edgetam_mask_decoder.mlpackage")

# Segment with point prompt
image = Image.open("image.jpg").resize((1024, 1024))
encoder_out = image_encoder.predict({"image": image})

# Note: Further steps involve using prompt_encoder and mask_decoder

  11. Perform image prediction with SAM2ImagePredictor

    main

    EdgeTAM provides a SAM2ImagePredictor class for static image segmentation, following an interface similar to SAM.

    To use it, initialize the predictor with build_sam2 using your checkpoint and config, then use set_image and predict within an inference mode block.

    import torch
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

checkpoint = "./checkpoints/edgetam.pt"
model_cfg = "configs/edgetam.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>)

  12. Perform video prediction and tracking

    main

    For video segmentation and tracking, use build_sam2_video_predictor. This API allows you to initialize a state for a video and then either add new prompts (points or boxes) to specific frames or propagate existing masklets throughout the entire video sequence.

    EdgeTAM supports tracking multiple objects simultaneously using an inference state.

    import torch
from sam2.build_sam import build_sam2_video_predictor

checkpoint = "./checkpoints/edgetam.pt"
model_cfg = "configs/edgetam.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):
        ...