konfai 1.0.0__py3-none-any.whl

konfai/models/generation/gan.py

@@ -0,0 +1,134 @@
+from functools import partial
+import torch
+
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+from KonfAI.konfai.data.HDF5 import ModelPatch
+
+class Discriminator(network.Network):
+
+    class DiscriminatorNLayers(network.ModuleArgsDict):
+
+        def __init__(self, channels: list[int], strides: list[int], dim: int) -> None:
+            super().__init__()
+            blockConfig = partial(blocks.BlockConfig, kernel_size=4, padding=1, bias=False, activation=partial(torch.nn.LeakyReLU, negative_slope = 0.2, inplace=True), normMode=blocks.NormMode.SYNCBATCH)
+            for i, (in_channels, out_channels, stride) in enumerate(zip(channels, channels[1:], strides)):
+                self.add_module("Layer_{}".format(i), blocks.ConvBlock(in_channels, out_channels, 1, blockConfig(stride=stride), dim))
+    
+    class DiscriminatorHead(network.ModuleArgsDict):
+
+        def __init__(self, channels: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels=channels, out_channels=1, kernel_size=4, stride=1, padding=1))
+            self.add_module("AdaptiveAvgPool", blocks.getTorchModule("AdaptiveAvgPool", dim)(tuple([1]*dim)))
+            self.add_module("Flatten", torch.nn.Flatten(1))
+    
+    @config("Discriminator")
+    def __init__(self,
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    nb_batch_per_step: int = 64,
+                    dim : int = 3) -> None:
+        super().__init__(in_channels = 1, optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim=dim, nb_batch_per_step=nb_batch_per_step)
+        channels = [1, 16, 32, 64, 64]
+        strides = [2,2,2,1]
+        self.add_module("Layers", Discriminator.DiscriminatorNLayers(channels, strides, dim))
+        self.add_module("Head", Discriminator.DiscriminatorHead(channels[-1], dim))
+
+class Generator(network.Network):
+
+    class GeneratorStem(network.ModuleArgsDict):
+
+        def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("ConvBlock", blocks.ConvBlock(in_channels, out_channels, nb_conv=1, blockConfig=blocks.BlockConfig(bias=False, activation="ReLU", normMode="SYNCBATCH"), dim=dim))
+
+    class GeneratorHead(network.ModuleArgsDict):
+
+        def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("ConvBlock", blocks.ConvBlock(in_channels, in_channels, nb_conv=1, blockConfig=blocks.BlockConfig(bias=False, activation="ReLU", normMode="SYNCBATCH"), dim=dim))
+            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels, out_channels, kernel_size=1, bias=False))
+            self.add_module("Tanh", torch.nn.Tanh())
+
+    class GeneratorDownSample(network.ModuleArgsDict):
+
+        def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("ConvBlock", blocks.ConvBlock(in_channels, out_channels, nb_conv=1, blockConfig=blocks.BlockConfig(stride=2, bias=False, activation="ReLU", normMode="SYNCBATCH"), dim=dim))
+    
+    class GeneratorUpSample(network.ModuleArgsDict):
+
+        def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("ConvBlock", blocks.ConvBlock(in_channels, out_channels, nb_conv=1, blockConfig=blocks.BlockConfig(bias=False, activation="ReLU", normMode="SYNCBATCH"), dim=dim))
+            self.add_module("Upsample", torch.nn.Upsample(scale_factor=2, mode="bilinear" if dim < 3 else "trilinear"))
+    
+    class GeneratorEncoder(network.ModuleArgsDict):
+        def __init__(self, channels: list[int], dim: int) -> None:
+            super().__init__()
+            for i, (in_channels, out_channels) in enumerate(zip(channels, channels[1:])):
+                self.add_module("DownSample_{}".format(i), Generator.GeneratorDownSample(in_channels=in_channels, out_channels=out_channels, dim=dim))
+    
+    class GeneratorResnetBlock(network.ModuleArgsDict):
+
+        def __init__(self, channels : int, dim : int):
+            super().__init__()
+            self.add_module("Conv_0", blocks.getTorchModule("Conv", dim)(channels, channels, kernel_size=3, padding=1, bias=False))
+            self.add_module("Norm", torch.nn.LeakyReLU(0.2, inplace=True))
+            self.add_module("Conv_1", blocks.getTorchModule("Conv", dim)(channels, channels, kernel_size=3, padding=1, bias=False))
+            self.add_module("Residual", blocks.Add(), in_branch=[0,1])
+
+    class GeneratorNResnetBlock(network.ModuleArgsDict):
+
+        def __init__(self, channels: int, nb_conv: int, dim: int) -> None:
+            super().__init__()
+            for i in range(nb_conv):
+                self.add_module("ResnetBlock_{}".format(i), Generator.GeneratorResnetBlock(channels=channels, dim=dim))
+
+    class GeneratorDecoder(network.ModuleArgsDict):
+        def __init__(self, channels: list[int], dim: int) -> None:
+            super().__init__()
+            for i, (in_channels, out_channels) in enumerate(zip(reversed(channels), reversed(channels[:-1]))):
+                self.add_module("UpSample_{}".format(i), Generator.GeneratorUpSample(in_channels=in_channels, out_channels=out_channels, dim=dim))
+    
+    class GeneratorAutoEncoder(network.ModuleArgsDict):
+
+        def __init__(self, ngf: int, dim: int) -> None:
+            super().__init__()
+            channels = [ngf, ngf*2]
+            self.add_module("Encoder", Generator.GeneratorEncoder(channels, dim))
+            self.add_module("NResBlock", Generator.GeneratorNResnetBlock(channels=channels[-1], nb_conv=6, dim=dim))
+            self.add_module("Decoder", Generator.GeneratorDecoder(channels, dim))
+            
+    @config("Generator")
+    def __init__(self, 
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    patch : ModelPatch = ModelPatch(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    nb_batch_per_step: int = 64,
+                    dim : int = 3) -> None:
+        super().__init__(optimizer=optimizer, in_channels=1, schedulers=schedulers, patch=patch, outputsCriterions=outputsCriterions, dim=dim, nb_batch_per_step=nb_batch_per_step)
+        ngf=32
+        self.add_module("Stem", Generator.GeneratorStem(1, ngf, dim))
+        self.add_module("AutoEncoder", Generator.GeneratorAutoEncoder(ngf, dim))
+        self.add_module("Head", Generator.GeneratorHead(in_channels=ngf, out_channels=1, dim=dim))
+
+    def getName(self):
+        return "Generator"
+
+class Gan(network.Network):
+
+    @config("Gan")
+    def __init__(self, generator : Generator = Generator(), discriminator : Discriminator = Discriminator()) -> None:
+        super().__init__()
+        self.add_module("Discriminator_B", discriminator, in_branch=[1], out_branch=[-1], requires_grad=True)
+        self.add_module("Generator_A_to_B", generator, in_branch=[0], out_branch=["pB"])
+        
+        self.add_module("detach", blocks.Detach(), in_branch=["pB"], out_branch=["pB_detach"])
+        self.add_module("Discriminator_pB_detach", discriminator, in_branch=["pB_detach"], out_branch=[-1])
+          
+        self.add_module("Discriminator_pB", discriminator, in_branch=["pB"], out_branch=[-1], requires_grad=False)
+

konfai/models/generation/vae.py

@@ -0,0 +1,72 @@
+import torch
+
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+
+class VAE(network.Network):
+
+    class AutoEncoderBlock(network.ModuleArgsDict):
+
+        def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, dim: int, block: type, i : int = 0) -> None:
+            super().__init__()
+            if i > 0:
+                self.add_module(downSampleMode.name, blocks.downSample(in_channels=channels[0], out_channels=channels[1], downSampleMode=downSampleMode, dim=dim))
+            self.add_module("DownBlock", block(in_channels=channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0], out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim))
+            if len(channels) > 2:
+                self.add_module("AutoEncoder_{}".format(i+1), VAE.AutoEncoderBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, dim, block, i+1))
+                self.add_module("UpBlock", block(in_channels=channels[2] if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1], out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim))
+            if i > 0:
+               self.add_module(upSampleMode.name, blocks.upSample(in_channels=channels[1], out_channels=channels[0], upSampleMode=upSampleMode, dim=dim))
+    
+    class VAE_Head(network.ModuleArgsDict):
+
+        def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = out_channels, kernel_size = 3, stride = 1, padding = 1))
+            self.add_module("Tanh", torch.nn.Tanh()) 
+
+    @config("VAE")
+    def __init__(self,
+                    optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers: network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    dim : int = 3,
+                    channels: list[int]=[1, 64, 128, 256, 512, 1024],
+                    blockConfig: blocks.BlockConfig = blocks.BlockConfig(),
+                    nb_conv_per_stage: int = 2,
+                    downSampleMode: str = "MAXPOOL",
+                    upSampleMode: str = "CONV_TRANSPOSE",
+                    blockType: str = "Conv") -> None:
+        
+        super().__init__(in_channels = channels[0], init_type="normal", optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim=dim, nb_batch_per_step=1)
+        self.add_module("AutoEncoder_0", VAE.AutoEncoderBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], dim=dim, block = blocks.ConvBlock if blockType == "Conv" else blocks.ResBlock))
+        self.add_module("Head", VAE.VAE_Head(channels[1], channels[0], dim))
+
+
+class LinearVAE(network.Network):
+
+    class LinearVAE_DenseLayer(network.ModuleArgsDict):
+
+        def __init__(self, in_features: int, out_features: int) -> None:
+            super().__init__()
+            self.add_module("Linear", torch.nn.Linear(in_features, out_features))
+            #self.add_module("Norm", torch.nn.BatchNorm1d(out_features))
+            self.add_module("Activation", torch.nn.LeakyReLU())
+
+    class LinearVAE_Head(network.ModuleArgsDict):
+
+        def __init__(self, in_features: int, out_features: int) -> None:
+            super().__init__()
+            self.add_module("Linear", torch.nn.Linear(in_features, out_features))
+            self.add_module("Tanh", torch.nn.Tanh())
+
+    @config("LinearVAE")
+    def __init__(self,
+                    optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers: network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},) -> None:
+        super().__init__(in_channels = 1, init_type="normal", optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim=1, nb_batch_per_step=1)
+        self.add_module("DenseLayer_0", LinearVAE.LinearVAE_DenseLayer(23343, 5))
+        #self.add_module("Head", LinearVAE.DenseLayer(100, 28590))
+        self.add_module("Head", LinearVAE.LinearVAE_Head(5, 23343))
+        #self.add_module("DenseLayer_5", LinearVAE.DenseLayer(5000, 28590))

konfai/models/registration/registration.py

@@ -0,0 +1,136 @@
+import torch
+from torch.nn.parameter import Parameter
+import torch.nn.functional as F
+
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+from KonfAI.konfai.models.segmentation import UNet
+
+class VoxelMorph(network.Network):
+                
+    @config("VoxelMorph")
+    def __init__(   self,
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    dim : int = 3,
+                    channels : list[int] = [4, 16,32,32,32],
+                    blockConfig: blocks.BlockConfig = blocks.BlockConfig(),
+                    nb_conv_per_stage: int = 2,
+                    downSampleMode: str = "MAXPOOL",
+                    upSampleMode: str = "CONV_TRANSPOSE",
+                    attention : bool = False,
+                    shape : list[int] = [192, 192, 192],
+                    int_steps : int = 7,
+                    int_downsize : int = 2,
+                    nb_batch_per_step : int = 1,
+                    rigid: bool = False):
+        super().__init__(in_channels = channels[0], optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim = dim, nb_batch_per_step=nb_batch_per_step)
+        self.add_module("Concat", blocks.Concat(), in_branch=[0,1], out_branch=["input_concat"])
+        self.add_module("UNetBlock_0", UNet.UNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, block = blocks.ConvBlock, nb_class=0, dim=dim), in_branch=["input_concat"], out_branch=["unet"])
+
+        if rigid:
+            self.add_module("Flow", Rigid(channels[1], dim), in_branch=["unet"], out_branch=["pos_flow"])
+        else:
+            self.add_module("Flow", Flow(channels[1], int_steps, int_downsize, shape, dim), in_branch=["unet"], out_branch=["pos_flow"])
+        self.add_module("MovingImageResample", SpatialTransformer(shape, rigid=rigid), in_branch=[1, "pos_flow"], out_branch=["moving_image_resample"])
+
+class Flow(network.ModuleArgsDict):
+
+    def __init__(self, in_channels: int, int_steps: int, int_downsize: int, shape: list[int], dim: int) -> None:
+        super().__init__()
+        self.add_module("Head", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = dim, kernel_size = 3, stride = 1, padding = 1))
+        self["Head"].weight = Parameter(torch.distributions.Normal(0, 1e-5).sample(self["Head"].weight.shape))
+        self["Head"].bias = Parameter(torch.zeros(self["Head"].bias.shape))
+    
+        if int_steps > 0 and int_downsize > 1:
+            self.add_module("DownSample", ResizeTransform(int_downsize))
+
+        if int_steps > 0:
+            self.add_module("Integrate_pos_flow", VecInt([int(dim / int_downsize) for dim in shape], int_steps))
+        
+        if int_steps > 0 and int_downsize > 1:
+            self.add_module("Upsample_pos_flow", ResizeTransform(1 / int_downsize))
+
+class Rigid(network.ModuleArgsDict):
+
+    def __init__(self, in_channels: int, dim: int) -> None:
+        super().__init__()
+        self.add_module("ToFeatures", torch.nn.Flatten(1))
+        self.add_module("Head", torch.nn.Linear(in_channels*512*512, 2))
+
+    def init(self, init_type: str, init_gain: float):
+        self["Head"].weight.data.fill_(0)
+        self["Head"].bias.data.copy_(torch.tensor([0, 0], dtype=torch.float))
+        
+class MaskFlow(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+    
+    def forward(self, mask: torch.Tensor, *flows: torch.Tensor):
+        result = torch.zeros_like(flows[0])
+        for i, flow in enumerate(flows):
+            result = result+torch.where(mask == i+1, flow, torch.tensor(0))
+        return result
+
+class SpatialTransformer(torch.nn.Module):
+    
+    def __init__(self, size : list[int], rigid: bool = False):
+        super().__init__()
+        self.rigid = rigid
+        if not rigid:
+            vectors = [torch.arange(0, s) for s in size]
+            grids = torch.meshgrid(vectors, indexing='ij')
+            grid = torch.stack(grids)
+            grid = torch.unsqueeze(grid, 0)
+            grid = grid.type(torch.float)
+            self.register_buffer('grid', grid)
+
+    def forward(self, src: torch.Tensor, flow: torch.Tensor):
+        if self.rigid:
+            new_locs = torch.zeros((flow.shape[0], 2, 3)).to(flow.device)
+            new_locs[:, 0,0] = 1
+            new_locs[:, 1,1] = 1
+            new_locs[:, 0,2] = flow[:, 0]
+            new_locs[:, 1,2] = flow[:, 1]
+            print(new_locs)
+            return F.grid_sample(src, F.affine_grid(new_locs, src.size()), align_corners=True, mode="bilinear")
+        else:
+            new_locs = self.grid + flow
+            shape = flow.shape[2:]
+            for i in range(len(shape)):
+                new_locs[:, i, ...] = 2 * (new_locs[:, i, ...] / (shape[i] - 1) - 0.5)
+            new_locs = new_locs.permute(0, 2, 3, 1)
+            return F.grid_sample(src, new_locs[..., [1, 0]], align_corners=True, mode="bilinear")
+
+class VecInt(torch.nn.Module):
+
+    def __init__(self, inshape, nsteps):
+        super().__init__()
+        assert nsteps >= 0, 'nsteps should be >= 0, found: %d' % nsteps
+        self.nsteps = nsteps
+        self.scale = 1.0 / (2 ** self.nsteps)
+        self.transformer = SpatialTransformer(inshape)
+
+    def forward(self, vec: torch.Tensor):
+        vec = vec * self.scale
+        for _ in range(self.nsteps):
+            vec = vec + self.transformer(vec, vec)
+        return vec
+
+
+class ResizeTransform(torch.nn.Module):
+   
+    def __init__(self, size):
+        super().__init__()
+        self.factor = 1.0 / size
+
+    def forward(self, x: torch.Tensor):
+        if self.factor < 1:
+            x = F.interpolate(x, align_corners=True, scale_factor=self.factor, mode="bilinear", recompute_scale_factor = True)
+            x = self.factor * x
+        elif self.factor > 1:
+            x = self.factor * x
+            x = F.interpolate(x, align_corners=True, scale_factor=self.factor, mode="bilinear", recompute_scale_factor = True)
+        return x

konfai/models/representation/representation.py

@@ -0,0 +1,57 @@
+import torch
+
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+
+class ConvBlock(torch.nn.Module):
+    
+    def __init__(self, in_channels : int, out_channels : int) -> None:
+        super().__init__()
+        self.Conv_0 = torch.nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=True)
+        self.Norm_0 = torch.nn.InstanceNorm3d(num_features=out_channels, affine=True)
+        self.Activation_0 = torch.nn.LeakyReLU(negative_slope=0.01)
+        self.Conv_1 = torch.nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=True)
+        self.Norm_1 = torch.nn.InstanceNorm3d(num_features=out_channels, affine=True)
+        self.Activation_1 = torch.nn.LeakyReLU(negative_slope=0.01)
+        
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        output = self.Conv_0(input)
+        output = self.Norm_0(output)
+        output = self.Activation_0(output)
+        output = self.Conv_1(output)
+        output = self.Norm_1(output)
+        output = self.Activation_1(output)
+        return output
+
+class UnetCPP_1_Layers(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.DownConvBlock_0 = ConvBlock(in_channels=1, out_channels=32)
+        
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return self.DownConvBlock_0(input)
+
+class Adaptation(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.Encoder_1 = UnetCPP_1_Layers()
+        self.ToFeatures = blocks.ToFeatures(3)
+        self.FCT_1 = torch.nn.Linear(32, 32, bias=True)
+
+    def forward(self, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> list[torch.Tensor]:
+        self.Encoder_1.requires_grad_(False)
+        self.FCT_1.requires_grad_(True)
+        return self.FCT_1(self.ToFeatures(self.Encoder_1(A))), self.FCT_1(self.ToFeatures(self.Encoder_1(B))), self.FCT_1(self.ToFeatures(self.Encoder_1(C)))
+
+class Representation(network.Network):
+
+    @config("Representation")
+    def __init__(   self,
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    dim : int = 3):
+        super().__init__(in_channels = 1, optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim=dim, init_type="kaiming")
+        self.add_module("Model", Adaptation(), in_branch=[0,1,2])

konfai/models/segmentation/NestedUNet.py

@@ -0,0 +1,53 @@
+from typing import Union
+import torch
+
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+from KonfAI.konfai.data.HDF5 import ModelPatch
+
+     
+class NestedUNetBlock(network.ModuleArgsDict):
+
+    def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, attention : bool, block: type, dim: int, i : int = 0) -> None:
+        super().__init__()
+        if i > 0:
+            self.add_module(downSampleMode.name, blocks.downSample(in_channels=channels[0], out_channels=channels[1], downSampleMode=downSampleMode, dim=dim))
+        
+        self.add_module("X_{}_{}".format(i, 0), block(in_channels=channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0], out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim), out_branch=["X_{}_{}".format(i, 0)])
+        if len(channels) > 2:
+            self.add_module("UNetBlock_{}".format(i+1), NestedUNetBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, block, dim, i+1), in_branch=["X_{}_{}".format(i, 0)], out_branch=["X_{}_{}".format(i+1, j) for j in range(len(channels)-2)])
+            for j in range(len(channels)-2):
+                self.add_module("X_{}_{}_{}".format(i, j+1, upSampleMode.name), blocks.upSample(in_channels=channels[2], out_channels=channels[1], upSampleMode=upSampleMode, dim=dim), in_branch=["X_{}_{}".format(i+1, j)], out_branch=["X_{}_{}".format(i+1, j)])
+                self.add_module("SkipConnection_{}_{}".format(i, j+1), blocks.Concat(), in_branch=["X_{}_{}".format(i+1, j)]+["X_{}_{}".format(i, r) for r in range(j+1)], out_branch=["X_{}_{}".format(i, j+1)])
+                self.add_module("X_{}_{}".format(i, j+1), block(in_channels=(channels[1]*(j+1)+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*(j+2), out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim), in_branch=["X_{}_{}".format(i, j+1)], out_branch=["X_{}_{}".format(i, j+1)])
+
+class NestedUNet(network.Network):
+
+    class NestedUNetHead(network.ModuleArgsDict):
+
+        def __init__(self, in_channels: int, nb_class: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = nb_class, kernel_size = 1, stride = 1, padding = 0))
+            self.add_module("Softmax", torch.nn.Softmax(dim=1))
+            self.add_module("Argmax", blocks.ArgMax(dim=1))
+
+    @config("NestedUNet")
+    def __init__(   self,
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    patch : Union[ModelPatch, None] = None,
+                    dim : int = 3,
+                    channels: list[int]=[1, 64, 128, 256, 512, 1024],
+                    nb_class: int = 2,
+                    blockConfig: blocks.BlockConfig = blocks.BlockConfig(),
+                    nb_conv_per_stage: int = 2,
+                    downSampleMode: str = "MAXPOOL",
+                    upSampleMode: str = "CONV_TRANSPOSE",
+                    attention : bool = False,
+                    blockType: str = "Conv") -> None:
+        super().__init__(in_channels = channels[0], optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, patch=patch, dim = dim)
+                
+        self.add_module("UNetBlock_0", NestedUNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, block = blocks.ConvBlock if blockType == "Conv" else blocks.ResBlock, dim=dim), out_branch=["X_0_{}".format(j+1) for j in range(len(channels)-2)])    
+        for j in range(len(channels)-2):
+            self.add_module("Head_{}".format(j), NestedUNet.NestedUNetHead(in_channels=channels[1], nb_class=nb_class, dim=dim), in_branch=["X_0_{}".format(j+1)], out_branch=[-1])

konfai/models/segmentation/UNet.py

@@ -0,0 +1,58 @@
+import torch
+from typing import Union
+
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+from KonfAI.konfai.data.HDF5 import ModelPatch
+
+class UNetHead(network.ModuleArgsDict):
+
+    def __init__(self, in_channels: int, nb_class: int, dim: int) -> None:
+        super().__init__()
+        self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = nb_class, kernel_size = 1, stride = 1, padding = 0))
+        self.add_module("Softmax", torch.nn.Softmax(dim=1))
+        self.add_module("Argmax", blocks.ArgMax(dim=1))
+
+class UNetBlock(network.ModuleArgsDict):
+
+    def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, attention : bool, block: type, nb_class: int, dim: int, i : int = 0, mri: bool = False) -> None:
+        super().__init__()
+        blockConfig_stride = blockConfig
+        if i > 0:
+            if downSampleMode != blocks.DownSampleMode.CONV_STRIDE:
+                self.add_module(downSampleMode.name, blocks.downSample(in_channels=channels[0], out_channels=channels[1], downSampleMode=downSampleMode, dim=dim))
+            else:
+                blockConfig_stride = blocks.BlockConfig(blockConfig.kernel_size, (1,2,2) if mri and i > 4 else 2, blockConfig.padding, blockConfig.bias, blockConfig.activation, blockConfig.normMode)
+        self.add_module("DownConvBlock", block(in_channels=channels[0], out_channels=channels[1], blockConfigs=[blockConfig_stride]+[blockConfig]*(nb_conv_per_stage-1), dim=dim))
+        if len(channels) > 2:
+            self.add_module("UNetBlock_{}".format(i+1), UNetBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, block, nb_class, dim, i+1, mri=mri))
+            self.add_module("UpConvBlock", block(in_channels=(channels[1]+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*2, out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim))
+            if nb_class > 0:
+                self.add_module("Head", UNetHead(channels[1], nb_class, dim), out_branch=[-1])
+        if i > 0:
+            if attention:
+                self.add_module("Attention", blocks.Attention(F_g=channels[1], F_l=channels[0], F_int=channels[0], dim=dim), in_branch=[1, 0], out_branch=[1])
+            self.add_module(upSampleMode.name, blocks.upSample(in_channels=channels[1], out_channels=channels[0], upSampleMode=upSampleMode, dim=dim, kernel_size=(1,2,2) if mri and i > 4 else 2, stride=(1,2,2) if mri and i > 4 else 2))
+            self.add_module("SkipConnection", blocks.Concat(), in_branch=[0, 1])
+
+class UNet(network.Network):
+
+    @config("UNet")
+    def __init__(   self,
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    patch : Union[ModelPatch, None] = None,
+                    dim : int = 3,
+                    channels: list[int]=[1, 64, 128, 256, 512, 1024],
+                    nb_class: int = 2,
+                    blockConfig: blocks.BlockConfig = blocks.BlockConfig(),
+                    nb_conv_per_stage: int = 2,
+                    downSampleMode: str = "MAXPOOL",
+                    upSampleMode: str = "CONV_TRANSPOSE",
+                    attention : bool = False,
+                    blockType: str = "Conv",
+                    mri: bool = False) -> None:
+        super().__init__(in_channels = channels[0], optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, patch=patch, dim = dim)        
+        self.add_module("UNetBlock_0", UNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, block = blocks.ConvBlock if blockType == "Conv" else blocks.ResBlock, nb_class=nb_class, dim=dim, mri = mri))    
+