konfai 1.1.1__py3-none-any.whl → 1.1.3__py3-none-any.whl

konfai/__init__.py

@@ -12,5 +12,5 @@ CONFIG_FILE = lambda : os.environ["KONFAI_CONFIG_FILE"]
 KONFAI_STATE = lambda : os.environ["KONFAI_STATE"]
 KONFAI_ROOT = lambda : os.environ["KONFAI_ROOT"]
 CUDA_VISIBLE_DEVICES = lambda : os.environ["CUDA_VISIBLE_DEVICES"]
-
+KONFAI_NB_CORES = lambda : os.environ["KONFAI_NB_CORES"]
 DATE = lambda : datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")

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/data/data_manager.py

@@ -62,12 +62,25 @@ class GroupTransform:
         for transform in self.post_transforms:
             transform.setDevice(device)
 
+class GroupTransformMetric(GroupTransform):
+
+    @config()
+    def __init__(self,  pre_transforms : Union[dict[str, TransformLoader], list[Transform]] = {"default:Normalize:Standardize:Unsqueeze:TensorCast:ResampleIsotropic:ResampleResize": TransformLoader()},
+                        post_transforms : Union[dict[str, TransformLoader], list[Transform]] = {"default:Normalize:Standardize:Unsqueeze:TensorCast:ResampleIsotropic:ResampleResize": TransformLoader()}):
+        super().__init__(pre_transforms, post_transforms)
+
 class Group(dict[str, GroupTransform]):
 
     @config()
     def __init__(self, groups_dest: dict[str, GroupTransform] = {"default:group_dest": GroupTransform()}):
         super().__init__(groups_dest)
 
+class GroupMetric(dict[str, GroupTransformMetric]):
+
+    @config()
+    def __init__(self, groups_dest: dict[str, GroupTransformMetric] = {"default:group_dest": GroupTransformMetric()}):
+        super().__init__(groups_dest)
+
 class CustomSampler(Sampler[int]):
 
     def __init__(self, size: int, shuffle: bool = False) -> None:
@@ -109,32 +122,33 @@ class DatasetIter(data.Dataset):
     def getDatasetFromIndex(self, group_dest: str, index: int) -> DatasetManager:
         return self.data[group_dest][index]
     
-    def resetAugmentation(self):
-        if self.inlineAugmentations:
+    def resetAugmentation(self, label):
+        if self.inlineAugmentations and len(self.dataAugmentationsList) > 0:
             for index in range(self.nb_dataset):
-                self._unloadData(index)
                 for group_src in self.groups_src:
                     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")
 
-    def load(self):
+    def load(self, label: str):
         if self.use_cache:
             memory_init = getMemory()
 
-            indexs = [index for index in range(self.nb_dataset) if index not in self._index_cache]
+            indexs = [index for index in range(self.nb_dataset)]
             if len(indexs) > 0:
                 memory_lock = threading.Lock()
+                desc = lambda : "Caching "+ label +": {} | {} | {}".format(memoryInfo(), memoryForecast(memory_init, 0, self.nb_dataset), cpuInfo())
                 pbar = tqdm.tqdm(
                     total=len(indexs),
-                    desc="Caching : init | {} | {}".format(memoryForecast(memory_init, 0, self.nb_dataset), cpuInfo()),
-                    leave=False,
-                    disable=self.rank != 0 and "KONFAI_CLUSTER" not in os.environ
+                    desc=desc(),
+                    leave=False
                 )
 
                 def process(index):
                     self._loadData(index)
                     with memory_lock:
-                        pbar.set_description("Caching : {} | {} | {}".format(memoryInfo(), memoryForecast(memory_init, index, self.nb_dataset), cpuInfo()))
+                        pbar.set_description(desc())
                         pbar.update(1)
                 with ThreadPoolExecutor(max_workers=os.cpu_count()//(device_count() if device_count() > 0 else 1)) as executor:
                     futures = [executor.submit(process, index) for index in indexs]
@@ -170,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)
 
@@ -257,10 +271,10 @@ class Data(ABC):
                         groups_src : dict[str, Group],
                         patch : Union[DatasetPatch, None],
                         use_cache : bool,
-                        subset : Union[Subset, dict[str, Subset]],
+                        subset : Subset,
                         num_workers : int,
                         batch_size : int,
-                        validation: Union[float, str, list[int], list[str]] = 1,
+                        validation: Union[float, str, list[int], list[str], None] = None,
                         inlineAugmentations: bool = False,
                         dataAugmentationsList: dict[str, DataAugmentationsList]= {}) -> None:
         self.dataset_filenames = dataset_filenames
@@ -343,7 +357,7 @@ class Data(ABC):
                 append = flag == "a"
 
             if format not in SUPPORTED_EXTENSIONS:
-                    raise DatasetManagerError(f"Unsupported file format '{format}'.",
+                raise DatasetManagerError(f"Unsupported file format '{format}'.",
                         f"Supported extensions are: {', '.join(SUPPORTED_EXTENSIONS)}")
             
             dataset = Dataset(filename, format)
@@ -362,7 +376,7 @@ class Data(ABC):
                 raise DatasetManagerError(
                     f"Group source '{group_src}' not found in any dataset.",
                     f"Dataset filenames provided: {self.dataset_filenames}",
-                    f"Available groups across all datasets: {sorted(list(datasets.keys()))}",
+                    "Available groups across all datasets: {}".format(["{} {}".format(f, d.getGroup()) for f, d in self.datasets.items()]),
                     f"Please check that an entry in the dataset with the name '{group_src}.{format}' exists."
                 )
            
@@ -376,34 +390,44 @@ 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]]] = {}
         dataset_info : dict[str, dict[str, dict[str, Attribute]]] = {}
         for group in self.groups_src:
+            namesByGroup = set()
             if group not in dataset_name:
                 dataset_name[group] = {}
                 dataset_info[group] = {}
             for filename, _ in datasets[group]:
-                names.update(self.datasets[filename].getNames(group))
+                namesByGroup.update(self.datasets[filename].getNames(group))
                 dataset_name[group][filename] = self.datasets[filename].getNames(group)
                 dataset_info[group][filename] = {name: self.datasets[filename].getInfos(group, name) for name in dataset_name[group][filename]}
+            if len(names) == 0:
+                names.update(namesByGroup)
+            else: 
+                names = names.intersection(namesByGroup)
+        if len(names) == 0:
+           raise DatasetManagerError(
+                f"No data was found for groups {list(self.groups_src.keys())}: although each group contains data from a dataset, there are no common dataset names shared across all groups, the intersection is empty."
+            )   
+                    
         subset_names = set()
-
-        if isinstance(self.subset, dict):
-            for filename, subset in self.subset.items():
-                subset_names.update(subset([dataset_name[group][filename] for group in dataset_name], [dataset_info[group][filename] for group in dataset_name]))
-        else:
-             for group in dataset_name:
-                for filename, append in datasets[group]:
-                    if append:
-                        subset_names.update(self.subset([dataset_name[group][filename]], [dataset_info[group][filename]]))
+        for group in dataset_name:
+            subset_names_bygroup = set()
+            for filename, append in datasets[group]:
+                if append:
+                    subset_names_bygroup.update(self.subset([dataset_name[group][filename]], [dataset_info[group][filename]]))
+                else:
+                    if len(subset_names_bygroup) == 0:
+                        subset_names_bygroup.update(self.subset([dataset_name[group][filename]], [dataset_info[group][filename]])) 
                     else:
-                        if len(subset_names) == 0:
-                            subset_names.update(self.subset([dataset_name[group][filename]], [dataset_info[group][filename]])) 
-                        else:
-                            subset_names.intersection(self.subset([dataset_name[group][filename]], [dataset_info[group][filename]]))
+                        subset_names_bygroup = subset_names_bygroup.intersection(self.subset([dataset_name[group][filename]], [dataset_info[group][filename]]))
+            if len(subset_names) == 0:
+                subset_names.update(subset_names_bygroup)
+            else: 
+                subset_names = subset_names.intersection(subset_names_bygroup)
         if len(subset_names) == 0:
             raise DatasetManagerError("All data entries were excluded by the subset filter.",
                 f"Dataset entries found: {', '.join(names)}",
@@ -424,7 +448,7 @@ class Data(ABC):
         validate_map = []
         if isinstance(self.validation, float) or isinstance(self.validation, int):
             if self.validation <= 0 or self.validation >= 1:
-                raise DatasetManagerError("validation must be a float between 0 and 1.", f"→ Received: {self.validation}", "→ Example: validation = 0.2  # for a 20% validation split")
+                raise DatasetManagerError("Validation must be a float between 0 and 1.", f"Received: {self.validation}", "Example: validation = 0.2  # for a 20% validation split")
             
             train_map, validate_map = map[:int(math.floor(len(map)*(1-self.validation)))], map[int(math.floor(len(map)*(1-self.validation))):]
         elif isinstance(self.validation, str):
@@ -527,20 +551,19 @@ class DataPrediction(Data):
                         groups_src : dict[str, Group] = {"default" : Group()},
                         augmentations : Union[dict[str, DataAugmentationsList], None] = {"DataAugmentation_0" : DataAugmentationsList()},
                         patch : Union[DatasetPatch, None] = DatasetPatch(),
-                        use_cache : bool = True,
                         subset : Union[PredictionSubset, dict[str, PredictionSubset]] = PredictionSubset(),
                         num_workers : int = 4,
                         batch_size : int = 1) -> None:
 
-        super().__init__(dataset_filenames, groups_src, patch, use_cache, subset, num_workers, batch_size, dataAugmentationsList=augmentations if augmentations else {})
+        super().__init__(dataset_filenames, groups_src, patch, False, subset, num_workers, batch_size, dataAugmentationsList=augmentations if augmentations else {})
 
 class DataMetric(Data):
 
     @config("Dataset")
     def __init__(self,  dataset_filenames : list[str] = ["default:./Dataset"], 
-                        groups_src : dict[str, Group] = {"default" : Group()},
+                        groups_src : dict[str, GroupMetric] = {"default" : GroupMetric()},
                         subset : Union[PredictionSubset, dict[str, PredictionSubset]] = PredictionSubset(),
                         validation: Union[str, None] = None,
                         num_workers : int = 4) -> None:
 
-        super().__init__(dataset_filenames=dataset_filenames, groups_src=groups_src, patch=None, use_cache=False, subset=subset, num_workers=num_workers, batch_size=1, validation=1 if validation is None else validation)
+        super().__init__(dataset_filenames=dataset_filenames, groups_src=groups_src, patch=None, use_cache=False, subset=subset, num_workers=num_workers, batch_size=1, validation=validation)

konfai/data/patching.py

@@ -15,7 +15,6 @@ from konfai.utils.dataset import Dataset, Attribute
 from konfai.data.transform import Transform, Save
 from konfai.data.augmentation import DataAugmentationsList
 
-
 class PathCombine(ABC):
 
     def __init__(self) -> None:
@@ -240,6 +239,7 @@ class DatasetManager():
         self.index = index
         self.dataset = dataset
         self.loaded = False
+        self.augmentationLoaded = False
         self.cache_attributes: list[Attribute] = []
         _shape, cache_attribute =  self.dataset.getInfos(self.group_src, name)
         self.cache_attributes.append(cache_attribute)
@@ -258,6 +258,7 @@ class DatasetManager():
         self.cache_attributes_bak = copy.deepcopy(self.cache_attributes)
     
     def resetAugmentation(self):
+        self.cache_attributes[:] = self.cache_attributes[:1]
         i = 1
         for dataAugmentations in self.dataAugmentationsList:
             shape = []
@@ -272,15 +273,24 @@ class DatasetManager():
                 self.cache_attributes.append(caches_attribute[it])
                 self.patch.load(s, i)
                 i+=1
-
+    
     def load(self, pre_transform : list[Transform], dataAugmentationsList : list[DataAugmentationsList], device: torch.device) -> None:
-        if self.loaded:
-            return
+        if not self.loaded:
+            self._load(pre_transform)
+        if not self.augmentationLoaded:
+            self._loadAugmentation(dataAugmentationsList, device)
+            
+    def _load(self, pre_transform : list[Transform]):
+        self.cache_attributes = copy.deepcopy(self.cache_attributes_bak)
         i = len(pre_transform)
         data = None
         for transformFunction in reversed(pre_transform):
             if isinstance(transformFunction, Save):
-                filename, format = transformFunction.save.split(":")
+                if len(transformFunction.dataset.split(":")) > 1:
+                    filename, format = transformFunction.dataset.split(":")
+                else:
+                    filename = transformFunction.dataset.split(":")
+                    format = "mha"
                 dataset = Dataset(filename, format)
                 if dataset.isDatasetExist(self.group_dest, self.name):
                     data, attrib = dataset.readData(self.group_dest, self.name)
@@ -298,27 +308,41 @@ class DatasetManager():
             for transformFunction in pre_transform[i:]:
                 data = transformFunction(self.name, data, self.cache_attributes[0])
                 if isinstance(transformFunction, Save):
-                    filename, format = transformFunction.save.split(":")
+                    if len(transformFunction.dataset.split(":")) > 1:
+                        filename, format = transformFunction.dataset.split(":")
+                    else:
+                        filename = transformFunction.dataset.split(":")
+                        format = "mha"
                     dataset = Dataset(filename, format)
                     dataset.write(self.group_dest, self.name, data.numpy(), self.cache_attributes[0])
         self.data : list[torch.Tensor] = list()
         self.data.append(data)
 
+        for i in range(len(self.cache_attributes)-1):
+            self.cache_attributes[i+1].update(self.cache_attributes[0])
+        self.loaded = True
+    
+    def _loadAugmentation(self, dataAugmentationsList : list[DataAugmentationsList], device: torch.device) -> None:
         for dataAugmentations in dataAugmentationsList:
-            a_data = [data.clone() for _ in range(dataAugmentations.nb)]
+            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)
-        self.loaded = True
+        self.augmentationLoaded = True
 
     def unload(self) -> None:
-        if hasattr(self, "data"):
-            del self.data
-        self.cache_attributes = copy.deepcopy(self.cache_attributes_bak)
+        self.data.clear()
         self.loaded = False
-    
+        self.augmentationLoaded = False
+        
+    def unloadAugmentation(self) -> None:
+        self.data[:] = self.data[:1]
+        self.augmentationLoaded = False
+
+
     def getData(self, index : int, a : int, post_transforms : list[Transform], isInput: bool) -> torch.Tensor:
         data = self.patch.getData(self.data[a], index, a, isInput)
         for transformFunction in post_transforms: