micro_sam.models.simple_sam_3d_wrapper

  1import os
  2from contextlib import nullcontext
  3from typing import Any, List, Dict, Union, Optional
  4
  5import torch
  6import torch.nn as nn
  7
  8from ..util import get_sam_model
  9from .peft_sam import LoRASurgery
 10
 11
 12def get_simple_sam_3d_model(
 13    device: Union[str, torch.device],
 14    n_classes: int,
 15    image_size: int,
 16    lora_rank: Optional[int] = None,
 17    freeze_encoder: bool = False,
 18    model_type: str = "vit_b",
 19    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
 20):
 21    if lora_rank is None:
 22        peft_kwargs = {}
 23    else:
 24        peft_kwargs = {"rank": lora_rank, "peft_module": LoRASurgery}
 25
 26    _, sam = get_sam_model(
 27        model_type=model_type,
 28        device=device,
 29        checkpoint_path=checkpoint_path,
 30        return_sam=True,
 31        image_size=image_size,
 32        flexible_load_checkpoint=True,
 33        peft_kwargs=peft_kwargs,
 34    )
 35
 36    # Make sure not to freeze the encoder when using LoRA.
 37    freeze_encoder_ = freeze_encoder if lora_rank is None else False
 38    sam_3d = SimpleSam3DWrapper(sam, num_classes=n_classes, freeze_encoder=freeze_encoder_)
 39    sam_3d.to(device)
 40
 41    return sam_3d
 42
 43
 44class BasicBlock(nn.Module):
 45    def __init__(
 46        self,
 47        in_channels,
 48        out_channels,
 49        kernel_size=(3, 3, 3),
 50        stride=(1, 1, 1),
 51        padding=(1, 1, 1),
 52        bias=True,
 53        mode="nearest"
 54    ):
 55        super().__init__()
 56
 57        self.conv1 = nn.Sequential(
 58            nn.Conv3d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias),
 59            nn.InstanceNorm3d(out_channels),
 60            nn.LeakyReLU()
 61        )
 62
 63        self.conv2 = nn.Sequential(
 64            nn.Conv3d(out_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias),
 65            nn.InstanceNorm3d(out_channels)
 66        )
 67
 68        self.downsample = nn.Sequential(
 69            nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=stride, padding=0, bias=bias),
 70            nn.InstanceNorm3d(out_channels)
 71        )
 72
 73        self.leakyrelu = nn.LeakyReLU()
 74
 75        self.up = nn.Upsample(scale_factor=(1, 2, 2), mode=mode)
 76
 77    def forward(self, x):
 78        residual = self.downsample(x)
 79
 80        out = self.conv1(x)
 81        out = self.conv2(out)
 82        out += residual
 83
 84        out = self.leakyrelu(out)
 85        out = self.up(out)
 86        return out
 87
 88
 89class SegmentationHead(nn.Sequential):
 90    def __init__(
 91        self,
 92        in_channels,
 93        out_channels,
 94        kernel_size=(3, 3, 3),
 95        stride=(1, 1, 1),
 96        padding=(1, 1, 1),
 97        bias=True
 98    ):
 99        super().__init__()
100
101        self.conv_pred = nn.Sequential(
102            nn.Conv3d(
103                in_channels, in_channels // 2, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias
104            ),
105            nn.InstanceNorm3d(in_channels // 2),
106            nn.LeakyReLU()
107        )
108        self.segmentation_head = nn.Conv3d(in_channels // 2, out_channels, kernel_size=1)
109
110    def forward(self, x):
111        x = self.conv_pred(x)
112        return self.segmentation_head(x)
113
114
115class SimpleSam3DWrapper(nn.Module):
116    def __init__(self, sam, num_classes, freeze_encoder):
117        super().__init__()
118
119        self.sam = sam
120        self.freeze_encoder = freeze_encoder
121        if self.freeze_encoder:
122            for param in self.sam.image_encoder.parameters():
123                param.requires_grad = False
124            self.no_grad = torch.no_grad
125
126        else:
127            self.no_grad = nullcontext
128
129        self.decoders = nn.ModuleList([
130            BasicBlock(in_channels=256, out_channels=128),
131            BasicBlock(in_channels=128, out_channels=64),
132            BasicBlock(in_channels=64, out_channels=32),
133            BasicBlock(in_channels=32, out_channels=16),
134        ])
135        self.out_conv = SegmentationHead(in_channels=16, out_channels=num_classes)
136
137    def _apply_image_encoder(self, x, D):
138        encoder_features = []
139        for d in range(D):
140            image = x[:, :, d]
141            feature = self.sam.image_encoder(image)
142            encoder_features.append(feature)
143        encoder_features = torch.stack(encoder_features, 2)
144        return encoder_features
145
146    def forward(
147        self, batched_input: List[Dict[str, Any]], 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