konfai 1.1.2__py3-none-any.whl → 1.1.4__py3-none-any.whl

konfai/data/augmentation.py

@@ -9,8 +9,7 @@ import os
 from konfai import KONFAI_ROOT
 from konfai.utils.config import config
 from konfai.utils.utils import _getModule, AugmentationError
-from konfai.utils.dataset import Attribute, data_to_image
-
+from konfai.utils.dataset import Attribute, data_to_image, Dataset
 
 def _translate2DMatrix(t: torch.Tensor) -> torch.Tensor:
     return torch.cat((torch.cat((torch.eye(2), torch.tensor([[t[0]], [t[1]]])), dim=1), torch.Tensor([[0,0,1]])), dim=0)
@@ -56,23 +55,29 @@ class DataAugmentationsList():
         self.dataAugmentations : list[DataAugmentation] = []
         self.dataAugmentationsLoader = dataAugmentations
 
-    def load(self, key: str):
+    def load(self, key: str, datasets: list[Dataset]):
         for augmentation, prob in self.dataAugmentationsLoader.items():
             module, name = _getModule(augmentation, "data.augmentation")
-            dataAugmentation: DataAugmentation = getattr(importlib.import_module(module), name)(config = None, DL_args="{}.Dataset.augmentations.{}.dataAugmentations".format(KONFAI_ROOT(), key))
+            dataAugmentation: DataAugmentation = config("{}.Dataset.augmentations.{}.dataAugmentations.{}".format(KONFAI_ROOT(), key, augmentation))(getattr(importlib.import_module(module), name))(config = None)
             dataAugmentation.load(prob.prob)
+            dataAugmentation.setDatasets(datasets)
             self.dataAugmentations.append(dataAugmentation)
-    
+
 class DataAugmentation(ABC):
 
-    def __init__(self) -> None:
+    def __init__(self, groups: Union[list[str], None] = None) -> None:
         self.who_index: dict[int, list[int]] = {}
         self.shape_index: dict[int, list[list[int]]] = {}
         self._prob: float = 0
+        self.groups = groups
+        self.datasets : list[Dataset] = []
 
     def load(self, prob: float):
         self._prob = prob
 
+    def setDatasets(self, datasets: list[Dataset]):
+        self.datasets = datasets
+    
     def state_init(self, index: Union[None, int], shapes: list[list[int]], caches_attribute: list[Attribute]) -> list[list[int]]:
         if index is not None:
             if index not in self.who_index:
@@ -93,14 +98,14 @@ class DataAugmentation(ABC):
     def _state_init(self, index : int, shapes: list[list[int]], caches_attribute: list[Attribute]) -> list[list[int]]:
         pass
 
-    def __call__(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def __call__(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         if len(self.who_index[index]) > 0:
-            for i, result in enumerate(self._compute(index, [inputs[i] for i in self.who_index[index]], device)):
+            for i, result in enumerate(self._compute(name, index, [inputs[i] for i in self.who_index[index]], device)):
                 inputs[self.who_index[index][i]] = result if device is None else result.cpu()
         return inputs
     
     @abstractmethod
-    def _compute(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def _compute(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         pass
 
     def inverse(self, index: int, a: int, input : torch.Tensor) -> torch.Tensor:
@@ -118,7 +123,7 @@ class EulerTransform(DataAugmentation):
         super().__init__()
         self.matrix: dict[int, list[torch.Tensor]] = {}
 
-    def _compute(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def _compute(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         results = []
         for input, matrix in zip(inputs, self.matrix[index]):
             results.append(F.grid_sample(input.unsqueeze(0).type(torch.float32), F.affine_grid(matrix[:, :-1,...], [1]+list(input.shape), align_corners=True).to(input.device), align_corners=True, mode="bilinear", padding_mode="reflection").type(input.dtype).squeeze(0))
@@ -126,10 +131,9 @@ class EulerTransform(DataAugmentation):
     
     def _inverse(self, index: int, a: int, input : torch.Tensor) -> torch.Tensor:
         return F.grid_sample(input.unsqueeze(0).type(torch.float32), F.affine_grid(self.matrix[index][a].inverse()[:, :-1,...], [1]+list(input.shape), align_corners=True).to(input.device), align_corners=True, mode="bilinear", padding_mode="reflection").type(input.dtype).squeeze(0)
-
+                                                                                   
 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
@@ -147,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
@@ -169,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
@@ -182,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
@@ -194,7 +195,7 @@ class Flip(DataAugmentation):
         self.flip[index] = [dims[mask].tolist() for mask in prob]
         return shapes
 
-    def _compute(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def _compute(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         results = []
         for input, flip in zip(inputs, self.flip[index]):
             results.append(torch.flip(input, dims=flip))
@@ -206,12 +207,11 @@ class Flip(DataAugmentation):
 
 class ColorTransform(DataAugmentation):
 
-    @config("ColorTransform")
-    def __init__(self) -> None:
-        super().__init__()
+    def __init__(self, groups: Union[list[str], None] = None) -> None:
+        super().__init__(groups)
         self.matrix: dict[int, list[torch.Tensor]] = {}
     
-    def _compute(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def _compute(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         results = []
         for input, matrix in zip(inputs, self.matrix[index]):
             result = input.reshape([*input.shape[:1], int(np.prod(input.shape[1:]))])
@@ -231,9 +231,8 @@ class ColorTransform(DataAugmentation):
 
 class Brightness(ColorTransform):
 
-    @config("Brightness")
-    def __init__(self, b_std: float) -> None:
-        super().__init__()
+    def __init__(self, b_std: float, groups: Union[list[str], None] = None) -> None:
+        super().__init__(groups)
         self.b_std = b_std
 
     def _state_init(self, index : int, shapes: list[list[int]], caches_attribute: list[Attribute]) -> list[list[int]]:
@@ -243,9 +242,8 @@ class Brightness(ColorTransform):
 
 class Contrast(ColorTransform):
 
-    @config("Contrast")
-    def __init__(self, c_std: float) -> None:
-        super().__init__()
+    def __init__(self, c_std: float, groups: Union[list[str], None] = None) -> None:
+        super().__init__(groups)
         self.c_std = c_std
 
     def _state_init(self, index : int, shapes: list[list[int]], caches_attribute: list[Attribute]) -> list[list[int]]:
@@ -255,9 +253,8 @@ class Contrast(ColorTransform):
 
 class LumaFlip(ColorTransform):
 
-    @config("LumaFlip")
-    def __init__(self) -> None:
-        super().__init__()
+    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)) 
     
     def _state_init(self, index : int, shapes: list[list[int]], caches_attribute: list[Attribute]) -> list[list[int]]:
@@ -268,8 +265,8 @@ class LumaFlip(ColorTransform):
 class HUE(ColorTransform):
 
     @config("HUE")
-    def __init__(self, hue_max: float) -> None:
-        super().__init__()
+    def __init__(self, hue_max: float, groups: Union[list[str], None] = None) -> None:
+        super().__init__(groups)
         self.hue_max = hue_max
         self.v = torch.tensor([1, 1, 1])/torch.sqrt(torch.tensor(3)) 
     
@@ -280,9 +277,8 @@ class HUE(ColorTransform):
     
 class Saturation(ColorTransform):
 
-    @config("Saturation")
-    def __init__(self, s_std: float) -> None:
-        super().__init__()
+    def __init__(self, s_std: float, groups: Union[list[str], None] = None) -> None:
+        super().__init__(groups)
         self.s_std = s_std
         self.v = torch.tensor([1, 1, 1, 0])/torch.sqrt(torch.tensor(3)) 
 
@@ -367,9 +363,8 @@ 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) -> None:
-        super().__init__()
+    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
         self.noise_step = noise_step
 
@@ -410,7 +405,7 @@ class Noise(DataAugmentation):
             self.ts[index] = [torch.randint(0, int(self.max_T), (1,)) for _ in shapes]
         return shapes
     
-    def _compute(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def _compute(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         results = []
         for input, t in zip(inputs, self.ts[index]):
             alpha_hat_t = self.alpha_hat[t].to(input.device).reshape(*[1 for _ in range(len(input.shape))])
@@ -422,9 +417,8 @@ class Noise(DataAugmentation):
 
 class CutOUT(DataAugmentation):
 
-    @config("CutOUT")
-    def __init__(self, c_prob: float, cutout_size: int, value: float) -> None:
-        super().__init__()
+    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
         self.cutout_size = cutout_size
         self.centers: dict[int, list[torch.Tensor]] = {}
@@ -434,7 +428,7 @@ class CutOUT(DataAugmentation):
         self.centers[index] = [torch.rand((3) if len(shape) == 3 else (2)) for shape in shapes]
         return shapes
     
-    def _compute(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def _compute(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         results = []
         for input, center in zip(inputs, self.centers[index]):
             masks = []
@@ -453,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
@@ -513,7 +506,7 @@ class Elastix(DataAugmentation):
             print("Compute in progress : {:.2f} %".format((i+1)/len(shapes)*100))
         return shapes
     
-    def _compute(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def _compute(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         results = []
         for input, displacement_field in zip(inputs, self.displacement_fields[index]):
             results.append(F.grid_sample(input.type(torch.float32).unsqueeze(0), displacement_field.to(input.device), align_corners=True, mode="bilinear", padding_mode="border").type(input.dtype).squeeze(0))
@@ -524,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]])
@@ -546,7 +538,7 @@ class Permute(DataAugmentation):
                     shapes[i] = [shapes[i][dim-1] for dim in permute[1:]]
         return shapes
     
-    def _compute(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def _compute(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         results = []
         for input, prob in zip(inputs, self.permute[index]):
             res = input
@@ -562,9 +554,8 @@ class Permute(DataAugmentation):
 
 class Mask(DataAugmentation):
 
-    @config("Mask")
-    def __init__(self, mask: str, value: float) -> None:
-        super().__init__()
+    def __init__(self, mask: str, value: float, groups: Union[list[str], None] = None) -> None:
+        super().__init__(groups)
         if mask is not None:
             if os.path.exists(mask):
                 self.mask = torch.tensor(sitk.GetArrayFromImage(sitk.ReadImage(mask)))
@@ -577,7 +568,7 @@ class Mask(DataAugmentation):
         self.positions[index] = [torch.rand((3) if len(shape) == 3 else (2))*(torch.tensor([max(s1-s2, 0) for s1, s2 in zip(torch.tensor(shape), torch.tensor(self.mask.shape))])) for shape in shapes]
         return [self.mask.shape for _ in shapes]
     
-    def _compute(self, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
+    def _compute(self, name: str, index: int, inputs : list[torch.Tensor], device: Union[torch.device, None]) -> list[torch.Tensor]:
         results = []
         for input, position in zip(inputs, self.positions[index]):
             slices = [slice(None, None)]+[slice(int(s1), int(s1)+s2) for s1, s2 in zip(position, self.mask.shape)]

konfai/data/data_manager.py

@@ -184,7 +184,7 @@ class DatasetIter(data.Dataset):
         data = {}
         x, a, p = self.map[index]
         if x not in self._index_cache:
-            if len(self._index_cache) >= self.buffer_size and not self.use_cache:
+            if x not in self._index_cache and len(self._index_cache) >= self.buffer_size and not self.use_cache:
                 self._unloadData(self._index_cache[0])
             self._loadData(x)
 
@@ -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):
@@ -390,7 +398,7 @@ class Data(ABC):
             )
 
         for key, dataAugmentations in self.dataAugmentationsList.items():
-            dataAugmentations.load(key)
+            dataAugmentations.load(key, [self.datasets[filename] for filename, _ in datasets[group_src]])
 
         names = set()
         dataset_name : dict[str, dict[str, list[str]]] = {}

konfai/data/patching.py

@@ -326,7 +326,8 @@ class DatasetManager():
         for dataAugmentations in dataAugmentationsList:
             a_data = [self.data[0].clone() for _ in range(dataAugmentations.nb)]
             for dataAugmentation in dataAugmentations.dataAugmentations:
-                a_data = dataAugmentation(self.index, a_data, device)
+                if dataAugmentation.groups is None or self.group_dest in dataAugmentation.groups:
+                    a_data = dataAugmentation(self.name, self.index, a_data, device)
             
             for d in a_data:
                 self.data.append(d)

konfai/data/transform.py

@@ -52,7 +52,7 @@ class Clip(Transform):
         input[torch.where(input < self.min_value)] = self.min_value
         input[torch.where(input > self.max_value)] = self.max_value
         if self.saveClip_min:
-            cache_attribute["Min"] = self .min_value
+            cache_attribute["Min"] = self.min_value
         if self.saveClip_max:
             cache_attribute["Max"] = self.max_value
         return input

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/main.py

@@ -9,6 +9,8 @@ import sys
 sys.path.insert(0, os.getcwd())
 
 def main():
+    import tracemalloc
+
     parser = argparse.ArgumentParser(description="KonfAI", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     try:
         with setup(parser) as distributedObject:
@@ -22,7 +24,6 @@ def main():
     except KeyboardInterrupt:
         print("\n[KonfAI] Manual interruption (Ctrl+C)")
 
-
 def cluster():
     parser = argparse.ArgumentParser(description="KonfAI", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 

konfai/models/segmentation/NestedUNet.py

@@ -31,7 +31,6 @@ class NestedUNet(network.Network):
             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(),

konfai/models/segmentation/UNet.py

@@ -7,37 +7,37 @@ from konfai.data.patching import ModelPatch
 
 class UNetHead(network.ModuleArgsDict):
 
-    def __init__(self, in_channels: int, nb_class: int, dim: int) -> None:
+    def __init__(self, in_channels: int, nb_class: int, dim: int, level: 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:
+    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) -> 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)
+                blockConfig_stride = blocks.BlockConfig(blockConfig.kernel_size, 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("UNetBlock_{}".format(i+1), UNetBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, block, nb_class, dim, i+1))
             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])
+                self.add_module("Head", UNetHead(channels[1], nb_class, dim, i), 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(upSampleMode.name, blocks.upSample(in_channels=channels[1], out_channels=channels[0], upSampleMode=upSampleMode, dim=dim, kernel_size=2, stride=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(),
@@ -51,8 +51,6 @@ class UNet(network.Network):
                     downSampleMode: str = "MAXPOOL",
                     upSampleMode: str = "CONV_TRANSPOSE",
                     attention : bool = False,
-                    blockType: str = "Conv",
-                    mri: bool = False) -> None:
+                    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", 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))    
-        
+        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))    

konfai/network/network.py

@@ -157,7 +157,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 +167,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 +190,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 +929,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 +948,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/predictor.py

@@ -91,6 +91,23 @@ class OutDataset(Dataset, NeedDevice, ABC):
         super().write(self.group, name, layer.numpy(), self.attributes[index][0][0])
         self.attributes.pop(index)
 
+class Reduction():
+
+    def __init__(self):
+        pass
+
+class ReductionMean():
+
+    def __init__(self):
+        pass
+    
+class ReductionMedian():
+
+    def __init__(self):
+        pass
+    
+    
+
 class OutSameAsGroupDataset(OutDataset):
 
     @config("OutDataset")
@@ -241,7 +258,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)
@@ -255,7 +271,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)
         

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
@@ -109,6 +109,8 @@ class _Trainer():
             for data in data_log:
                 self.data_log[data.split("/")[0].replace(":", ".")] = (DataLog.__getitem__(data.split("/")[1]).value[0], int(data.split("/")[2]))
         
+       
+        
     def __enter__(self):
         return self
     
@@ -314,6 +316,18 @@ class Trainer(DistributedObject):
         self.ema_decay = ema_decay
         self.modelEMA : Union[torch.optim.swa_utils.AveragedModel, None] = None
         self.data_log = data_log
+
+        modules = []
+        for i,_ in self.model.named_modules():
+            modules.append(i)
+        for k in self.data_log:
+            tmp = k.split("/")[0].replace(":", ".")
+            if tmp not in self.dataset.getGroupsDest() and tmp not in modules:
+                raise TrainerError( f"Invalid key '{tmp}' in `data_log`.",
+                                   f"This key is neither a destination group from the dataset ({self.dataset.getGroupsDest()})",
+                                    f"nor a valid module name in the model ({modules}).",
+                "Please check your `data_log` configuration — it should reference either a model output or a dataset group.")
+            
         self.gradient_checkpoints = gradient_checkpoints
         self.gpu_checkpoints = gpu_checkpoints
         self.save_checkpoint_mode = save_checkpoint_mode

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/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]

{konfai-1.1.2.dist-info → konfai-1.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.2
+Version: 1.1.4
 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.2.dist-info → konfai-1.1.4.dist-info}/RECORD

@@ -1,13 +1,13 @@
 konfai/__init__.py,sha256=YXG-wpSEXWs6Jt3BDI77V4r89gEUNX-6lxW9btj5VYI,851
-konfai/evaluator.py,sha256=rAhfdRemMjzC3VoaqyQKJR0SBekuLDiLT1nhblH8RQk,8293
-konfai/main.py,sha256=rTTJl-biaX4CkLNxPtqwwsrybXSDxjWlU39TE2ImU5o,2574
-konfai/predictor.py,sha256=e5V7awlIRDE1NkLTR_D6fUuDg93ZFGaGLTDR71KeqM8,22549
-konfai/trainer.py,sha256=NQDzp4hUKYAcDdKmeBcAE9QNFW9KGW2qV-7q_PXeWi8,19509
+konfai/evaluator.py,sha256=G8tXI-g8yTxBZXmESpRW7B5UEI4SMXq47fsGksU0G0M,8394
+konfai/main.py,sha256=xhCIs3VKXxNCgyyN616K-VANEFCsZZRR-l8LIMHlmPw,2597
+konfai/predictor.py,sha256=R9ohOqEcSoq6Wu9Yv3pkerVfm6qSotd4aICe7ctOsU0,22728
+konfai/trainer.py,sha256=gJ7qmU302M3Mevq3l1aIdPpey5C3prtibD69LyxbnAg,20240
 konfai/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-konfai/data/augmentation.py,sha256=ASmKWBpykLBHDB_YeTgoBTlqvJ06v5OUG7n7ugPN6NU,31718
-konfai/data/data_manager.py,sha256=9VpF4CTTAnqS1Hq0csbTE-XAocRwziFssITg-SZhA4A,29603
-konfai/data/patching.py,sha256=OtGNs99jKQSYmuj8B3MNGcbupKOcX5PUpMCIe0pKnX0,15636
-konfai/data/transform.py,sha256=dTZh2CgfxNitHW-HEkO0jQ9GG26-Wf-YlI-qoI2mqC0,26472
+konfai/data/augmentation.py,sha256=ruvS7ykFveOZBe2Mv_TOU5a8AzUgH7gfTipwqYgRM7k,32293
+konfai/data/data_manager.py,sha256=TCQFZ_zpRHYIomPOotG8U_I0DfpjhSaajb7ypd6fyRU,30103
+konfai/data/patching.py,sha256=oq43r3JlS9hw-pnzpjT3epvXJxEKqiQEtFVZPjY9Jic,15749
+konfai/data/transform.py,sha256=ZZoxft0V8t-DaKjF3OsIFezz41B8_xwPqBvYCYwF0dk,26471
 konfai/metric/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 konfai/metric/measure.py,sha256=K7s15wNJm1_Iaav9y7Oe_UhSnbLUiybhLTVumHZ8ZaY,22001
 konfai/metric/schedulers.py,sha256=UoSr1TW_hrus3DhvOEbDefxCSUGz7lJS_8vbz0GEye8,1370
@@ -20,20 +20,20 @@ konfai/models/generation/gan.py,sha256=-GoKxHm3W9NdD4U77UcJrG5TfOZ3NWFUZG663kt2X
 konfai/models/generation/vae.py,sha256=_3JYVT2ojZ0P98tYcD2ny7a-gWVUmnByLDhY7i-n_4g,4719
 konfai/models/registration/registration.py,sha256=18EiWt4RJIXLyFtqU-kHjV1sMnQRm9mxAA6_-2B1YqI,6313
 konfai/models/representation/representation.py,sha256=RwQYoxtdph440-t_ZLelykl0hkUAD1zdspQaLkgxb-0,2677
-konfai/models/segmentation/NestedUNet.py,sha256=6XGizAIc4bDL8vx4AHW8BBFjUvovRYcjdMBHsN4ViNo,4301
-konfai/models/segmentation/UNet.py,sha256=9-g63oNqaxSlGmrD1-IVzJ-kac9QphmM7i-3rEsH23I,4218
+konfai/models/segmentation/NestedUNet.py,sha256=wutcHYpZkOm8FqPPt5zifD6dGSDlhB9zIBlibzVyAYY,4275
+konfai/models/segmentation/UNet.py,sha256=Uei51-_0Le24ZJs9tbuL6hQn66kigkfbaLSMlH38els,4045
 konfai/network/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 konfai/network/blocks.py,sha256=RaTI0Lrvq1V-GIFei-WTUB6wlg4LydZksAyJ8DMk40M,13502
-konfai/network/network.py,sha256=SnwQUKFTZZU8zuix5sm0vW8s0G3h6Fa8pHoIcwC984Q,46512
+konfai/network/network.py,sha256=YCPfgQ7EuRiTgAlMlQTx5ospluO4aiqSbsdU4GdLiMA,47467
 konfai/utils/ITK.py,sha256=OxTieDNNYHGkn7zxJsAG-6ecRG1VYMvn1dlBbBe1DOs,13955
 konfai/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-konfai/utils/config.py,sha256=OsCUqgNaPX56nD_MpJ2WRpiQ7xxqHonfPYaS9VGtsl8,12413
+konfai/utils/config.py,sha256=f7o83ix5_oNbr2pki-Czqr-yHi-8n92ZL64nlo0XGwA,12514
 konfai/utils/dataset.py,sha256=6ZzevdhJ7e5zlXATAVwSh9O6acKXM7gYNxkMAa5DrmM,36351
 konfai/utils/registration.py,sha256=v1srEBOcgDnHrx0YtsK6bcj0yCMH7wNeaQ3wC7gEvOw,8898
-konfai/utils/utils.py,sha256=Laq8bGc5mGKFZlJIkHxa-BrC9uR2F7MTRuL4YFAIxQY,23439
-konfai-1.1.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-konfai-1.1.2.dist-info/METADATA,sha256=Usg7hzW1xmlHswHZb3fC2Hgl2V1Jkq43J0T5j5TDEIw,2515
-konfai-1.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-konfai-1.1.2.dist-info/entry_points.txt,sha256=fG82HRN5-g39ACSOCtij_I3N6EHxfYnMR0D7TI_8pW8,81
-konfai-1.1.2.dist-info/top_level.txt,sha256=xF470dkIlFoFqTZEOlRehKJr4WU_8OKGXrJqYm9vWKs,7
-konfai-1.1.2.dist-info/RECORD,,
+konfai/utils/utils.py,sha256=-W5fbsYsTh4wIOEZBEQY7krZvGU7JC8tXHYC2rZAkaM,23459
+konfai-1.1.4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+konfai-1.1.4.dist-info/METADATA,sha256=86Azx-ool5VrcYT48km-OOMvvBQzjJX1jZihzDxAhgs,2515
+konfai-1.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+konfai-1.1.4.dist-info/entry_points.txt,sha256=fG82HRN5-g39ACSOCtij_I3N6EHxfYnMR0D7TI_8pW8,81
+konfai-1.1.4.dist-info/top_level.txt,sha256=xF470dkIlFoFqTZEOlRehKJr4WU_8OKGXrJqYm9vWKs,7
+konfai-1.1.4.dist-info/RECORD,,