Train Meta's Segment Anything Model (SAM) to segment high-fidelity masks in any domain.
Advances in the release and fine-tuning of several powerful open-source foundational models have given rise to a new paradigm in machine learning and artificial intelligence, and at the heart of this revolution are Transformer models.
Whereas once high-accuracy domain-specific models were beyond the reach of all but the most well-funded companies, today the foundational model paradigm has made it possible for students and independent researchers with modest resources to achieve results comparable to state-of-the-art proprietary models.
In this article, we explain how to apply Meta's Segment Anything Model (SAM) to the remote sensing task of river pixel segmentation. If you want to jump right in and read the code, the source files for this project are available on GitHub and the data on HuggingFace, but we recommend reading the entire article first.
The first step is to find or create a suitable dataset. Existing literature suggests that a good fine-tuning dataset for SAM contains at least 200-800 images. A key lesson learned from the last decade of advances in deep learning is that more data is better, and you can't go wrong with a large fine-tuning dataset. However, the goal behind the underlying model is to enable strong performance even on relatively small datasets.
You will also need a HuggingFace account, which you can create here. HuggingFace makes it easy to store and retrieve datasets at any time and from any device, making it easy to collaborate and reproduce.
The final requirement is a device with a GPU capable of running the training workflow. An Nvidia T4 GPU, freely available from Google Colab, is powerful enough to train the largest SAM model checkpoint (sam-vit-huge) on 1000 images over 50 epochs in under 12 hours.
To avoid losing progress due to usage limits on the hosted runtime, mount your Google Drive and save checkpoints for each model there. Alternatively, deploy and connect to a GCP virtual machine to avoid the limits altogether. If you've never used GCP before, you'll receive $300 in free credits, enough to train your model at least 12 times.
Before we begin training, we need to understand the architecture of SAM. The model contains three components: an image encoder from a mask autoencoder with minimal modifications, a flexible prompt encoder that can handle a variety of prompt types, and a fast and lightweight mask decoder. One of the motivations behind this design is to enable fast and real-time segmentation on edge devices (e.g. browsers), because the image embeddings only need to be computed once and the mask decoder can run on a CPU in about 50 milliseconds.
In theory, the image encoder has already learned how to best embed an image by identifying shapes, edges, and other common visual features. Similarly, in theory, the prompt encoder can optimally encode a prompt. The mask decoder is the part of the model architecture that takes these image-prompt embeddings and manipulates the image-prompt embeddings to actually create the mask.
Therefore, one approach is to fix the model parameters associated with the image and prompt encoders during training and only update the weights of the mask decoder. This approach has the advantage that the control point and bounding box prompts can be automated or used by humans, allowing for both supervised and unsupervised downstream tasks.
An alternative approach is to overload the prompt encoder and freeze the image encoder and mask decoder, and not use the original SAM mask encoder. For example, the AutoSAM architecture uses a network based on a Harmonic Dense Net to generate prompt embeddings based on the image itself. This tutorial covers the first approach, freezing the image encoder and prompt encoder and training only the mask decoder, but the code for this alternative approach can be found on the AutoSAM GitHub and in the paper.
The next step is to determine what kind of prompts your model should receive at inference time, so that you can provide those types of prompts at training time. Given the unpredictable and inconsistent nature of natural language processing, I personally don't recommend using text prompts for a full-blown computer vision pipeline. In this case, you're left with points and bounding boxes, the choice of which ultimately depends on the specific nature of your particular dataset. However, the literature has shown that bounding boxes perform fairly consistently better than control points.
The reason for this is not entirely clear, but it could be due to one or a combination of the following factors:
- Good control points are harder to choose at inference time (when the ground truth mask is unknown) than bounding boxes.
- The space of possible point prompts is orders of magnitude larger than the space of possible bounding box prompts, and therefore has not been trained as thoroughly.
- The original SAM authors focused on the zero-shot and few-shot (counting in terms of human-prompted interactions) features of their model, so pre-training may have focused on bounding boxes.
In any case, river segmentation is a rare case where point prompts are actually better than bounding boxes (but only marginally, even in very favorable domains). In an image of a river, the water body stretches from one edge of the image to the other, so the bounding box that surrounds it will almost always cover a large portion of the image. Thus, bounding box prompts for very different parts of the river can be very similar. In theory, bounding boxes provide the model with significantly less information than control points, resulting in poorer performance.
Notice in the image above that although the actual segmentation masks for the two river sections are completely different, their respective bounding boxes are nearly identical, while their point prompts are (relatively) very different.
Another important factor to consider is how easily the input prompts can be generated during inference. Bounding boxes and control points are both fairly easy to obtain during inference if you expect a human in the loop. However, if you have a fully automated pipeline, this question becomes more complicated to answer.
Whether you use control points or bounding boxes, generating a prompt typically requires that you first estimate a rough mask for the object of interest. A bounding box is the smallest box that encloses the rough mask, while control points must be sampled from the rough mask. This means that a bounding box is easier to obtain when the ground truth mask is unknown. The estimated mask for the object of interest only needs to roughly match the size and position of the actual object, whereas for control points the estimated mask must more closely match the object's contours.
For river segmentation, if you have access to both RGB and NIR, you can use the spectral index thresholding method to get a rough mask. If you only have access to RGB, you can convert the image to HSV and threshold all pixels within a certain hue, saturation and value range. Then, remove connected components below a certain size threshold, erosion from skimage.morphology Make sure that only one pixel in the mask is close to the center of the big blue blob.
To train a model, we need a data loader that contains all the training data that we can iterate over for each training epoch. When we load the dataset from HuggingFace, it has the following format: datasets.Dataset Class. If your dataset is private, first install HuggingFace CLI, !huggingface-cli login.
from datasets import load_dataset, load_from_disk, Datasethf_dataset_name = "stodoran/elwha-segmentation-v1"
training_data = load_dataset(hf_dataset_name, split="train")
validation_data = load_dataset(hf_dataset_name, split="validation")
Next, you need to code your own custom dataset class that returns not only images and labels for any index, but also a prompt. Below is an implementation that can handle both control point and bounding box prompts. It requires a HuggingFace to be initialized. datasets.Dataset Instances and SAM processor instances.
from torch.utils.data import Datasetclass PromptType:
CONTROL_POINTS = "pts"
BOUNDING_BOX = "bbox"
class SAMDataset(Dataset):
def __init__(
self,
dataset,
processor,
prompt_type = PromptType.CONTROL_POINTS,
num_positive = 3,
num_negative = 0,
erode = True,
multi_mask = "mean",
perturbation = 10,
image_size = (1024, 1024),
mask_size = (256, 256),
):
# Asign all values to self
...
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
datapoint = self.dataset[idx]
input_image = cv2.resize(np.array(datapoint["image"]), self.image_size)
ground_truth_mask = cv2.resize(np.array(datapoint["label"]), self.mask_size)
if self.prompt_type == PromptType.CONTROL_POINTS:
inputs = self._getitem_ctrlpts(input_image, ground_truth_mask)
elif self.prompt_type == PromptType.BOUNDING_BOX:
inputs = self._getitem_bbox(input_image, ground_truth_mask)
inputs["ground_truth_mask"] = ground_truth_mask
return inputs
Also, SAMDataset._getitem_ctrlpts and SAMDataset._getitem_bbox It's a function, but if you only plan on using one prompt type, you can refactor your code to handle that type directly. SAMDataset.__getitem__ Remove the helper function.
class SAMDataset(Dataset):
...def _getitem_ctrlpts(self, input_image, ground_truth_mask):
# Get control points prompt. See the GitHub for the source
# of this function, or replace with your own point selection algorithm.
input_points, input_labels = generate_input_points(
num_positive=self.num_positive,
num_negative=self.num_negative,
mask=ground_truth_mask,
dynamic_distance=True,
erode=self.erode,
)
input_points = input_points.astype(float).tolist()
input_labels = input_labels.tolist()
input_labels = [[x] for x in input_labels]
# Prepare the image and prompt for the model.
inputs = self.processor(
input_image,
input_points=input_points,
input_labels=input_labels,
return_tensors="pt"
)
# Remove batch dimension which the processor adds by default.
inputs = {k: v.squeeze(0) for k, v in inputs.items()}
inputs["input_labels"] = inputs["input_labels"].squeeze(1)
return inputs
def _getitem_bbox(self, input_image, ground_truth_mask):
# Get bounding box prompt.
bbox = get_input_bbox(ground_truth_mask, perturbation=self.perturbation)
# Prepare the image and prompt for the model.
inputs = self.processor(input_image, input_boxes=[[bbox]], return_tensors="pt")
inputs = {k: v.squeeze(0) for k, v in inputs.items()} # Remove batch dimension which the processor adds by default.
return inputs
Putting all this together, we can write a function that creates and returns a PyTorch data loader given one of the splits of the HuggingFace dataset. Writing a function that returns a data loader, rather than running cells with the same code, is not only a good way to write flexible and maintainable code, but it is also necessary if you plan to use HuggingFace Accelerate to perform distributed training.
from transformers import SamProcessor
from torch.utils.data import DataLoaderdef get_dataloader(
hf_dataset,
model_size = "base", # One of "base", "large", or "huge"
batch_size = 8,
prompt_type = PromptType.CONTROL_POINTS,
num_positive = 3,
num_negative = 0,
erode = True,
multi_mask = "mean",
perturbation = 10,
image_size = (256, 256),
mask_size = (256, 256),
):
processor = SamProcessor.from_pretrained(f"facebook/sam-vit-{model_size}")
sam_dataset = SAMDataset(
dataset=hf_dataset,
processor=processor,
prompt_type=prompt_type,
num_positive=num_positive,
num_negative=num_negative,
erode=erode,
multi_mask=multi_mask,
perturbation=perturbation,
image_size=image_size,
mask_size=mask_size,
)
dataloader = DataLoader(sam_dataset, batch_size=batch_size, shuffle=True)
return dataloader
Subsequent training simply involves loading the model, freezing the image and prompt encoders, and training for the desired number of iterations.
model = SamModel.from_pretrained(f"facebook/sam-vit-{model_size}")
optimizer = AdamW(model.mask_decoder.parameters(), lr=learning_rate, weight_decay=weight_decay)# Train only the decoder.
for name, param in model.named_parameters():
if name.startswith("vision_encoder") or name.startswith("prompt_encoder"):
param.requires_grad_(False)
Below is the basic outline of the training loop code. forward_pass, calculate loss, evaluate_modeland save_model_checkpoint Functions have been omitted for brevity, but the implementation is available on GitHub. The forward pass code differs slightly depending on the prompt type, and the loss calculation also requires special cases based on the prompt type. When using point prompts, SAM returns a predicted mask for each input point, so to get a single mask that can be compared to the ground truth, you need to either average the predicted masks or choose the best predicted mask (identified based on SAM's predicted IoU score).
train_losses = []
validation_losses = []
epoch_loop = tqdm(total=num_epochs, position=epoch, leave=False)
batch_loop = tqdm(total=len(train_dataloader), position=0, leave=True)while epoch < num_epochs:
epoch_losses = []
batch_loop.n = 0 # Loop Reset
for idx, batch in enumerate(train_dataloader):
# Forward Pass
batch = {k: v.to(accelerator.device) for k, v in batch.items()}
outputs = forward_pass(model, batch, prompt_type)
# Compute Loss
ground_truth_masks = batch["ground_truth_mask"].float()
train_loss = calculate_loss(outputs, ground_truth_masks, prompt_type, loss_fn, multi_mask="best")
epoch_losses.append(train_loss)
# Backward Pass & Optimizer Step
optimizer.zero_grad()
accelerator.backward(train_loss)
optimizer.step()
lr_scheduler.step()
batch_loop.set_description(f"Train Loss: {train_loss.item():.4f}")
batch_loop.update(1)
validation_loss = evaluate_model(model, validation_dataloader, accelerator.device, loss_fn)
train_losses.append(torch.mean(torch.Tensor(epoch_losses)))
validation_losses.append(validation_loss)
if validation_loss < best_loss:
save_model_checkpoint(
accelerator,
best_checkpoint_path,
model,
optimizer,
lr_scheduler,
epoch,
train_history,
validation_loss,
train_losses,
validation_losses,
loss_config,
model_descriptor=model_descriptor,
)
best_loss = validation_loss
epoch_loop.set_description(f"Best Loss: {best_loss:.4f}")
epoch_loop.update(1)
epoch += 1
The Elwha River project had an optimal setup, training the “sam-vit-base” model on a GCP instance in under 12 hours using a dataset of over 1,000 segmentation masks.
Compared to the baseline SAM, fine-tuning significantly improved performance, taking the central mask from unusable to highly accurate.
One important fact to note is that the training dataset of 1,000 river images is incomplete and the number of pixels correctly classified by the segmentation labels varies widely, so the metrics above were calculated on a pixel-perfect dataset of 225 river images.
An interesting behavior observed is that the model learned to generalize from incomplete training data. Evaluating data points where the training examples contained obvious misclassifications shows that the model's predictions avoided errors. Notice that the top image showing training examples contains a mask that does not completely fill the river to its banks, while in the bottom image the model predictions more closely segment the river boundary.
