konfai 1.0.8__py3-none-any.whl → 1.1.0__py3-none-any.whl

konfai/__init__.py

@@ -1,16 +1,16 @@
 import os
 import datetime
 
-MODELS_DIRECTORY = lambda : os.environ["DL_API_MODELS_DIRECTORY"]
-CHECKPOINTS_DIRECTORY =lambda : os.environ["DL_API_CHECKPOINTS_DIRECTORY"]
-MODEL = lambda : os.environ["DL_API_MODEL"]
-PREDICTIONS_DIRECTORY =lambda : os.environ["DL_API_PREDICTIONS_DIRECTORY"]
-EVALUATIONS_DIRECTORY =lambda : os.environ["DL_API_EVALUATIONS_DIRECTORY"]
-STATISTICS_DIRECTORY = lambda : os.environ["DL_API_STATISTICS_DIRECTORY"]
-SETUPS_DIRECTORY = lambda : os.environ["DL_API_SETUPS_DIRECTORY"]
-CONFIG_FILE = lambda : os.environ["DEEP_LEARNING_API_CONFIG_FILE"]
-DL_API_STATE = lambda : os.environ["DL_API_STATE"]
-DEEP_LEARNING_API_ROOT = lambda : os.environ["DEEP_LEARNING_API_ROOT"]
+MODELS_DIRECTORY = lambda : os.environ["KONFAI_MODELS_DIRECTORY"]
+CHECKPOINTS_DIRECTORY =lambda : os.environ["KONFAI_CHECKPOINTS_DIRECTORY"]
+MODEL = lambda : os.environ["KONFAI_MODEL"]
+PREDICTIONS_DIRECTORY =lambda : os.environ["KONFAI_PREDICTIONS_DIRECTORY"]
+EVALUATIONS_DIRECTORY =lambda : os.environ["KONFAI_EVALUATIONS_DIRECTORY"]
+STATISTICS_DIRECTORY = lambda : os.environ["KONFAI_STATISTICS_DIRECTORY"]
+SETUPS_DIRECTORY = lambda : os.environ["KONFAI_SETUPS_DIRECTORY"]
+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"]
 
 DATE = lambda : datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")

konfai/data/augmentation.py

@@ -6,7 +6,7 @@ import SimpleITK as sitk
 import torch.nn.functional as F
 from typing import Union
 import os
-from konfai import DEEP_LEARNING_API_ROOT
+from konfai import KONFAI_ROOT
 from konfai.utils.config import config
 from konfai.utils.utils import _getModule
 from konfai.utils.dataset import Attribute, data_to_image
@@ -51,7 +51,7 @@ class Prob():
 class DataAugmentationsList():
 
     @config()
-    def __init__(self, nb : int = 10, dataAugmentations: dict[str, Prob] = {"default:RandomElastixTransform" : Prob(1)}) -> None:
+    def __init__(self, nb : int = 10, dataAugmentations: dict[str, Prob] = {"default:Flip" : Prob(1)}) -> None:
         self.nb = nb
         self.dataAugmentations : list[DataAugmentation] = []
         self.dataAugmentationsLoader = dataAugmentations
@@ -59,7 +59,7 @@ class DataAugmentationsList():
     def load(self, key: str):
         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(DEEP_LEARNING_API_ROOT(), key))
+            dataAugmentation: DataAugmentation = getattr(importlib.import_module(module), name)(config = None, DL_args="{}.Dataset.augmentations.{}.dataAugmentations".format(KONFAI_ROOT(), key))
             dataAugmentation.load(prob.prob)
             self.dataAugmentations.append(dataAugmentation)
     
@@ -525,7 +525,7 @@ class Elastix(DataAugmentation):
 class Permute(DataAugmentation):
 
     @config("Permute")
-    def __init__(self, prob_permute: Union[list[float], None] = [0.33 ,0.33]) -> None:
+    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]])
         self.prob_permute = prob_permute

konfai/data/{dataset.py → data_manager.py}

@@ -12,8 +12,8 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 import threading
 from torch.cuda import device_count
 
-from konfai import DL_API_STATE, DEEP_LEARNING_API_ROOT
-from konfai.data.HDF5 import DatasetPatch, DatasetManager
+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.dataset import Dataset, Attribute
@@ -36,7 +36,7 @@ class GroupTransform:
         if self._pre_transforms is not None:
             if isinstance(self._pre_transforms, dict):
                 for classpath, transform in self._pre_transforms.items():
-                    transform = transform.getTransform(classpath, DL_args =  "{}.Dataset.groups_src.{}.groups_dest.{}.pre_transforms".format(DEEP_LEARNING_API_ROOT(), group_src, group_dest))
+                    transform = transform.getTransform(classpath, DL_args =  "{}.Dataset.groups_src.{}.groups_dest.{}.pre_transforms".format(KONFAI_ROOT(), group_src, group_dest))
                     transform.setDatasets(datasets)
                     self.pre_transforms.append(transform)
             else:
@@ -47,7 +47,7 @@ class GroupTransform:
         if self._post_transforms is not None:
             if isinstance(self._post_transforms, dict):
                 for classpath, transform in self._post_transforms.items():
-                    transform = transform.getTransform(classpath, DL_args = "{}.Dataset.groups_src.{}.groups_dest.{}.post_transforms".format(DEEP_LEARNING_API_ROOT(), group_src, group_dest))
+                    transform = transform.getTransform(classpath, DL_args = "{}.Dataset.groups_src.{}.groups_dest.{}.post_transforms".format(KONFAI_ROOT(), group_src, group_dest))
                     transform.setDatasets(datasets)
                     self.post_transforms.append(transform)
             else:
@@ -64,7 +64,7 @@ class GroupTransform:
 class Group(dict[str, GroupTransform]):
 
     @config()
-    def __init__(self, groups_dest: dict[str, GroupTransform] = {"default": GroupTransform()}):
+    def __init__(self, groups_dest: dict[str, GroupTransform] = {"default:group_dest": GroupTransform()}):
         super().__init__(groups_dest)
 
 class CustomSampler(Sampler[int]):
@@ -127,7 +127,7 @@ class DatasetIter(data.Dataset):
                     total=len(indexs),
                     desc="Caching : init | {} | {}".format(memoryForecast(memory_init, 0, self.nb_dataset), cpuInfo()),
                     leave=False,
-                    disable=self.rank != 0 and "DL_API_CLUSTER" not in os.environ
+                    disable=self.rank != 0 and "KONFAI_CLUSTER" not in os.environ
                 )
 
                 def process(index):
@@ -135,7 +135,7 @@ class DatasetIter(data.Dataset):
                     with memory_lock:
                         pbar.set_description("Caching : {} | {} | {}".format(memoryInfo(), memoryForecast(memory_init, index, self.nb_dataset), cpuInfo()))
                         pbar.update(1)
-                with ThreadPoolExecutor(max_workers=os.cpu_count()//device_count()) as executor:
+                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]
                     for _ in as_completed(futures):
                         pass
@@ -257,14 +257,14 @@ class Data(ABC):
                         subset : Union[Subset, dict[str, Subset]],
                         num_workers : int,
                         batch_size : int,
-                        train_size: Union[float, str, list[int], list[str]] = 1,
+                        validation: Union[float, str, list[int], list[str]] = 1,
                         inlineAugmentations: bool = False,
                         dataAugmentationsList: dict[str, DataAugmentationsList]= {}) -> None:
         self.dataset_filenames = dataset_filenames
         self.subset = subset
         self.groups_src = groups_src
         self.patch = patch
-        self.train_size = train_size
+        self.validation = validation
         self.dataAugmentationsList = dataAugmentationsList
         self.batch_size = batch_size
         self.dataSet_args = dict(groups_src=self.groups_src, inlineAugmentations=inlineAugmentations, dataAugmentationsList = list(self.dataAugmentationsList.values()), use_cache = use_cache, buffer_size=batch_size+1, patch_size=self.patch.patch_size if self.patch is not None else None, overlap=self.patch.overlap if self.patch is not None else None)
@@ -295,7 +295,7 @@ class Data(ABC):
             return [[] for _ in range(world_size)]
         
         maps = []
-        if DL_API_STATE() == str(State.PREDICTION) or DL_API_STATE() == str(State.EVALUATION):
+        if KONFAI_STATE() == str(State.PREDICTION) or KONFAI_STATE() == str(State.EVALUATION):
             np_map = np.asarray(map)
             unique_index = np.unique(np_map[:, 0])
             offset = int(np.ceil(len(unique_index)/world_size))
@@ -314,7 +314,11 @@ class Data(ABC):
     def getData(self, world_size: int) -> list[list[DataLoader]]:
         datasets: dict[str, list[(str, bool)]] = {}
         for dataset_filename in self.dataset_filenames:
-            if len(dataset_filename.split(":")) == 2:
+            if len(dataset_filename.split(":")) == 1:
+                filename = dataset_filename
+                format = "mha"
+                append = True
+            elif len(dataset_filename.split(":")) == 2:
                 filename, format = dataset_filename.split(":")
                 append = True
             else:
@@ -331,8 +335,9 @@ class Data(ABC):
                     else:
                         datasets[group] = [(filename, append)]
         for group_src in self.groups_src:
-            assert group_src in datasets, "Error group source {} not found".format(group_src)
-
+            if group_src not in datasets:
+                raise ValueError("[DatasetManager] Error: group source {} not found. Available groups: {}".format(group_src, list(datasets.keys())))
+                
             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]])
         for key, dataAugmentations in self.dataAugmentationsList.items():
@@ -367,17 +372,17 @@ class Data(ABC):
     
         train_map = map
         validate_map = []
-        if isinstance(self.train_size, float):
-            if self.train_size < 1.0 and int(math.floor(len(map)*(1-self.train_size))) > 0:
-                train_map, validate_map = map[:int(math.floor(len(map)*self.train_size))], map[int(math.floor(len(map)*self.train_size)):]
-        elif isinstance(self.train_size, str):
-            if ":" in self.train_size:
+        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)):]
+        elif isinstance(self.validation, str):
+            if ":" in self.validation:
                 index = list(range(int(self.subset.split(":")[0]), int(self.subset.split(":")[1])))
                 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]
-            elif os.path.exists(self.train_size):
+            elif os.path.exists(self.validation):
                 validation_names = []
-                with open(self.train_size, "r") as f:
+                with open(self.validation, "r") as f:
                     for name in f:
                         validation_names.append(name.strip())
                 index = [i for i, n in enumerate(subset_names) if n in validation_names]
@@ -385,13 +390,13 @@ class Data(ABC):
                 validate_map = [m for m in map if m[0] in index]
             else:
                 validate_map = train_map
-        elif isinstance(self.train_size, list):
-            if len(self.train_size) > 0:
-                if isinstance(self.train_size[0], int):
-                    train_map = [m for m in map if m[0] not in self.train_size]
-                    validate_map = [m for m in map if m[0] in self.train_size]
-                elif isinstance(self.train_size[0], str):
-                    index = [i for i, n in enumerate(subset_names) if n in self.train_size]
+        elif isinstance(self.validation, list):
+            if len(self.validation) > 0:
+                if isinstance(self.validation[0], int):
+                    train_map = [m for m in map if m[0] not in self.validation]
+                    validate_map = [m for m in map if m[0] in self.validation]
+                elif isinstance(self.validation[0], str):
+                    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]
 
@@ -422,8 +427,8 @@ class Data(ABC):
 class DataTrain(Data):
 
     @config("Dataset")
-    def __init__(self,  dataset_filenames : list[str] = ["default:Dataset.h5"], 
-                        groups_src : dict[str, Group] = {"default" : Group()},
+    def __init__(self,  dataset_filenames : list[str] = ["default:./Dataset"], 
+                        groups_src : dict[str, Group] = {"default:group_src" : Group()},
                         augmentations : Union[dict[str, DataAugmentationsList], None] = {"DataAugmentation_0" : DataAugmentationsList()},
                         inlineAugmentations: bool = False,
                         patch : Union[DatasetPatch, None] = DatasetPatch(),
@@ -431,13 +436,13 @@ class DataTrain(Data):
                         subset : Union[TrainSubset, dict[str, TrainSubset]] = TrainSubset(),
                         num_workers : int = 4,
                         batch_size : int = 1,
-                        train_size : Union[float, str, list[int], list[str]] = 0.8) -> None:
-        super().__init__(dataset_filenames, groups_src, patch, use_cache, subset, num_workers, batch_size, train_size, inlineAugmentations, augmentations if augmentations else {})
+                        validation : Union[float, str, list[int], list[str]] = 0.8) -> None:
+        super().__init__(dataset_filenames, groups_src, patch, use_cache, subset, num_workers, batch_size, validation, inlineAugmentations, augmentations if augmentations else {})
 
 class DataPrediction(Data):
 
     @config("Dataset")
-    def __init__(self,  dataset_filenames : list[str] = ["default:Dataset.h5"], 
+    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,
@@ -452,19 +457,10 @@ class DataPrediction(Data):
 class DataMetric(Data):
 
     @config("Dataset")
-    def __init__(self,  dataset_filenames : list[str] = ["default:Dataset.h5"], 
+    def __init__(self,  dataset_filenames : list[str] = ["default:./Dataset"], 
                         groups_src : dict[str, Group] = {"default" : Group()},
                         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, train_size=1 if validation is None else validation)
-
-class DataHyperparameter(Data):
-
-    @config("Dataset")
-    def __init__(self,  dataset_filenames : list[str] = ["default:Dataset.h5"], 
-                        groups_src : dict[str, Group] = {"default" : Group()},
-                        patch : Union[DatasetPatch, None] = DatasetPatch()) -> None:
-
-        super().__init__(dataset_filenames, groups_src, patch, False, PredictionSubset(), 0, False, 1)
+        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)

konfai/data/{HDF5.py → patching.py}

@@ -37,13 +37,13 @@ class PathCombine(ABC):
     3D : 
         AAA+AAB+ABA+ABB+BAA+BAB+BBA+BBB
         
-        AAC+ABC+BAC+BBC
-        ACA+ACB+BCA+BCB
         CAA+CAB+CBA+CBB
+        ACA+ACB+BCA+BCB
+        AAC+ABC+BAC+BBC
 
-        ACC+BCC
-        CAC+CBC
         CCA+CCB
+        CAC+CBC
+        ACC+BCC
     
     """
     def setPatchConfig(self, patch_size: list[int], overlap: int):
@@ -188,7 +188,7 @@ class Patch(ABC):
             pad_bottom = 0
             pad_top = 0
             if self._patch_slices[a][index][0].start-bottom < 0:
-                pad_bottom = bottom-s.start
+                pad_bottom = bottom-self._patch_slices[a][index][0].start
             if self._patch_slices[a][index][0].stop+top > data.shape[len(slices_pre)]:
                 pad_top = self._patch_slices[a][index][0].stop+top-data.shape[len(slices_pre)]
             data_sliced = F.pad(data_sliced, [0 for _ in range((len(slices)-1)*2)]+[pad_bottom, pad_top], 'reflect')
@@ -214,13 +214,13 @@ class Patch(ABC):
 class DatasetPatch(Patch):
 
     @config("Patch")
-    def __init__(self, patch_size : list[int] = [128, 256, 256], overlap : Union[int, None] = None, mask: Union[str, None] = None, padValue: float = 0, extend_slice: int = 0) -> None:
+    def __init__(self, patch_size : list[int] = [128, 128, 128], overlap : Union[int, None] = None, mask: Union[str, None] = None, padValue: float = 0, extend_slice: int = 0) -> None:
         super().__init__(patch_size, overlap, mask, padValue, extend_slice)
 
 class ModelPatch(Patch):
 
     @config("Patch")
-    def __init__(self, patch_size : list[int] = [128, 256, 256], overlap : Union[int, None] = None, patchCombine: Union[str, None] = None, mask: Union[str, None] = None, padValue: float = 0, extend_slice: int = 0) -> None:
+    def __init__(self, patch_size : list[int] = [128, 128, 128], overlap : Union[int, None] = None, patchCombine: Union[str, None] = None, mask: Union[str, None] = None, padValue: float = 0, extend_slice: int = 0) -> None:
         super().__init__(patch_size, overlap, mask, padValue, extend_slice)
         self.patchCombine = patchCombine
 

konfai/data/transform.py

@@ -130,7 +130,7 @@ class Standardize(Transform):
         
 class TensorCast(Transform):
 
-    def __init__(self, dtype : str = "default:float32,int64,int16") -> None:
+    def __init__(self, dtype : str = "float32") -> None:
         self.dtype : torch.dtype = getattr(torch, dtype)
 
     def __call__(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
@@ -142,7 +142,7 @@ class TensorCast(Transform):
 
 class Padding(Transform):
 
-    def __init__(self, padding : list[int] = [0,0,0,0,0,0], mode : str = "default:constant,reflect,replicate,circular") -> None:
+    def __init__(self, padding : list[int] = [0,0,0,0,0,0], mode : str = "constant") -> None:
         self.padding = padding
         self.mode = mode
 
@@ -332,7 +332,7 @@ class ResampleTransform(Transform):
 
 class Mask(Transform):
 
-    def __init__(self, path : str = "default:./default.mha", value_outside: int = 0) -> None:
+    def __init__(self, path : str = "./default.mha", value_outside: int = 0) -> None:
         self.path = path
         self.value_outside = value_outside
         

konfai/evaluator.py

@@ -7,10 +7,10 @@ import json
 import shutil
 import builtins
 import importlib
-from konfai import EVALUATIONS_DIRECTORY, PREDICTIONS_DIRECTORY, DEEP_LEARNING_API_ROOT, CONFIG_FILE
+from konfai import EVALUATIONS_DIRECTORY, PREDICTIONS_DIRECTORY, KONFAI_ROOT, CONFIG_FILE
 from konfai.utils.config import config
 from konfai.utils.utils import _getModule, DistributedObject, synchronize_data
-from konfai.data.dataset import DataMetric
+from konfai.data.data_manager import DataMetric
 
 class CriterionsAttr():
 
@@ -28,7 +28,7 @@ class CriterionsLoader():
         criterions = {}
         for module_classpath, criterionsAttr in self.criterionsLoader.items():
             module, name = _getModule(module_classpath, "metric.measure")
-            criterions[config("{}.metrics.{}.targetsCriterions.{}.criterionsLoader.{}".format(DEEP_LEARNING_API_ROOT(), output_group, target_group, module_classpath))(getattr(importlib.import_module(module), name))(config = None)] = criterionsAttr
+            criterions[config("{}.metrics.{}.targetsCriterions.{}.criterionsLoader.{}".format(KONFAI_ROOT(), output_group, target_group, module_classpath))(getattr(importlib.import_module(module), name))(config = None)] = criterionsAttr
         return criterions
 
 class TargetCriterionsLoader():
@@ -86,8 +86,8 @@ class Statistics():
 class Evaluator(DistributedObject):
 
     @config("Evaluator")
-    def __init__(self, train_name: str = "default:name", metrics: dict[str, TargetCriterionsLoader] = {"default": TargetCriterionsLoader()}, dataset : DataMetric = DataMetric(),) -> None:
-        if os.environ["DEEP_LEANING_API_CONFIG_MODE"] != "Done":
+    def __init__(self, train_name: str = "default:TRAIN_01", metrics: dict[str, TargetCriterionsLoader] = {"default": TargetCriterionsLoader()}, dataset : DataMetric = DataMetric(),) -> None:
+        if os.environ["KONFAI_CONFIG_MODE"] != "Done":
             exit(0)
         super().__init__(train_name)
         self.metric_path = EVALUATIONS_DIRECTORY()+self.name+"/"
@@ -111,7 +111,7 @@ class Evaluator(DistributedObject):
     
     def setup(self, world_size: int):
         if os.path.exists(self.metric_path):
-            if os.environ["DL_API_OVERWRITE"] != "True":
+            if os.environ["KONFAI_OVERWRITE"] != "True":
                 accept = builtins.input("The metric {} already exists ! Do you want to overwrite it (yes,no) : ".format(self.name))
                 if accept != "yes":
                     return

konfai/main.py

@@ -2,21 +2,25 @@ import argparse
 import os
 from torch.cuda import device_count
 import torch.multiprocessing as mp
-from konfai.utils.utils import setupAPI, TensorBoard, Log
+from konfai.utils.utils import setup, TensorBoard, Log
 
 import sys
 sys.path.insert(0, os.getcwd())
 
 def main():
     parser = argparse.ArgumentParser(description="KonfAI", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    with setupAPI(parser) as distributedObject:
-        with Log(distributedObject.name):
-            world_size = device_count()
-            if world_size == 0:
-                world_size = 1
-            distributedObject.setup(world_size)
-            with TensorBoard(distributedObject.name):
-                mp.spawn(distributedObject, nprocs=world_size)
+    try:
+        with setup(parser) as distributedObject:
+            with Log(distributedObject.name):
+                world_size = device_count()
+                if world_size == 0:
+                    world_size = 1
+                distributedObject.setup(world_size)
+                with TensorBoard(distributedObject.name):
+                    mp.spawn(distributedObject, nprocs=world_size)
+    except Exception as e:
+        print(e)
+        exit(1)
 
 
 def cluster():
@@ -29,17 +33,20 @@ def cluster():
     cluster_args.add_argument('--memory', type=int, default=16, help='Amount of memory per node')
     cluster_args.add_argument('--time-limit', '--time_limit', type=int, default=1440, help='Job time limit in minute')
     cluster_args.add_argument('--resubmit', action='store_true', help='Automatically resubmit job just before timout')
+    try:
+        with setup(parser) as distributedObject:
+            args = parser.parse_args()
+            config = vars(args)
+            os.environ["KONFAI_OVERWRITE"] = "True"
+            os.environ["KONFAI_CLUSTER"] = "True"
 
-    with setupAPI(parser) as distributedObject:
-        args = parser.parse_args()
-        config = vars(args)
-        os.environ["DL_API_OVERWRITE"] = "True"
-        os.environ["DL_API_CLUSTER"] = "True"
-
-        n_gpu = len(config["gpu"].split(","))
-        distributedObject.setup(n_gpu*int(config["num_nodes"]))
-        import submitit
-        executor = submitit.AutoExecutor(folder="./Cluster/")
-        executor.update_parameters(name=config["name"], mem_gb=config["memory"], gpus_per_node=n_gpu, tasks_per_node=n_gpu//distributedObject.size, cpus_per_task=config["num_workers"], nodes=config["num_nodes"], timeout_min=config["time_limit"])
-        with TensorBoard(distributedObject.name):
-            executor.submit(distributedObject)
+            n_gpu = len(config["gpu"].split(","))
+            distributedObject.setup(n_gpu*int(config["num_nodes"]))
+            import submitit
+            executor = submitit.AutoExecutor(folder="./Cluster/")
+            executor.update_parameters(name=config["name"], mem_gb=config["memory"], gpus_per_node=n_gpu, tasks_per_node=n_gpu//distributedObject.size, cpus_per_task=config["num_workers"], nodes=config["num_nodes"], timeout_min=config["time_limit"])
+            with TensorBoard(distributedObject.name):
+                executor.submit(distributedObject)
+    except Exception as e:
+        print(e)
+        exit(1)

konfai/metric/measure.py

@@ -17,7 +17,7 @@ from abc import abstractmethod
 
 from konfai.utils.config import config
 from konfai.utils.utils import _getModule
-from konfai.data.HDF5 import ModelPatch
+from konfai.data.patching import ModelPatch
 from konfai.network.blocks import LatentDistribution
 from konfai.network.network import ModelLoader, Network
 
@@ -92,7 +92,7 @@ class PSNR(MaskedLoss):
         return psnr 
 
     def __init__(self, dynamic_range: Union[float, None] = None) -> None:
-        dynamic_range = dynamic_range if dynamic_range else 1024+3000
+        dynamic_range = dynamic_range if dynamic_range else 1024+3071
         super().__init__(partial(PSNR._loss, dynamic_range), False)
     
 class SSIM(MaskedLoss):
@@ -143,8 +143,11 @@ class Dice(Criterion):
         target = self.flatten(target)
         return (2.*(input * target).sum() + self.smooth)/(input.sum() + target.sum() + self.smooth)
 
+
     def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
         target = targets[0]
+        if output.shape[1] == 1:
+            output = F.one_hot(output.type(torch.int64), num_classes=int(torch.max(output).item()+1)).permute(0, len(target.shape), *[i+1 for i in range(len(target.shape)-1)]).float()
         target = F.one_hot(target.type(torch.int64), num_classes=output.shape[1]).permute(0, len(target.shape), *[i+1 for i in range(len(target.shape)-1)]).float().squeeze(2)
         return 1-torch.mean(self.dice_per_channel(output, target))
 
@@ -239,7 +242,7 @@ class PerceptualLoss(Criterion):
         @config(None)
         def __init__(self, losses: dict[str, float] = {"Gram": 1, "torch_nn_L1Loss": 1}) -> None:
             self.losses = losses
-            self.DL_args = os.environ['DEEP_LEARNING_API_CONFIG_PATH'] if "DEEP_LEARNING_API_CONFIG_PATH" in os.environ else ""
+            self.DL_args = os.environ['KONFAI_CONFIG_PATH'] if "KONFAI_CONFIG_PATH" in os.environ else ""
 
         def getLoss(self) -> dict[torch.nn.Module, float]:
             result: dict[torch.nn.Module, float] = {}
@@ -252,7 +255,7 @@ class PerceptualLoss(Criterion):
         super().__init__()
         self.path_model = path_model
         if self.path_model not in modelsRegister:
-            self.model = modelLoader.getModel(train=False, DL_args=os.environ['DEEP_LEARNING_API_CONFIG_PATH'].split("PerceptualLoss")[0]+"PerceptualLoss.Model", DL_without=["optimizer", "schedulers", "nb_batch_per_step", "init_type", "init_gain", "outputsCriterions", "drop_p"])
+            self.model = modelLoader.getModel(train=False, DL_args=os.environ['KONFAI_CONFIG_PATH'].split("PerceptualLoss")[0]+"PerceptualLoss.Model", DL_without=["optimizer", "schedulers", "nb_batch_per_step", "init_type", "init_gain", "outputsCriterions", "drop_p"])
             if path_model.startswith("https"):
                 state_dict = torch.hub.load_state_dict_from_url(path_model)
                 state_dict = {"Model": {self.model.getName() : state_dict["model"]}}

konfai/models/classification/convNeXt.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn.functional as F
 from konfai.network import network, blocks
 from konfai.utils.config import config
-from konfai.data.HDF5 import ModelPatch
+from konfai.data.patching import ModelPatch
 
 """
 "convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth", depths=[3, 3, 9, 3], dims=[96, 192, 384, 768]

konfai/models/classification/resnet.py

@@ -3,7 +3,7 @@ from typing import Type
 import torch
 from konfai.network import network, blocks
 from konfai.utils.config import config
-from konfai.data.HDF5 import ModelPatch
+from konfai.data.patching import ModelPatch
 
 """
 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', dim = 2, in_channels = 3, depths=[2, 2, 2, 2], widths = [64, 64, 128, 256, 512],  num_classes=1000, useBottleneck=False

konfai/models/generation/cStyleGan.py

@@ -3,7 +3,7 @@ import torch
 
 from konfai.network import network, blocks
 from konfai.utils.config import config
-from konfai.data.HDF5 import ModelPatch    
+from konfai.data.patching import ModelPatch    
 
 class MappingNetwork(network.ModuleArgsDict):
     def __init__(self, z_dim: int, c_dim: int, w_dim: int, num_layers: int, embed_features: int, layer_features: int):

konfai/models/generation/ddpm.py

@@ -8,7 +8,7 @@ import numpy as np
 
 from konfai.network import network, blocks
 from konfai.utils.config import config
-from konfai.data.HDF5 import ModelPatch
+from konfai.data.patching import ModelPatch
 from konfai.utils.utils import gpuInfo
 from konfai.metric.measure import Criterion
 

konfai/models/generation/diffusionGan.py

@@ -5,7 +5,7 @@ import numpy as np
 
 from konfai.network import network, blocks
 from konfai.utils.config import config
-from konfai.data.HDF5 import ModelPatch, Attribute
+from konfai.data.patching import ModelPatch, Attribute
 from konfai.data import augmentation
 from konfai.models.segmentation import UNet, NestedUNet
 from konfai.models.generation.ddpm import DDPM

konfai/models/generation/gan.py

@@ -3,7 +3,7 @@ import torch
 
 from konfai.network import network, blocks
 from konfai.utils.config import config
-from konfai.data.HDF5 import ModelPatch
+from konfai.data.patching import ModelPatch
 
 class Discriminator(network.Network):
 

konfai/models/segmentation/NestedUNet.py

@@ -3,7 +3,7 @@ import torch
 
 from konfai.network import network, blocks
 from konfai.utils.config import config
-from konfai.data.HDF5 import ModelPatch
+from konfai.data.patching import ModelPatch
 
      
 class NestedUNetBlock(network.ModuleArgsDict):

konfai/models/segmentation/UNet.py

@@ -3,7 +3,7 @@ from typing import Union
 
 from konfai.network import network, blocks
 from konfai.utils.config import config
-from konfai.data.HDF5 import ModelPatch
+from konfai.data.patching import ModelPatch
 
 class UNetHead(network.ModuleArgsDict):