konfai 1.1.3__tar.gz → 1.1.5__tar.gz

{konfai-1.1.3 → konfai-1.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.3
+Version: 1.1.5
 Summary: Modular and configurable Deep Learning framework with YAML and PyTorch
 Author-email: Valentin Boussot <boussot.v@gmail.com>
 License-Expression: Apache-2.0

{konfai-1.1.3 → konfai-1.1.5}/konfai/data/augmentation.py

@@ -134,7 +134,6 @@ class EulerTransform(DataAugmentation):
                                                                                    
 class Translate(EulerTransform):
     
-    @config("Translate")
     def __init__(self, t_min: float = -10, t_max = 10, is_int: bool = False):
         super().__init__()
         self.t_min = t_min
@@ -152,7 +151,6 @@ class Translate(EulerTransform):
 
 class Rotate(EulerTransform):
 
-    @config("Rotate")
     def __init__(self, a_min: float = 0, a_max: float = 360, is_quarter: bool = False):
         super().__init__()
         self.a_min = a_min
@@ -174,7 +172,6 @@ class Rotate(EulerTransform):
 
 class Scale(EulerTransform):
 
-    @config("Scale")
     def __init__(self, s_std: float = 0.2):
         super().__init__()
         self.s_std = s_std
@@ -187,7 +184,6 @@ class Scale(EulerTransform):
 
 class Flip(DataAugmentation):
 
-    @config("Flip")
     def __init__(self, f_prob: Union[list[float], None] = [0.33, 0.33 ,0.33]) -> None:
         super().__init__()
         self.f_prob = f_prob
@@ -211,7 +207,6 @@ class Flip(DataAugmentation):
 
 class ColorTransform(DataAugmentation):
 
-    @config("ColorTransform")
     def __init__(self, groups: Union[list[str], None] = None) -> None:
         super().__init__(groups)
         self.matrix: dict[int, list[torch.Tensor]] = {}
@@ -236,7 +231,6 @@ class ColorTransform(DataAugmentation):
 
 class Brightness(ColorTransform):
 
-    @config("Brightness")
     def __init__(self, b_std: float, groups: Union[list[str], None] = None) -> None:
         super().__init__(groups)
         self.b_std = b_std
@@ -248,7 +242,6 @@ class Brightness(ColorTransform):
 
 class Contrast(ColorTransform):
 
-    @config("Contrast")
     def __init__(self, c_std: float, groups: Union[list[str], None] = None) -> None:
         super().__init__(groups)
         self.c_std = c_std
@@ -260,7 +253,6 @@ class Contrast(ColorTransform):
 
 class LumaFlip(ColorTransform):
 
-    @config("LumaFlip")
     def __init__(self, groups: Union[list[str], None] = None) -> None:
         super().__init__(groups)
         self.v = torch.tensor([1, 1, 1, 0])/torch.sqrt(torch.tensor(3)) 
@@ -285,7 +277,6 @@ class HUE(ColorTransform):
     
 class Saturation(ColorTransform):
 
-    @config("Saturation")
     def __init__(self, s_std: float, groups: Union[list[str], None] = None) -> None:
         super().__init__(groups)
         self.s_std = s_std
@@ -372,7 +363,6 @@ class Saturation(ColorTransform):
 
 class Noise(DataAugmentation):
 
-    @config("Noise")
     def __init__(self, n_std: float, noise_step: int=1000, beta_start: float = 1e-4, beta_end: float = 0.02, groups: Union[list[str], None] = None) -> None:
         super().__init__(groups)
         self.n_std = n_std
@@ -427,7 +417,6 @@ class Noise(DataAugmentation):
 
 class CutOUT(DataAugmentation):
 
-    @config("CutOUT")
     def __init__(self, c_prob: float, cutout_size: int, value: float, groups: Union[list[str], None] = None) -> None:
         super().__init__(groups)
         self.c_prob = c_prob
@@ -458,7 +447,6 @@ class CutOUT(DataAugmentation):
 
 class Elastix(DataAugmentation):
 
-    @config("Elastix")
     def __init__(self, grid_spacing: int = 16, max_displacement: int = 16) -> None:
         super().__init__()
         self.grid_spacing = grid_spacing
@@ -529,7 +517,6 @@ class Elastix(DataAugmentation):
 
 class Permute(DataAugmentation):
 
-    @config("Permute")
     def __init__(self, prob_permute: Union[list[float], None] = [0.5 ,0.5]) -> None:
         super().__init__()
         self._permute_dims = torch.tensor([[0, 2, 1, 3], [0, 3, 1, 2]])
@@ -567,7 +554,6 @@ class Permute(DataAugmentation):
 
 class Mask(DataAugmentation):
 
-    @config("Mask")
     def __init__(self, mask: str, value: float, groups: Union[list[str], None] = None) -> None:
         super().__init__(groups)
         if mask is not None:

{konfai-1.1.3 → konfai-1.1.5}/konfai/data/data_manager.py

@@ -129,7 +129,7 @@ class DatasetIter(data.Dataset):
                     for group_dest in self.groups_src[group_src]:
                         self.data[group_dest][index].unloadAugmentation()
                         self.data[group_dest][index].resetAugmentation()
-        self.load(label + " Augmentation")
+            self.load(label + " Augmentation")
 
     def load(self, label: str):
         if self.use_cache:
@@ -217,7 +217,6 @@ class Subset():
         else:
             names_filtred = names
         size = len(names_filtred)
-        
         index = []
         if self.subset is None:
             index = list(range(0, size))
@@ -226,15 +225,24 @@ class Subset():
                 r = np.clip(np.asarray([int(self.subset.split(":")[0]), int(self.subset.split(":")[1])]), 0, size)
                 index = list(range(r[0], r[1]))
             elif os.path.exists(self.subset):
-                validation_names = []
+                train_names = []
                 with open(self.subset, "r") as f:
                     for name in f:
-                        validation_names.append(name.strip())
+                        train_names.append(name.strip())
                 index = []
                 for i, name in enumerate(names_filtred):
-                    if name in validation_names:
+                    if name in train_names:
                         index.append(i)
-            
+            elif self.subset.startswith("~") and os.path.exists(self.subset[1:]):
+                exclude_names = []
+                with open(self.subset[1:], "r") as f:
+                    for name in f:
+                        exclude_names.append(name.strip())
+                index = []
+                for i, name in enumerate(names_filtred):
+                    if name not in exclude_names:
+                        index.append(i)
+
         elif isinstance(self.subset, list):
             if len(self.subset) > 0:
                 if isinstance(self.subset[0], int):

{konfai-1.1.3 → konfai-1.1.5}/konfai/data/patching.py

@@ -125,23 +125,20 @@ class Accumulator():
         return len(self.patch_slices) == len([v for v in self._layer_accumulator if v is not None])
 
     def assemble(self) -> torch.Tensor:
-        if all([self._layer_accumulator[0].shape[i-len(self.patch_size)] == size for i, size in enumerate(self.patch_size)]): 
-            N = 2 if self.batch else 1
-            
-            result = torch.zeros((list(self._layer_accumulator[0].shape[:N])+list(max([[v.stop for v in patch] for patch in self.patch_slices]))), dtype=self._layer_accumulator[0].dtype).to(self._layer_accumulator[0].device)
-            for patch_slice, data in zip(self.patch_slices, self._layer_accumulator):
-                slices_dest = tuple([slice(result.shape[i]) for i in range(N)] + list(patch_slice))
-
-                for dim, s in enumerate(patch_slice):
-                    if s.stop-s.start == 1:
-                        data = data.unsqueeze(dim=dim+N)
-                if self.patchCombine is not None:
-                    result[slices_dest] += self.patchCombine(data)
-                else:
-                    result[slices_dest] = data
-            result = result[tuple([slice(None, None)]+[slice(0, s) for s in self.shape])]
-        else:
-            result = torch.cat(tuple(self._layer_accumulator), dim=0)
+        N = 2 if self.batch else 1
+        result = torch.zeros((list(self._layer_accumulator[0].shape[:N])+list(max([[v.stop for v in patch] for patch in self.patch_slices]))), dtype=self._layer_accumulator[0].dtype).to(self._layer_accumulator[0].device)
+        for patch_slice, data in zip(self.patch_slices, self._layer_accumulator):
+            slices_dest = tuple([slice(result.shape[i]) for i in range(N)] + list(patch_slice))
+
+            for dim, s in enumerate(patch_slice):
+                if s.stop-s.start == 1:
+                    data = data.unsqueeze(dim=dim+N)
+            if self.patchCombine is not None:
+                result[slices_dest] += self.patchCombine(data)
+            else:
+                result[slices_dest] = data
+        result = result[tuple([slice(None, None)]+[slice(0, s) for s in self.shape])]
+
         self._layer_accumulator.clear()
         return result
 

{konfai-1.1.3 → konfai-1.1.5}/konfai/evaluator.py

@@ -102,7 +102,7 @@ class Evaluator(DistributedObject):
         result = {}
         for output_group in self.metrics:
             for target_group in self.metrics[output_group]:
-                targets = [data_dict[group][0].to(0) if torch.cuda.is_available() else data_dict[group][0] for group in target_group.split("/") if group in data_dict]
+                targets = [data_dict[group][0].to(0) if torch.cuda.is_available() else data_dict[group][0] for group in target_group.split(";") if group in data_dict]
                 name = data_dict[output_group][1][0]
                 for metric in self.metrics[output_group][target_group]:
                     result["{}:{}:{}".format(output_group, target_group, metric.__class__.__name__)] = metric(data_dict[output_group][0].to(0) if torch.cuda.is_available() else data_dict[output_group][0], *targets).item()
@@ -135,8 +135,11 @@ class Evaluator(DistributedObject):
                 f"Available groups: {sorted(groupsDest)}"
             )
 
-        target_groups = {target for targets in self.metrics.values() for target in targets}
-        missing_targets = target_groups - set(groupsDest)
+        target_groups = []
+        for i in {target for targets in self.metrics.values() for target in targets}:
+            for u in i.split(";"):
+                target_groups.append(u)
+        missing_targets = set(target_groups) - set(groupsDest)
         if missing_targets:
             raise EvaluatorError(
                 f"The following metric target groups are missing from 'groupsDest': {sorted(missing_targets)}. ",

{konfai-1.1.3 → konfai-1.1.5}/konfai/metric/measure.py

@@ -3,6 +3,9 @@ import importlib
 import numpy as np
 import torch
 
+from torchvision.models import inception_v3, Inception_V3_Weights
+from torchvision.transforms import functional as F
+from scipy import linalg    
 
 import torch.nn.functional as F
 import os
@@ -401,21 +404,42 @@ class L1LossRepresentation(Criterion):
     def forward(self, output: torch.Tensor) -> torch.Tensor:
         return self.loss(output[0], output[1])+ self._variance(output[0]) + self._variance(output[1])
 
-"""import torch
-import torch.nn as nn
-from torchvision.models import inception_v3, Inception_V3_Weights
-from torchvision.transforms import functional as F
-from scipy import linalg
-import numpy as np
+class FocalLoss(Criterion):
+
+    def __init__(self, gamma: float = 2.0, alpha: list[float] = [0.5, 2.0, 0.5, 0.5, 1], reduction: str = "mean"):
+        super().__init__()
+        raw_alpha = torch.tensor(alpha, dtype=torch.float32)
+        self.alpha = raw_alpha / raw_alpha.sum() * len(raw_alpha)
+        self.gamma = gamma
+        self.reduction = reduction
+
+    def forward(self, output: torch.Tensor, targets: list[torch.Tensor]) -> torch.Tensor:
+        target = targets[0].long()
+        
+        logpt = F.log_softmax(output, dim=1)
+        pt = torch.exp(logpt)
+
+        logpt = logpt.gather(1, target.unsqueeze(1)) 
+        pt = pt.gather(1, target.unsqueeze(1))
+
+        at = self.alpha[target].unsqueeze(1)
+        loss = -at * ((1 - pt) ** self.gamma) * logpt
+
+        if self.reduction == "mean":
+            return loss.mean()
+        elif self.reduction == "sum":
+            return loss.sum()
+        return loss
     
+
 class FID(Criterion):
 
-    class InceptionV3(nn.Module):
+    class InceptionV3(torch.nn.Module):
 
         def __init__(self) -> None:
             super().__init__()
             self.model = inception_v3(weights=Inception_V3_Weights.DEFAULT, transform_input=False)
-            self.model.fc = nn.Identity()
+            self.model.fc = torch.nn.Identity()
             self.model.eval()
 
         def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -456,9 +480,9 @@ class FID(Criterion):
         generated_features = FID.get_features(generated_images, self.inception_model)
 
         return FID.calculate_fid(real_features, generated_features)
-        """
+        
 
-"""class MutualInformationLoss(torch.nn.Module):
+class MutualInformationLoss(torch.nn.Module):
     def __init__(self, num_bins: int = 23, sigma_ratio: float = 0.5, smooth_nr: float = 1e-7, smooth_dr: float = 1e-7) -> None:
         super().__init__()
         bin_centers = torch.linspace(0.0, 1.0, num_bins)
@@ -469,13 +493,13 @@ class FID(Criterion):
         self.smooth_nr = float(smooth_nr)
         self.smooth_dr = float(smooth_dr)
 
-    def parzen_windowing(self, pred: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    def parzen_windowing(self, pred: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         pred_weight, pred_probability = self.parzen_windowing_gaussian(pred)
         target_weight, target_probability = self.parzen_windowing_gaussian(target)
         return pred_weight, pred_probability, target_weight, target_probability
 
 
-    def parzen_windowing_gaussian(self, img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         img = torch.clamp(img, 0, 1)
         img = img.reshape(img.shape[0], -1, 1)  # (batch, num_sample, 1)
         weight = torch.exp(-self.preterm.to(img) * (img - self.bin_centers.to(img)) ** 2)  # (batch, num_sample, num_bin)
@@ -488,4 +512,4 @@ class FID(Criterion):
         pab = torch.bmm(wa.permute(0, 2, 1), wb.to(wa)).div(wa.shape[1])  # (batch, num_bins, num_bins)
         papb = torch.bmm(pa.permute(0, 2, 1), pb.to(pa))  # (batch, num_bins, num_bins)
         mi = torch.sum(pab * torch.log((pab + self.smooth_nr) / (papb + self.smooth_dr) + self.smooth_dr), dim=(1, 2))  # (batch)
-        return torch.mean(mi).neg()  # average over the batch and channel ndims"""
+        return torch.mean(mi).neg()  # average over the batch and channel ndims

konfai-1.1.5/konfai/models/segmentation/NestedUNet.py

@@ -0,0 +1,115 @@
+from konfai.network import network, blocks
+from typing import Union
+import torch
+from konfai.data.patching 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, activation: str, 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))
+            if activation == "Softmax":
+                self.add_module("Softmax", torch.nn.Softmax(dim=1))
+                self.add_module("Argmax", blocks.ArgMax(dim=1))
+            elif activation == "Tanh":
+                self.add_module("Tanh", torch.nn.Tanh())
+            
+    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",
+                    activation: str = "Softmax") -> 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, activation=activation, dim=dim), in_branch=["X_0_{}".format(j+1)], out_branch=[-1])
+
+
+
+        
+class ResNetEncoderLayer(network.ModuleArgsDict):
+
+    def __init__(self, in_channel: int, out_channel: int, nb_block: int, dim: int, downSampleMode : blocks.DownSampleMode):
+        super().__init__()
+        for i in range(nb_block):
+            if downSampleMode == blocks.DownSampleMode.MAXPOOL and i == 0:
+                self.add_module("DownSample", blocks.getTorchModule("MaxPool", dim)(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False))
+            self.add_module("ResBlock_{}".format(i), blocks.ResBlock(in_channel, out_channel, [blocks.BlockConfig(3, 2 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i == 0 else 1, 1, False, "ReLU;True", blocks.NormMode.BATCH), blocks.BlockConfig(3, 1, 1, False, None, blocks.NormMode.BATCH)], dim=dim))
+            in_channel = out_channel
+        
+def resNetEncoder(channels: list[int], layers: list[int], dim: int) -> list[torch.nn.Module]:
+    modules = []
+    modules.append(blocks.ConvBlock(channels[0], channels[1], [blocks.BlockConfig(7, 2, 3, False, "ReLU", blocks.NormMode.BATCH)], dim=dim))
+    for i, (in_channel, out_channel, layer) in enumerate(zip(channels[1:], channels[2:], layers)):
+        modules.append(ResNetEncoderLayer(in_channel, out_channel, layer, dim, blocks.DownSampleMode.MAXPOOL if i==0 else blocks.DownSampleMode.CONV_STRIDE))
+    return modules
+
+class NestedUNetBlock(network.ModuleArgsDict):
+
+    def __init__(self, encoder_channels: list[int], decoder_channels: list[int], encoders: list[torch.nn.Module], upSampleMode: blocks.UpSampleMode, dim: int, i : int = 0) -> None:
+        super().__init__()
+        self.add_module("X_{}_{}".format(i, 0), encoders[0], out_branch=["X_{}_{}".format(i, 0)])
+        if len(encoder_channels) > 2:
+            self.add_module("UNetBlock_{}".format(i+1), NestedUNetBlock(encoder_channels[1:], decoder_channels[1:], encoders[1:], upSampleMode, dim, i+1), in_branch=["X_{}_{}".format(i, 0)], out_branch=["X_{}_{}".format(i+1, j) for j in range(len(encoder_channels)-2)])
+            for j in range(len(encoder_channels)-2):
+                self.add_module("X_{}_{}_{}".format(i, j+1, upSampleMode.name), blocks.upSample(in_channels=encoder_channels[2], out_channels=decoder_channels[2] if j == len(encoder_channels)-3 else encoder_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), blocks.ConvBlock(in_channels=encoder_channels[1]*j+((encoder_channels[1]+encoder_channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else (decoder_channels[1]+decoder_channels[2] if j == len(encoder_channels)-3 else encoder_channels[1]+encoder_channels[1])), out_channels=decoder_channels[2] if j == len(encoder_channels)-3 else encoder_channels[1], blockConfigs=[blocks.BlockConfig(3,1,1,False, "ReLU;True", blocks.NormMode.BATCH)]*2, dim=dim), in_branch=["X_{}_{}".format(i, j+1)], out_branch=["X_{}_{}".format(i, j+1)])
+
+class NestedUNetHead(network.ModuleArgsDict):
+
+    def __init__(self, in_channels: int, out_channels: int, nb_class: int, dim: int) -> None:
+        super().__init__()
+        self.add_module("Upsample", blocks.upSample(in_channels=in_channels, out_channels=out_channels, upSampleMode=blocks.UpSampleMode.UPSAMPLE_BILINEAR, dim=dim))
+        self.add_module("ConvBlock", blocks.ConvBlock(in_channels=in_channels, out_channels=out_channels, blockConfigs=[blocks.BlockConfig(3,1,1,False, "ReLU;True", blocks.NormMode.BATCH)]*2, dim=dim))
+        self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = out_channels, out_channels = nb_class, kernel_size = 3, stride = 1, padding = 1))
+        if nb_class > 1:
+            self.add_module("Softmax", torch.nn.Softmax(dim=1))
+            self.add_module("Argmax", blocks.ArgMax(dim=1))
+        else:
+            self.add_module("Tanh", torch.nn.Tanh())
+        
+class UNetpp(network.Network):
+
+
+    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,
+                    encoder_channels: list[int] = [1,64,64,128,256,512],
+                    decoder_channels: list[int] = [256,128,64,32,16,1],
+                    layers: list[int] = [3,4,6,3],
+                    dim : int = 2) -> None:
+        super().__init__(in_channels = encoder_channels[0], optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, patch=patch, dim = dim)
+        self.add_module("Block_0", NestedUNetBlock(encoder_channels, decoder_channels[::-1], resNetEncoder(encoder_channels, layers, dim), blocks.UpSampleMode.UPSAMPLE_BILINEAR, dim=dim), out_branch=["X_0_{}".format(j+1) for j in range(len(encoder_channels)-2)])    
+        self.add_module("Head", NestedUNetHead(in_channels=decoder_channels[-3], out_channels=decoder_channels[-2], nb_class=decoder_channels[-1], dim=dim), in_branch=["X_0_{}".format(len(encoder_channels)-2)], out_branch=[-1])

{konfai-1.1.3 → konfai-1.1.5}/konfai/models/segmentation/UNet.py

@@ -38,7 +38,6 @@ class UNetBlock(network.ModuleArgsDict):
 
 class UNet(network.Network):
 
-    @config("UNet")
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
                     schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),

{konfai-1.1.3 → konfai-1.1.5}/konfai/network/blocks.py

@@ -19,7 +19,7 @@ class NormMode(Enum):
     SYNCBATCH = 5
     INSTANCE_AFFINE = 6
 
-def getNorm(normMode: Enum, channels : int, dim: int) -> torch.nn.Module:
+def getNorm(normMode: Enum, channels : int, dim: int) -> Union[torch.nn.Module, None]:
     if normMode == NormMode.BATCH:
         return getTorchModule("BatchNorm", dim = dim)(channels, affine=True, track_running_stats=True)
     if normMode == NormMode.INSTANCE:
@@ -32,7 +32,7 @@ def getNorm(normMode: Enum, channels : int, dim: int) -> torch.nn.Module:
         return torch.nn.GroupNorm(num_groups=32, num_channels=channels)
     if normMode == NormMode.LAYER:
         return torch.nn.GroupNorm(num_groups=1, num_channels=channels)
-    return torch.nn.Identity()
+    return None
 
 class UpSampleMode(Enum):
     CONV_TRANSPOSE = 0,
@@ -59,9 +59,9 @@ class BlockConfig():
         self.stride = stride
         self.padding = padding
         self.activation = activation
-        self.normMode = normMode
-
-        if isinstance(normMode, str):
+        if normMode is None:
+            self.norm = None
+        elif isinstance(normMode, str):
             self.norm = NormMode._member_map_[normMode]
         elif isinstance(normMode, NormMode):
             self.norm = normMode
@@ -70,9 +70,13 @@ class BlockConfig():
         return getTorchModule("Conv", dim = dim)(in_channels = in_channels, out_channels = out_channels, kernel_size = self.kernel_size, stride = self.stride, padding = self.padding, bias=self.bias)
     
     def getNorm(self, channels : int, dim: int) -> torch.nn.Module:
+        if self.norm is None:
+            return None
         return getNorm(self.norm, channels, dim) if isinstance(self.norm, NormMode) else self.norm(channels)
 
     def getActivation(self) -> torch.nn.Module:
+        if self.activation is None:
+            return None
         if isinstance(self.activation, str):
             return getTorchModule(self.activation.split(";")[0])(*[ast.literal_eval(value) for value in self.activation.split(";")[1:]]) if self.activation != "None" else torch.nn.Identity()
         return self.activation()
@@ -83,21 +87,35 @@ class ConvBlock(network.ModuleArgsDict):
         super().__init__()
         for i, blockConfig in enumerate(blockConfigs):
             self.add_module("Conv_{}".format(i), blockConfig.getConv(in_channels, out_channels, dim), alias=alias[0])
-            self.add_module("Norm_{}".format(i), blockConfig.getNorm(out_channels, dim), alias=alias[1])
-            self.add_module("Activation_{}".format(i), blockConfig.getActivation(), alias=alias[2])
+            norm = blockConfig.getNorm(out_channels, dim)
+            if norm is not None:
+                self.add_module("Norm_{}".format(i), norm, alias=alias[1])
+            activation = blockConfig.getActivation()
+            if activation is not None:
+                self.add_module("Activation_{}".format(i), activation, alias=alias[2])
             in_channels = out_channels
 
 class ResBlock(network.ModuleArgsDict):
     
-    def __init__(self, in_channels : int, out_channels : int, nb_conv: int, blockConfig : BlockConfig, dim : int, alias : list[list[str]]=[[], [], [], []]) -> None:
+    def __init__(self, in_channels : int, out_channels : int, blockConfigs : list[BlockConfig], dim : int, alias : list[list[str]]=[[], [], [], [], []]) -> None:
         super().__init__()
-        for i in range(nb_conv):
-            self.add_module("ConvBlock_{}".format(i), ConvBlock(in_channels, out_channels, nb_conv=1, blockConfig=blockConfig, dim=dim, alias=alias[:3]))
+        for i, blockConfig in enumerate(blockConfigs):
+            self.add_module("Conv_{}".format(i), blockConfig.getConv(in_channels, out_channels, dim), alias=alias[0])
+            norm = blockConfig.getNorm(out_channels, dim)
+            if norm is not None:
+                self.add_module("Norm_{}".format(i), norm, alias=alias[1])
+            activation = blockConfig.getActivation()
+            if activation is not None:
+                self.add_module("Activation_{}".format(i), activation, alias=alias[2])
+
             if in_channels != out_channels:
-                self.add_module("Conv_skip_{}".format(i), blockConfig.getConv(in_channels, out_channels, dim), alias=alias[3], in_branch=[1], out_branch=[1])    
+                self.add_module("Conv_skip", getTorchModule("Conv", dim)(in_channels, out_channels, 1, blockConfig.stride, bias=blockConfig.bias), alias=alias[3], in_branch=[1], out_branch=[1])    
+                self.add_module("Norm_skip", blockConfig.getNorm(out_channels, dim), alias=alias[4], in_branch=[1], out_branch=[1])    
             in_channels = out_channels
-            self.add_module("Add_{}".format(i), Add(), in_branch=[0,1], out_branch=[0,1])
-    
+            
+        self.add_module("Add", Add(), in_branch=[0,1])
+        self.add_module("Norm_{}".format(i+1), torch.nn.ReLU(inplace=True))
+            
 def downSample(in_channels: int, out_channels: int, downSampleMode: DownSampleMode, dim: int) -> torch.nn.Module:
     if downSampleMode == DownSampleMode.MAXPOOL:
         return getTorchModule("MaxPool", dim = dim)(2)
@@ -176,6 +194,7 @@ class Print(torch.nn.Module):
         super().__init__()
     
     def forward(self, input: torch.Tensor) -> torch.Tensor:
+        print(input.shape)
         return input
 
 class Write(torch.nn.Module):

{konfai-1.1.3 → konfai-1.1.5}/konfai/network/network.py

@@ -133,9 +133,10 @@ class Measure():
             self._loss.clear()
 
         def add(self, weight: float, value: torch.Tensor) -> None:
-            self._loss.append(value if self.isLoss else value.detach())
-            self._values.append(value.item())
-            self._weight.append(weight)
+            if self.isLoss or value != 0:
+                self._loss.append(value if self.isLoss else value.detach())
+                self._values.append(value.item())
+                self._weight.append(weight)
 
         def getLastLoss(self) -> torch.Tensor:
             return self._loss[-1]*self._weight[-1] if len(self._loss) else torch.zeros((1), requires_grad = True)
@@ -157,7 +158,7 @@ class Measure():
         outputs_group_rename = {}
         
         modules = []
-        for i,_ in model.named_modules():
+        for i,_,_ in model.named_ModuleArgsDict():
             modules.append(i)
 
         for output_group in self.outputsCriterions.keys():
@@ -167,11 +168,12 @@ class Measure():
                     "Please check that the name matches exactly a submodule or output of your model architecture."
                 )
             for target_group in self.outputsCriterions[output_group]:
-                if target_group not in group_dest:
-                    raise MeasureError(
-                        f"The target_group '{target_group}' defined in 'outputsCriterions.{output_group}.targetsCriterions' was not found in the available destination groups.",
-                        "This target_group is expected for loss or metric computation, but was not loaded in 'group_dest'.",
-                        f"Please make sure that the group '{target_group}' is defined in 'Dataset:groups_src:...:groups_dest:'{target_group}'' and correctly loaded from the dataset.")
+                for target_group_tmp in target_group.split(";"): 
+                    if target_group_tmp not in group_dest:
+                        raise MeasureError(
+                            f"The target_group '{target_group_tmp}' defined in 'outputsCriterions.{output_group}.targetsCriterions' was not found in the available destination groups.",
+                            "This target_group is expected for loss or metric computation, but was not loaded in 'group_dest'.",
+                            f"Please make sure that the group '{target_group_tmp}' is defined in 'Dataset:groups_src:...:groups_dest:'{target_group_tmp}'' and correctly loaded from the dataset.")
                 for criterion in self.outputsCriterions[output_group][target_group]:
                     if not self.outputsCriterions[output_group][target_group][criterion].isTorchCriterion:
                         outputs_group_rename[output_group] = criterion.init(model, output_group, target_group)
@@ -189,7 +191,7 @@ class Measure():
 
     def update(self, output_group: str, output : torch.Tensor, data_dict: dict[str, torch.Tensor], it: int, nb_patch: int, training: bool) -> None:
         for target_group in self.outputsCriterions[output_group]:
-            target = [data_dict[group].to(output[0].device).detach() for group in target_group.split("/") if group in data_dict]
+            target = [data_dict[group].to(output[0].device).detach() for group in target_group.split(";") if group in data_dict]
 
             for criterion, criterionsAttr in self.outputsCriterions[output_group][target_group].items():
                 if it >= criterionsAttr.stepStart and (criterionsAttr.stepStop is None or it <= criterionsAttr.stepStop):
@@ -928,6 +930,16 @@ class Network(ModuleArgsDict, ABC):
             if isinstance(module, ModuleArgsDict):
                 module._training = state
 
+class MinimalModel(Network):
+    
+    def __init__(self, model: Network, optimizer : OptimizerLoader = OptimizerLoader(),
+                    schedulers : LRSchedulersLoader = LRSchedulersLoader(),
+                    outputsCriterions: dict[str, TargetCriterionsLoader] = {"default" : TargetCriterionsLoader()},
+                    patch : Union[ModelPatch, None] = None,
+                    dim : int = 3, nb_batch_per_step = 1, init_type = "normal", init_gain = 0.02):
+        super().__init__(1, optimizer, schedulers, outputsCriterions, patch, nb_batch_per_step, init_type, init_gain, dim)
+        self.add_module("Model", model)
+
 class ModelLoader():
 
     @config("Model")
@@ -937,11 +949,13 @@ class ModelLoader():
     def getModel(self, train : bool = True, DL_args: Union[str, None] = None, DL_without=["optimizer", "schedulers", "nb_batch_per_step", "init_type", "init_gain"]) -> Network:
         if not DL_args:
             DL_args="{}.Model".format(KONFAI_ROOT())
-        model = partial(getattr(importlib.import_module(self.module), self.name), config = None, DL_args=DL_args)
-        if not train: 
-            model = partial(model, DL_without = DL_without)
-        return model()
-
+        DL_args += "."+self.name
+        model = config(DL_args)(getattr(importlib.import_module(self.module), self.name))(config = None,  DL_without = DL_without if not train else [])
+        if not isinstance(model, Network):
+            model = config(DL_args)(partial(MinimalModel, model))(config = None,  DL_without = DL_without+["model"] if not train else [])
+            model.setName(self.name)
+        return model
+    
 class CPU_Model():
 
     def __init__(self, model: Network) -> None:

{konfai-1.1.3 → konfai-1.1.5}/konfai/predictor.py

@@ -8,7 +8,7 @@ import os
 
 from konfai import MODELS_DIRECTORY, PREDICTIONS_DIRECTORY, CONFIG_FILE, MODEL, KONFAI_ROOT
 from konfai.utils.config import config
-from konfai.utils.utils import State, get_patch_slices_from_nb_patch_per_dim, NeedDevice, _getModule, DistributedObject, DataLog, description
+from konfai.utils.utils import State, get_patch_slices_from_nb_patch_per_dim, NeedDevice, _getModule, DistributedObject, DataLog, description, PredictorError
 from konfai.utils.dataset import Dataset, Attribute
 from konfai.data.data_manager import DataPrediction, DatasetIter
 from konfai.data.patching import Accumulator, PathCombine
@@ -46,7 +46,7 @@ class OutDataset(Dataset, NeedDevice, ABC):
         self.names: dict[int, str] = {}
         self.nb_data_augmentation = 0
 
-    def load(self, name_layer: str, datasets: list[Dataset]):
+    def load(self, name_layer: str, datasets: list[Dataset], groups: dict[str, str]):
         transforms_type = ["pre_transforms", "post_transforms", "final_transforms"]
         for name, _transform_type, transform_type in [(k, getattr(self, "_{}".format(k)), getattr(self, k)) for k in transforms_type]:
             
@@ -137,9 +137,22 @@ class OutSameAsGroupDataset(OutDataset):
         if self.inverse_transform:
             for transform in reversed(dataset.groups_src[self.group_src][self.group_dest].post_transforms):
                 layer = transform.inverse(self.names[index_dataset], layer, self.attributes[index_dataset][index_augmentation][index_patch])
-                
         self.output_layer_accumulator[index_dataset][index_augmentation].addLayer(index_patch, layer)
 
+
+    def load(self, name_layer: str, datasets: list[Dataset], groups: dict[str, str]):
+        super().load(name_layer, datasets, groups)
+
+        if self.group_src not in groups.keys():
+            raise PredictorError(
+                f"Source group '{self.group_src}' not found. Available groups: {list(groups.keys())}."
+            )
+
+        if self.group_dest not in groups[self.group_src]:
+            raise PredictorError(
+                f"Destination group '{self.group_dest}' not found. Available groups: {groups[self.group_src]}."
+            )
+    
     def _getOutput(self, index: int, index_augmentation: int, dataset: DatasetIter) -> torch.Tensor:
         layer = self.output_layer_accumulator[index][index_augmentation].assemble()
         name = self.names[index]
@@ -258,7 +271,6 @@ class _Predictor():
         desc = lambda : "Prediction : {}".format(description(self.modelComposite))
         self.dataloader_prediction.dataset.load("Prediction")
         with tqdm.tqdm(iterable = enumerate(self.dataloader_prediction), leave=True, desc = desc(), total=len(self.dataloader_prediction), ncols=0) as batch_iter:
-            dist.barrier()
             for it, data_dict in batch_iter:
                 with torch.amp.autocast('cuda', enabled=self.autocast):
                     input = self.getInput(data_dict)
@@ -272,7 +284,7 @@ class _Predictor():
 
                     batch_iter.set_description(desc())
                     self.it += 1
-
+        
     def _predict_log(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]):
         measures = DistributedObject.getMeasure(self.world_size, self.global_rank, self.local_rank, {"" : self.modelComposite.module}, 1)
         
@@ -404,9 +416,20 @@ class Predictor(DistributedObject):
 
         shutil.copyfile(CONFIG_FILE(), self.predict_path+"Prediction.yml")
 
+        
         self.model.init(self.autocast, State.PREDICTION, self.dataset.getGroupsDest())
         self.model.init_outputsGroup()
         self.model._compute_channels_trace(self.model, self.model.in_channels, None, self.gpu_checkpoints)
+        
+        modules = []
+        for i,_,_ in self.model.named_ModuleArgsDict():
+            modules.append(i)
+        for output_group in self.outsDataset.keys():
+            if output_group not in modules:
+                raise PredictorError("The output group '{}' defined in 'outputsCriterions' does not correspond to any module in the model.".format(output_group),
+                    "Available modules: {}".format(modules),
+                    "Please check that the name matches exactly a submodule or output of your model architecture."
+                )
         self.modelComposite = ModelComposite(self.model, len(MODEL().split(":")), self.combine)
         self.modelComposite.load(self._load())
 
@@ -416,7 +439,7 @@ class Predictor(DistributedObject):
         self.size = (len(self.gpu_checkpoints)+1 if self.gpu_checkpoints else 1)
         self.dataloader = self.dataset.getData(world_size//self.size)
         for name, outDataset in self.outsDataset.items():
-            outDataset.load(name.replace(".", ":"), list(self.dataset.datasets.values()))
+            outDataset.load(name.replace(".", ":"), list(self.dataset.datasets.values()), {src : dest for src, inner in self.dataset.groups_src.items() for dest in inner})
            
            
     def run_process(self, world_size: int, global_rank: int, local_rank: int, dataloaders: list[DataLoader]):

{konfai-1.1.3 → konfai-1.1.5}/konfai/trainer.py

@@ -34,7 +34,7 @@ class EarlyStoppingBase:
 class EarlyStopping(EarlyStoppingBase):
 
     @config("EarlyStopping")
-    def __init__(self, monitor: Union[list[str], None] = [], patience=10, min_delta=0.0, mode="min"):
+    def __init__(self, monitor: Union[list[str], None] = [], patience: int=10, min_delta: float=0.0, mode: str="min"):
         super().__init__()
         self.monitor = [] if monitor is None else monitor
         self.patience = patience
@@ -120,7 +120,8 @@ class _Trainer():
 
     def run(self) -> None:
         self.dataloader_training.dataset.load("Train")
-        self.dataloader_validation.dataset.load("Validation")
+        if self.dataloader_validation is not None:
+            self.dataloader_validation.dataset.load("Validation")
         with tqdm.tqdm(iterable = range(self.epoch, self.epochs), leave=False, total=self.epochs, initial=self.epoch, desc="Progress") as epoch_tqdm:
             for self.epoch in epoch_tqdm:
                 self.train()

{konfai-1.1.3 → konfai-1.1.5}/konfai/utils/config.py

@@ -3,7 +3,7 @@ import ruamel.yaml
 import inspect 
 import collections
 from copy import deepcopy
-from typing import Union, Literal, get_origin, get_args
+from typing import Union, Literal, get_origin, get_args, Any
 import torch
 from konfai import CONFIG_FILE
 from konfai.utils.utils import ConfigError
@@ -175,8 +175,11 @@ def config(key : Union[str, None] = None):
                     os.environ['KONFAI_CONFIG_VARIABLE'] = "False"
                     kwargs = {} 
                     for param in list(inspect.signature(function).parameters.values())[len(args):]:
-
+                        if param.name in without:
+                            continue
+                        
                         annotation = param.annotation
+
                         # --- support Literal ---
                         if get_origin(annotation) is Literal:
                             allowed_values = get_args(annotation)
@@ -192,11 +195,10 @@ def config(key : Union[str, None] = None):
                             for i in annotation.__args__:
                                 annotation = i
                                 break
-                        if param.name in without:
-                            continue
+                        
                         if not annotation == inspect._empty:
                             if annotation not in [int, str, bool, float, torch.Tensor]:
-                                if str(annotation).startswith("list") or str(annotation).startswith("tuple") or str(annotation).startswith("typing.Tuple") or str(annotation).startswith("typing.List"):
+                                if str(annotation).startswith("list") or str(annotation).startswith("tuple") or str(annotation).startswith("typing.Tuple") or str(annotation).startswith("typing.List") or str(annotation).startswith("typing.Sequence"):
                                     elem_type = annotation.__args__[0]
                                     values = config.getValue(param.name, param.default)
                                     if getattr(elem_type, '__origin__', None) is Union:
@@ -221,7 +223,7 @@ def config(key : Union[str, None] = None):
                                 elif str(annotation).startswith("dict"):
                                     if annotation.__args__[0] == str:
                                         values = config.getValue(param.name, param.default)
-                                        if values is not None and annotation.__args__[1] not in [int, str, bool, float]:
+                                        if values is not None and annotation.__args__[1] not in [int, str, bool, float, Any]:
                                             try:
                                                 kwargs[param.name] = {value : annotation.__args__[1](config = filename, DL_args = key_tmp+"."+param.name+"."+value) for value in values}
                                             except ValueError as e:
@@ -229,6 +231,7 @@ def config(key : Union[str, None] = None):
                                             except Exception as e:
                                                 raise ConfigError("{} {}".format(values, e))
                                         else:
+
                                             kwargs[param.name] = values
                                     else: 
                                         raise ConfigError("Config: The config only supports types : config(Object), int, str, bool, float, list[int], list[str], list[bool], list[float], dict[str, Object]")
@@ -236,7 +239,7 @@ def config(key : Union[str, None] = None):
                                     try:
                                         kwargs[param.name] = annotation(config = filename, DL_args = key_tmp)
                                     except Exception as e:
-                                        raise ConfigError("Failed to instantiate {} with type {}, error {} ".format(param.name, annotation.__name__, e))
+                                        raise ConfigError("Failed to instantiate {} with type {}, error {} ".format(param.name, annotation, e))
                                     
                                     if os.environ['KONFAI_CONFIG_VARIABLE'] == "True":
                                         os.environ['KONFAI_CONFIG_VARIABLE'] = "False"

{konfai-1.1.3 → konfai-1.1.5}/konfai/utils/utils.py

@@ -21,7 +21,7 @@ import sys
 import re
 
 
-def description(model, modelEMA = None, showMemory: bool = True) -> str:
+def description(model, modelEMA = None, showMemory: bool = True, train: bool = True) -> str:
     values_desc = lambda weights, values: " ".join(["{}({:.2f}) : {:.6f}".format(name.split(":")[-1], weight, value) for (name, value), weight in zip(values.items(), weights.values())])
     model_desc = lambda model : "("+" ".join(["{}({:.6f}) : {}".format(name, network.optimizer.param_groups[0]['lr'] if network.optimizer is not None else 0, values_desc(network.measure.getLastWeights(), network.measure.getLastValues())) for name, network in model.module.getNetworks().items() if network.measure is not None])+")"
     result = "Loss {}".format(model_desc(model))
@@ -33,9 +33,9 @@ def description(model, modelEMA = None, showMemory: bool = True) -> str:
     return result
 
 def _getModule(classpath : str, type : str) -> tuple[str, str]:
-    if len(classpath.split("_")) > 1:
-        module = ".".join(classpath.split("_")[:-1])
-        name = classpath.split("_")[-1] 
+    if len(classpath.split(":")) > 1:
+        module = ".".join(classpath.split(":")[:-1])
+        name = classpath.split(":")[-1] 
     else:
         module = "konfai."+type+("." if len(classpath.split(".")) > 2 else "")+".".join(classpath.split(".")[:-1])
         name = classpath.split(".")[-1]
@@ -586,6 +586,11 @@ class EvaluatorError(KonfAIError):
     def __init__(self, *message) -> None:
         super().__init__("Evaluator", message)
 
+class PredictorError(KonfAIError):
+
+    def __init__(self, *message) -> None:
+        super().__init__("Predictor", message)
+
 class TransformError(KonfAIError):
 
     def __init__(self, *message) -> None:

{konfai-1.1.3 → konfai-1.1.5}/konfai.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.3
+Version: 1.1.5
 Summary: Modular and configurable Deep Learning framework with YAML and PyTorch
 Author-email: Valentin Boussot <boussot.v@gmail.com>
 License-Expression: Apache-2.0

{konfai-1.1.3 → konfai-1.1.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "konfai"
-version = "1.1.3"
+version = "1.1.5"
 description = "Modular and configurable Deep Learning framework with YAML and PyTorch"
 readme = "README.md"
 requires-python = ">=3.8"

konfai-1.1.3/konfai/models/segmentation/NestedUNet.py

@@ -1,53 +0,0 @@
-from typing import Union
-import torch
-
-from konfai.network import network, blocks
-from konfai.utils.config import config
-from konfai.data.patching 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])