micro_sam.training.trainable_sam

  1from typing import Any, Dict, List, Tuple
  2
  3import torch
  4from torch import nn
  5from torch.nn import functional as F
  6
  7from segment_anything.modeling import Sam
  8from segment_anything.utils.transforms import ResizeLongestSide
  9
 10
 11# simple wrapper around SAM in order to keep things trainable
 12class TrainableSAM(nn.Module):
 13    """Wrapper to make the SegmentAnything model trainable.
 14
 15    Args:
 16        sam: The SegmentAnything Model.
 17    """
 18    def __init__(self, sam: Sam) -> None:
 19        super().__init__()
 20        self.sam = sam
 21        self.transform = ResizeLongestSide(sam.image_encoder.img_size)
 22
 23    def preprocess(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
 24        """Resize, normalize pixel values and pad to a square input.
 25
 26        Args:
 27            x: The input tensor.
 28
 29        Returns:
 30            The resized, normalized and padded tensor.
 31            The shape of the image after resizing.
 32        """
 33
 34        # Resize longest side to match the image encoder.
 35        x = self.transform.apply_image_torch(x)
 36        input_size = x.shape[-2:]
 37
 38        # Normalize colors
 39        x = (x - self.sam.pixel_mean.unsqueeze(0)) / self.sam.pixel_std.unsqueeze(0)
 40
 41        # Pad
 42        h, w = x.shape[-2:]
 43        padh = self.sam.image_encoder.img_size - h
 44        padw = self.sam.image_encoder.img_size - w
 45        x = F.pad(x, (0, padw, 0, padh))
 46        return x, input_size
 47
 48    def image_embeddings_oft(self, batched_inputs):
 49        # Compute the input images.
 50        input_images, input_size = self.preprocess(
 51            torch.stack([x["image"] for x in batched_inputs], dim=0).to(self.sam.device, non_blocking=True)
 52        )
 53        # Update the input size for each input in the batch.
 54        for i in range(len(batched_inputs)):
 55            batched_inputs[i]["input_size"] = input_size
 56        # Compute the image embeddings.
 57        image_embeddings = self.sam.image_encoder(input_images)
 58        return image_embeddings, batched_inputs
 59
 60    # batched inputs follow the same syntax as the input to sam.forward
 61    def forward(
 62        self, batched_inputs: List[Dict[str, Any]], image_embeddings: torch.Tensor, multimask_output: bool = False,
 63    ) -> List[Dict[str, Any]]:
 64        """Forward pass.
 65
 66        Args:
 67            batched_inputs: The batched input images and prompts.
 68            image_embeddings: The precompute image embeddings. If not passed then they will be computed.
 69            multimask_output: Whether to predict mutiple or just a single mask.
 70
 71        Returns:
 72            The predicted segmentation masks and iou values.
 73        """
 74        outputs = []
 75        for image_record, curr_embedding in zip(batched_inputs, image_embeddings):
 76            if "point_coords" in image_record:
 77                points = (
 78                    image_record["point_coords"].to(self.sam.device, non_blocking=True),
 79                    image_record["point_labels"].to(self.sam.device, non_blocking=True)
 80                )
 81            else:
 82                points = None
 83
 84            if "boxes" in image_record:
 85                boxes = image_record.get("boxes").to(self.sam.device, non_blocking=True)
 86            else:
 87                boxes = None
 88
 89            if "mask_inputs" in image_record:
 90                masks = image_record.get("mask_inputs").to(self.sam.device, non_blocking=True)
 91            else:
 92                masks = None
 93
 94            sparse_embeddings, dense_embeddings = self.sam.prompt_encoder(
 95                points=points,
 96                boxes=boxes,
 97                masks=masks,
 98            )
 99
100            low_res_masks, iou_predictions = self.sam.mask_decoder(
101                image_embeddings=curr_embedding.unsqueeze(0),
102                image_pe=self.sam.prompt_encoder.get_dense_pe(),
103                sparse_prompt_embeddings=sparse_embeddings,
104                dense_prompt_embeddings=dense_embeddings,
105                multimask_output=multimask_output,
106            )
107
108            masks = self.sam.postprocess_masks(
109                low_res_masks,
110                input_size=image_record["input_size"],
111                original_size=image_record["original_size"],
112            )
113
114            outputs.append(
115                {
116                    "low_res_masks": low_res_masks,
117                    "masks": masks,
118                    "iou_predictions": iou_predictions
119                }
120            )
121
122        return outputs