konfai 1.1.0__py3-none-any.whl → 1.1.2__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

@@ -8,7 +8,7 @@ from typing import Union
 import os
 from konfai import KONFAI_ROOT
 from konfai.utils.config import config
-from konfai.utils.utils import _getModule
+from konfai.utils.utils import _getModule, AugmentationError
 from konfai.utils.dataset import Attribute, data_to_image
 
 
@@ -222,7 +222,7 @@ class ColorTransform(DataAugmentation):
                 matrix = matrix[:, :3, :].mean(dim=1, keepdims=True).to(input.device)
                 result = result.float() * matrix[:, :, :3].sum(dim=2, keepdims=True) + matrix[:, :, 3:]
             else:
-                raise ValueError('Image must be RGB (3 channels) or L (1 channel)')
+                raise AugmentationError('Image must be RGB (3 channels) or L (1 channel)')
             results.append(result.reshape(input.shape))
         return results
     

konfai/data/data_manager.py

@@ -11,11 +11,12 @@ from typing import Union, Iterator
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import threading
 from torch.cuda import device_count
+import SimpleITK as sitk
 
 from konfai import KONFAI_STATE, KONFAI_ROOT
 from konfai.data.patching import DatasetPatch, DatasetManager
 from konfai.utils.config import config
-from konfai.utils.utils import memoryInfo, cpuInfo, memoryForecast, getMemory, State
+from konfai.utils.utils import memoryInfo, cpuInfo, memoryForecast, getMemory, State, SUPPORTED_EXTENSIONS, DatasetManagerError
 from konfai.utils.dataset import Dataset, Attribute
 from konfai.data.transform import TransformLoader, Transform
 from konfai.data.augmentation import DataAugmentationsList
@@ -61,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:
@@ -108,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]
@@ -236,6 +251,8 @@ class Subset():
             index = random.sample(index, len(index))
         return set([names_filtred[i] for i in index])
     
+    def __str__(self):
+        return "Subset : " + str(self.subset) + " shuffle : "+ str(self.shuffle) + " filter : "+ str(self.filter)
 class TrainSubset(Subset):
 
     @config()
@@ -254,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
@@ -290,6 +307,13 @@ class Data(ABC):
                     map.append((x, y, z))
         return data, map
 
+    def getGroupsDest(self):
+        groupsDest = []
+        for group_src in self.groups_src:
+            for group_dest in self.groups_src[group_src]:
+                groupsDest.append(group_dest)
+        return groupsDest
+    
     def _split(map: list[tuple[int, int, int]], world_size: int) -> list[list[tuple[int, int, int]]]:
         if len(map) == 0:
             return [[] for _ in range(world_size)]
@@ -313,7 +337,14 @@ class Data(ABC):
     
     def getData(self, world_size: int) -> list[list[DataLoader]]:
         datasets: dict[str, list[(str, bool)]] = {}
+        if self.dataset_filenames is None or len(self.dataset_filenames) == 0:
+            raise DatasetManagerError("No dataset filenames were provided")
         for dataset_filename in self.dataset_filenames:
+            if dataset_filename is None:
+                raise DatasetManagerError("Invalid dataset entry: 'None' received.",
+                    "Each dataset must be a valid path string (e.g., './Dataset/', './Dataset/:mha, './Dataset/:a:mha', './Dataset/:i:mha').",
+                    "Please check your 'dataset_filenames' list for missing or null entries."
+                )
             if len(dataset_filename.split(":")) == 1:
                 filename = dataset_filename
                 format = "mha"
@@ -324,9 +355,13 @@ class Data(ABC):
             else:
                 filename, flag, format = dataset_filename.split(":")
                 append = flag == "a"
-                
-            dataset = Dataset(filename, format) 
+
+            if format not in SUPPORTED_EXTENSIONS:
+                raise DatasetManagerError(f"Unsupported file format '{format}'.",
+                        f"Supported extensions are: {', '.join(SUPPORTED_EXTENSIONS)}")
             
+            dataset = Dataset(filename, format)
+
             self.datasets[filename] = dataset
             for group in self.groups_src:
                 if dataset.isGroupExist(group):
@@ -334,47 +369,88 @@ class Data(ABC):
                         datasets[group].append((filename, append)) 
                     else:
                         datasets[group] = [(filename, append)]
+        modelHaveInput = False
         for group_src in self.groups_src:
             if group_src not in datasets:
-                raise ValueError("[DatasetManager] Error: group source {} not found. Available groups: {}".format(group_src, list(datasets.keys())))
                 
+                raise DatasetManagerError(
+                    f"Group source '{group_src}' not found in any dataset.",
+                    f"Dataset filenames provided: {self.dataset_filenames}",
+                    "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."
+                )
+           
             for group_dest in self.groups_src[group_src]:
                 self.groups_src[group_src][group_dest].load(group_src, group_dest, [self.datasets[filename] for filename, _ in datasets[group_src]])
+                modelHaveInput |= self.groups_src[group_src][group_dest].isInput
+
+        if not modelHaveInput:
+            raise DatasetManagerError(
+                "At least one group must be defined with 'isInput: true' to provide input to the network."
+            )
+
         for key, dataAugmentations in self.dataAugmentationsList.items():
             dataAugmentations.load(key)
 
-        names : list[list[str]] = []
+        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.append(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)}",
+                f"Subset object applied: {self.subset}",
+                f"Subset requested : {', '.join(subset_names)}",
+                "None of the dataset entries matched the given subset.",
+                "Please check your 'subset' configuration — it may be too restrictive or incorrectly formatted.",
+                "Examples of valid subset formats:",
+                "\tsubset: [0, 1]            # explicit indices",
+                "\tsubset: 0:10              # slice notation",
+                "\tsubset: ./Validation.txt  # external file",
+                "\tsubset: None              # to disable filtering"
+            )
+                                        
         data, map = self._getDatasets(list(subset_names), dataset_name)
-    
+
         train_map = map
         validate_map = []
-        if isinstance(self.validation, float):
-            if self.validation < 1.0 and int(math.floor(len(map)*(1-self.validation))) > 0:
-                train_map, validate_map = map[:int(math.floor(len(map)*self.validation))], map[int(math.floor(len(map)*self.validation)):]
+        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")
+            
+            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):
             if ":" in self.validation:
                 index = list(range(int(self.subset.split(":")[0]), int(self.subset.split(":")[1])))
@@ -389,7 +465,17 @@ class Data(ABC):
                 train_map = [m for m in map if m[0] not in index]
                 validate_map = [m for m in map if m[0] in index]
             else:
-                validate_map = train_map
+                raise DatasetManagerError(
+                    f"Invalid string value for 'validation': '{self.validation}'",
+                    "Expected one of the following formats:",
+                    "\t• A slice string like '0:10'",
+                    "\t• A path to a text file listing validation sample names (e.g., './val.txt')",
+                    "\t• A float between 0 and 1 (e.g., 0.2)",
+                    "\t• A list of sample names or indices",
+                    "The provided value is neither a valid slice nor a readable file.",
+                    "Please fix your 'validation' setting in the configuration."
+                    )
+            
         elif isinstance(self.validation, list):
             if len(self.validation) > 0:
                 if isinstance(self.validation[0], int):
@@ -399,10 +485,29 @@ class Data(ABC):
                     index = [i for i, n in enumerate(subset_names) if n in self.validation]
                     train_map = [m for m in map if m[0] not in index]
                     validate_map = [m for m in map if m[0] in index]
-
+                else:
+                    raise DatasetManagerError(f"Invalid list type for 'validation': elements of type '{type(self.validation[0]).__name__}' are not supported.",
+                            "Supported list element types are:",
+                            "\t• int  → list of indices (e.g., [0, 1, 2])",
+                            "\t• str  → list of sample names (e.g., ['patient01', 'patient02'])",
+                            f"Received list: {self.validation}"
+                        )
+        if len(train_map) == 0:
+            raise DatasetManagerError("No data left for training after applying the validation split.",
+                f"Dataset size: {len(map)}",
+                f"Validation setting: {self.validation}",
+                "Please reduce the validation size, increase the dataset, or disable validation."
+            )
+
+        if self.validation is not None and len(validate_map) == 0:
+            raise DatasetManagerError("No data left for validation after applying the validation split.",
+                f"Dataset size: {len(map)}",
+                f"Validation setting: {self.validation}",
+                "Please increase the validation size, increase the dataset, or disable validation."
+            )
         train_maps = Data._split(train_map, world_size)
         validate_maps = Data._split(validate_map, world_size)
-
+        
         for i, (train_map, validate_map) in enumerate(zip(train_maps, validate_maps)):
             maps = [train_map]
             if len(validate_map):
@@ -436,7 +541,7 @@ class DataTrain(Data):
                         subset : Union[TrainSubset, dict[str, TrainSubset]] = TrainSubset(),
                         num_workers : int = 4,
                         batch_size : int = 1,
-                        validation : Union[float, str, list[int], list[str]] = 0.8) -> None:
+                        validation : Union[float, str, list[int], list[str]] = 0.2) -> None:
         super().__init__(dataset_filenames, groups_src, patch, use_cache, subset, num_workers, batch_size, validation, inlineAugmentations, augmentations if augmentations else {})
 
 class DataPrediction(Data):
@@ -445,22 +550,20 @@ class DataPrediction(Data):
     def __init__(self,  dataset_filenames : list[str] = ["default:./Dataset"], 
                         groups_src : dict[str, Group] = {"default" : Group()},
                         augmentations : Union[dict[str, DataAugmentationsList], None] = {"DataAugmentation_0" : DataAugmentationsList()},
-                        inlineAugmentations: bool = False,
                         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, inlineAugmentations=inlineAugmentations, 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:
@@ -150,7 +149,10 @@ class Patch(ABC):
 
     def __init__(self, patch_size: list[int], overlap: Union[int, None], path_mask: Union[str, None] = None, padValue: float = 0, extend_slice: int = 0) -> None:
         self.patch_size = patch_size
-        self.overlap = overlap
+        self.overlap = overlap 
+        if isinstance(self.overlap, int):
+             if self.overlap < 0:
+                 self.overlap = None
         self._patch_slices : dict[int, list[tuple[slice]]] = {}
         self._nb_patch_per_dim: dict[int, list[tuple[int, bool]]] = {}
         self.path_mask = path_mask
@@ -237,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)
@@ -255,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 = []
@@ -269,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)
@@ -295,27 +308,40 @@ 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)
             
             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:

konfai/data/transform.py

@@ -6,7 +6,7 @@ from abc import ABC, abstractmethod
 import torch.nn.functional as F
 from typing import Any, Union
 
-from konfai.utils.utils import _getModule, NeedDevice, _resample_affine, _affine_matrix
+from konfai.utils.utils import _getModule, NeedDevice, _resample_affine, _affine_matrix, TransformError
 from konfai.utils.dataset import Dataset, Attribute, data_to_image, image_to_data
 from konfai.utils.config import config
 
@@ -211,39 +211,67 @@ class Resample(Transform, ABC):
         _ = cache_attribute.pop_np_array("Spacing")
         return self._resample(input, [int(size) for size in size_1])
 
-class ResampleIsotropic(Resample):
+class ResampleToResolution(Resample):
+
+    def __init__(self, spacing : list[Union[float, None]] = [1., 1., 1.]) -> None:
+        self.spacing = torch.tensor([0 if s < 0 else s for s in spacing])
 
-    def __init__(self, spacing : list[float] = [1., 1., 1.]) -> None:
-        self.spacing = torch.tensor(spacing, dtype=torch.float64)
-        
     def transformShape(self, shape: list[int], cache_attribute: Attribute) -> list[int]:
-        assert "Spacing" in cache_attribute, "Error no spacing"
-        resize_factor = self.spacing/cache_attribute.get_tensor("Spacing").flip(0)
-        return  [int(x) for x in (torch.tensor(shape) * 1/resize_factor)]
+        if "Spacing" not in cache_attribute:
+            TransformError("Missing 'Spacing' in cache attributes, the data is likely not a valid image.",
+                        "Make sure your input is a image (e.g., .nii, .mha) with proper metadata.")
+        if len(shape) != len(self.spacing):
+            TransformError("Shape and spacing dimensions do not match: shape={shape}, spacing={self.spacing}")
+        image_spacing = cache_attribute.get_tensor("Spacing").flip(0)
+        spacing = self.spacing
+        
+        for i, s in enumerate(self.spacing):
+            if s == 0:
+                spacing[i] = image_spacing[i]
+        resize_factor = spacing/cache_attribute.get_tensor("Spacing").flip(0)
+        return [int(x) for x in (torch.tensor(shape) * 1/resize_factor)]
 
     def __call__(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        assert "Spacing" in cache_attribute, "Error no spacing"
-        resize_factor = self.spacing/cache_attribute.get_tensor("Spacing").flip(0)
-        cache_attribute["Spacing"] = self.spacing.flip(0)
+        image_spacing = cache_attribute.get_tensor("Spacing").flip(0)
+        spacing = self.spacing
+        for i, s in enumerate(self.spacing):
+            if s == 0:
+                spacing[i] = image_spacing[i]
+        resize_factor = spacing/cache_attribute.get_tensor("Spacing").flip(0)
+        cache_attribute["Spacing"] = spacing.flip(0)
         cache_attribute["Size"] = np.asarray([int(x) for x in torch.tensor(input.shape[1:])])
         size = [int(x) for x in (torch.tensor(input.shape[1:]) * 1/resize_factor)]
         cache_attribute["Size"] = np.asarray(size)
         return self._resample(input, size)
 
-class ResampleResize(Resample):
+class ResampleToSize(Resample):
 
     def __init__(self, size : list[int] = [100,512,512]) -> None:
         self.size = size
 
     def transformShape(self, shape: list[int], cache_attribute: Attribute) -> list[int]:
-        return self.size
+        if "Spacing" not in cache_attribute:
+            TransformError("Missing 'Spacing' in cache attributes, the data is likely not a valid image.",
+                        "Make sure your input is a image (e.g., .nii, .mha) with proper metadata.")
+        if len(shape) != len(self.size):
+            TransformError("Shape and spacing dimensions do not match: shape={shape}, spacing={self.spacing}")
+        size = self.size
+        for i, s in enumerate(self.size):
+            if s == -1:
+                size[i] = shape[i]
+        return size
     
     def __call__(self, name: str, input: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
+        size = self.size
+        image_size =  np.asarray([int(x) for x in torch.tensor(input.shape[1:])])
+        for i, s in enumerate(self.size):
+            if s is None:
+                size[i] = image_size[i]
         if "Spacing" in cache_attribute:
-            cache_attribute["Spacing"] = torch.flip(torch.tensor(list(input.shape[1:]))/torch.tensor(self.size)*torch.flip(cache_attribute.get_tensor("Spacing"), dims=[0]), dims=[0])
-        cache_attribute["Size"] = np.asarray([int(x) for x in torch.tensor(input.shape[1:])])
-        cache_attribute["Size"] = self.size
-        return self._resample(input, self.size)
+            cache_attribute["Spacing"] = torch.flip(torch.tensor(image_size)/torch.tensor(size)*torch.flip(cache_attribute.get_tensor("Spacing"), dims=[0]), dims=[0])
+        cache_attribute["Size"] = image_size
+        cache_attribute["Size"] = size
+        return self._resample(input, size)
 
 class ResampleTransform(Transform):
 
@@ -412,8 +440,8 @@ class FlatLabel(Transform):
 
 class Save(Transform):
 
-    def __init__(self, save: str) -> None:
-        self.save = save
+    def __init__(self, dataset: str) -> None:
+        self.dataset = dataset
     
     def __call__(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
         return input
@@ -528,8 +556,7 @@ class OneHot(Transform):
         self.num_classes = num_classes
 
     def __call__(self, name: str, input: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        result =  F.one_hot(input.type(torch.int64), num_classes=self.num_classes).permute(0, len(input.shape), *[i+1 for i in range(len(input.shape)-1)]).float().squeeze(2)
-        print(result.shape)
+        result =  F.one_hot(input.type(torch.int64), num_classes=self.num_classes).permute(0, len(input.shape), *[i+1 for i in range(len(input.shape)-1)]).float().squeeze(0)
         return result
     
     def inverse(self, name: str, input: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: