micro_sam.models.simple_sam_3d_wrapper

  1from contextlib import nullcontext
  2from typing import Any, List, Dict
  3
  4import torch
  5import torch.nn as nn
  6
  7from ..util import get_sam_model
  8from .peft_sam import LoRASurgery
  9
 10
 11def get_simple_sam_3d_model(
 12    device,
 13    n_classes,
 14    image_size,
 15    lora_rank=None,
 16    freeze_encoder=False,
 17    model_type="vit_b",
 18    checkpoint_path=None,
 19):
 20    if lora_rank is None:
 21        peft_kwargs = {}
 22    else:
 23        peft_kwargs = {"rank": lora_rank, "peft_module": LoRASurgery}
 24
 25    _, sam = get_sam_model(
 26        model_type=model_type,
 27        device=device,
 28        checkpoint_path=checkpoint_path,
 29        return_sam=True,
 30        image_size=image_size,
 31        flexible_load_checkpoint=True,
 32        peft_kwargs=peft_kwargs,
 33    )
 34
 35    # Make sure not to freeze the encoder when using LoRA.
 36    freeze_encoder_ = freeze_encoder if lora_rank is None else False
 37    sam_3d = SimpleSam3DWrapper(sam, num_classes=n_classes, freeze_encoder=freeze_encoder_)
 38    sam_3d.to(device)
 39    return sam_3d
 40
 41
 42class BasicBlock(nn.Module):
 43    def __init__(
 44        self,
 45        in_channels,
 46        out_channels,
 47        kernel_size=(3, 3, 3),
 48        stride=(1, 1, 1),
 49        padding=(1, 1, 1),
 50        bias=True,
 51        mode="nearest"
 52    ):
 53        super().__init__()
 54
 55        self.conv1 = nn.Sequential(
 56            nn.Conv3d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias),
 57            nn.InstanceNorm3d(out_channels),
 58            nn.LeakyReLU()
 59        )
 60
 61        self.conv2 = nn.Sequential(
 62            nn.Conv3d(out_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias),
 63            nn.InstanceNorm3d(out_channels)
 64        )
 65
 66        self.downsample = nn.Sequential(
 67            nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=stride, padding=0, bias=bias),
 68            nn.InstanceNorm3d(out_channels)
 69        )
 70
 71        self.leakyrelu = nn.LeakyReLU()
 72
 73        self.up = nn.Upsample(scale_factor=(1, 2, 2), mode=mode)
 74
 75    def forward(self, x):
 76        residual = self.downsample(x)
 77
 78        out = self.conv1(x)
 79        out = self.conv2(out)
 80        out += residual
 81
 82        out = self.leakyrelu(out)
 83        out = self.up(out)
 84        return out
 85
 86
 87class SegmentationHead(nn.Sequential):
 88    def __init__(
 89        self,
 90        in_channels,
 91        out_channels,
 92        kernel_size=(3, 3, 3),
 93        stride=(1, 1, 1),
 94        padding=(1, 1, 1),
 95        bias=True
 96    ):
 97        super().__init__()
 98
 99        self.conv_pred = nn.Sequential(
100            nn.Conv3d(
101                in_channels, in_channels // 2, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias
102            ),
103            nn.InstanceNorm3d(in_channels // 2),
104            nn.LeakyReLU()
105        )
106        self.segmentation_head = nn.Conv3d(in_channels // 2, out_channels, kernel_size=1)
107
108    def forward(self, x):
109        x = self.conv_pred(x)
110        return self.segmentation_head(x)
111
112
113class SimpleSam3DWrapper(nn.Module):
114    def __init__(self, sam, num_classes, freeze_encoder):
115        super().__init__()
116
117        self.sam = sam
118        self.freeze_encoder = freeze_encoder
119        if self.freeze_encoder:
120            for param in self.sam.image_encoder.parameters():
121                param.requires_grad = False
122            self.no_grad = torch.no_grad
123
124        else:
125            self.no_grad = nullcontext
126
127        self.decoders = nn.ModuleList([
128            BasicBlock(in_channels=256, out_channels=128),
129            BasicBlock(in_channels=128, out_channels=64),
130            BasicBlock(in_channels=64, out_channels=32),
131            BasicBlock(in_channels=32, out_channels=16),
132        ])
133        self.out_conv = SegmentationHead(in_channels=16, out_channels=num_classes)
134
135    def _apply_image_encoder(self, x, D):
136        encoder_features = []
137        for d in range(D):
138            image = x[:, :, d]
139            feature = self.sam.image_encoder(image)
140            encoder_features.append(feature)
141        encoder_features = torch.stack(encoder_features, 2)
142        return encoder_features
143
144    def forward(
145        self,
146        batched_input: List[Dict[str, Any]],
147        multimask_output: bool
148    ) -> List[Dict[str, torch.Tensor]]:
149        """Predict 3D masks for the current inputs.
150
151        Unlike original SAM this model only supports automatic segmentation and does not support prompts.
152
153        Args:
154            batched_input: A list over input images, each a dictionary with the following keys.L
155                'image': The image as a torch tensor in 3xDxHxW format. Already transformed for the input to the model.
156            multimask_output: Wheterh to predict with the multi- or single-mask head of the maks decoder.
157
158        Returns:
159            A list over input images, where each element is as dictionary with the following keys:
160                'masks': Mask prediction for this object.
161        """
162        x = torch.stack([inp["image"] for inp in batched_input], dim=0)
163
164        B, C, D, H, W = x.shape
165        assert C == 3
166
167        with self.no_grad():
168            features = self._apply_image_encoder(x, D)
169
170        out = features
171        for decoder in self.decoders:
172            out = decoder(out)
173        logits = self.out_conv(out)
174
175        outputs = [{"masks": mask.unsqueeze(0)} for mask in logits]
176        return outputs