Getting Started

Installation

First, install groundingdino from its repository, this is a dependency for segmate:

NOTE: There is an issue with the setup script in the GroundingDINO repository causing it not able to install torch properly, please manually install PyTorch for now. For other issues, refer to the installation guide:

pip install -U git+https://github.com/IDEA-Research/GroundingDINO.git

Then, install segmate from PyPI:

pip install segmate

Example Usage

To use the provided code snippets, follow the steps below:

  1. Import the required modules and initialize the necessary objects:

import torch

from segmate.segmenter import SAM
from segmate.object_detector import GroundingDINO
import segmate.utils as utils

# Model checkpoint path for GroundingDINO is optional. If no path provided, it will download from HuggingFace
od = GroundingDINO(device='cuda', ckpt_path='PATH_TO_CHECKPOINT')
sm = SAM(model_type='MODEL_TYPE', checkpoint='PATH_to_CHECKPOINT', device='cuda')

  1. Perform segmentation with bounding box prompts:

masks = sm.segment(image=input_image, boxes_prompt=bbox)
utils.show_masks(image, masks)

  1. Perform segmentation with a text prompt:

bbox, _, _ = od.detect(input_image, "building", 0.30, 0.25)
masks = sm.segment(image=input_image, boxes_prompt=bbox)
utils.show_masks(masks)

  1. Perform segmentation with point prompts:

masks = sm.segment(image=input_image, points_prompt=(point_coords, point_labels))
utils.show_masks(masks)

  1. Perform segmentation with a mask prompt:

masks = sm.segment(image=input_image, mask_prompt=mask)
utils.show_masks(masks)

  1. Generate masks automatically without prompts:

masks = utils.visualize_automask(image=input_image, mask_input=input_masks)
utils.show_masks(masks)

  1. Fine-tune the SAM model on a custom dataset:

sm.fine_tune(
    train_data=train_dataset,
    original_input_size=500,
    criterion=loss,
    optimizer=optim,
    lr=1e-5,
    num_epochs=10)