konfai 1.1.2__tar.gz → 1.1.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.2
+Version: 1.1.3
 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 → konfai-1.1.3}/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,7 +131,7 @@ 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")
@@ -194,7 +199,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))
@@ -207,11 +212,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:]))])
@@ -232,8 +237,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]]:
@@ -244,8 +249,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]]:
@@ -256,8 +261,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 +273,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)) 
     
@@ -281,8 +286,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)) 
 
@@ -368,8 +373,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 +415,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))])
@@ -423,8 +428,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 +439,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 = []
@@ -513,7 +518,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))
@@ -546,7 +551,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
@@ -563,8 +568,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 +582,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-1.1.2 → konfai-1.1.3}/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)
 
@@ -390,7 +390,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-1.1.2 → konfai-1.1.3}/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-1.1.2 → konfai-1.1.3}/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-1.1.2 → konfai-1.1.3}/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-1.1.2 → konfai-1.1.3}/konfai/models/segmentation/UNet.py

@@ -7,32 +7,33 @@ 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):
@@ -51,8 +52,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-1.1.2 → konfai-1.1.3}/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")

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

@@ -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-1.1.2 → konfai-1.1.3}/konfai.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.2
+Version: 1.1.3
 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 → konfai-1.1.3}/pyproject.toml

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