Added preptraining model example

nnunet/loader.py

@@ -3,6 +3,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from collections import OrderedDict
 from contextlib import contextmanager
 
 os.environ["NNUNET_LOADER_USED"] = "True"
@@ -95,5 +96,92 @@ def export_onnx(
             )
 
 
+class FeatureHook(nn.Module):
+    def __init__(self, module: nn.Module, layers: list[str]):
+        super().__init__()
+        self.module = module
+        self.layers = layers
+        self._features = {layer: torch.empty(0) for layer in layers}
+
+        for layer_id in layers:
+            layer = self.module._modules[layer_id]
+            layer.register_forward_hook(self.save_outputs_hook(layer_id))
+
+    def save_outputs_hook(self, layer_id: str):
+        def fn(module, input, output):
+            self._features[layer_id] = output
+        return fn
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        x = self.module(x)
+        return x, self._features
+
+
+class UNet(nn.Module):
+    def __init__(
+        self,
+        num_channels: int = 1,
+        image_size: int = 512,
+        pretrained: str = "/ws/fold_0",
+        deep_supervision: bool = False,
+    ):
+        super().__init__()
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.deep_supervision = deep_supervision
+
+        fe_layers = nn.ModuleDict()
+
+        nnunet = nnUNet(pretrained)
+        for i, block in enumerate(nnunet.model.conv_blocks_context[:-1]):
+            fe_layers[f"stack_{i}"] = block
+            fe_layers[f"stack_{i}"].out_features = block.output_channels
+
+        self.bottleneck = nnunet.model.conv_blocks_context[-1]
+
+        self.fe_layers = nn.Sequential(OrderedDict(**fe_layers))
+        self.fe_keys = list(fe_layers.keys())
+        self.fe_hook = FeatureHook(self.fe_layers, self.fe_keys)
+
+        # Localization and Upscale layers
+        self.upscale_layers = nn.ModuleDict()
+        self.loc_layers = nn.ModuleDict()
+        self.seg_layers = nn.ModuleDict()
+        for i, key in enumerate(self.fe_keys[::-1]):
+            self.upscale_layers[key] = nnunet.model.tu[i]
+            self.loc_layers[key] = nnunet.model.conv_blocks_localization[i]
+            self.seg_layers[key] = nnunet.model.seg_outputs[i]
+
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        segmentation_outputs = []
+        b, c, h, w = x.size()
+        assert c == self.num_channels
+        assert h == w == self.image_size
+
+        x, downscale_features = self.fe_hook(x)
+
+        x = self.bottleneck(x)
+
+        for key in self.fe_keys[::-1]:
+            x = torch.cat((self.upscale_layers[key](x), downscale_features[key]), dim=1)
+            x = self.loc_layers[key](x)
+            segmentation = self.seg_layers[key](x)
+            segmentation_outputs.append(segmentation)
+
+        if self.deep_supervision:
+            return segmentation_outputs
+        else:
+            return segmentation_outputs[-1]
+
+
 if __name__ == "__main__":
-    nnunet_pretrained = nnUNet("/path/to/nnUNet_trained_models/nnUNet/2d/TaskXXX_NAME/nnUNetTrainerV2__nnUNetPlansv2.1/fold_0")
+
+    dev = torch.device("cpu")
+    model = UNet(pretrained="...Task512_MtSinaiBinBkg/nnUNetTrainerV2__nnUNetPlansv2.1/fold_0")
+    input = torch.rand(4, 1, 512, 512).to(dev)
+    model = model.to(dev)
+    out = model(input)
+    print(out.shape)
+# "/Users/artem/pyproj/ws/tcbs/backbones/nnunet/Task510_BBOXSem/nnUNetTrainerV2__nnUNetPlansv2.1/fold_0"
+