konfai 1.1.5__py3-none-any.whl → 1.1.7__py3-none-any.whl

konfai/data/augmentation.py

@@ -57,7 +57,7 @@ class DataAugmentationsList():
 
     def load(self, key: str, datasets: list[Dataset]):
         for augmentation, prob in self.dataAugmentationsLoader.items():
-            module, name = _getModule(augmentation, "data.augmentation")
+            module, name = _getModule(augmentation, "konfai.data.augmentation")
             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)
@@ -264,7 +264,6 @@ class LumaFlip(ColorTransform):
 
 class HUE(ColorTransform):
 
-    @config("HUE")
     def __init__(self, hue_max: float, groups: Union[list[str], None] = None) -> None:
         super().__init__(groups)
         self.hue_max = hue_max

konfai/data/data_manager.py

@@ -24,11 +24,11 @@ from konfai.data.augmentation import DataAugmentationsList
 class 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()},
+    def __init__(self,  transforms : Union[dict[str, TransformLoader], list[Transform]] = {"default:Normalize:Standardize:Unsqueeze:TensorCast:ResampleIsotropic:ResampleResize": TransformLoader()},
+                        patch_transforms : Union[dict[str, TransformLoader], list[Transform]] = {"default:Normalize:Standardize:Unsqueeze:TensorCast:ResampleIsotropic:ResampleResize": TransformLoader()},
                         isInput: bool = True) -> None:
-        self._pre_transforms = pre_transforms
-        self._post_transforms = post_transforms
+        self._pre_transforms = transforms
+        self._post_transforms = patch_transforms
         self.pre_transforms : list[Transform] = []
         self.post_transforms : list[Transform] = []
         self.isInput = isInput
@@ -37,7 +37,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(KONFAI_ROOT(), group_src, group_dest))
+                    transform = transform.getTransform(classpath, DL_args =  "{}.Dataset.groups_src.{}.groups_dest.{}.transforms".format(KONFAI_ROOT(), group_src, group_dest))
                     transform.setDatasets(datasets)
                     self.pre_transforms.append(transform)
             else:
@@ -48,7 +48,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(KONFAI_ROOT(), group_src, group_dest))
+                    transform = transform.getTransform(classpath, DL_args = "{}.Dataset.groups_src.{}.groups_dest.{}.patch_transforms".format(KONFAI_ROOT(), group_src, group_dest))
                     transform.setDatasets(datasets)
                     self.post_transforms.append(transform)
             else:
@@ -65,9 +65,8 @@ class GroupTransform:
 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)
+    def __init__(self, transforms : Union[dict[str, TransformLoader], list[Transform]] = {"default:Normalize:Standardize:Unsqueeze:TensorCast:ResampleIsotropic:ResampleResize": TransformLoader()}):
+        super().__init__(transforms, None)
 
 class Group(dict[str, GroupTransform]):
 
@@ -184,7 +183,7 @@ class DatasetIter(data.Dataset):
         data = {}
         x, a, p = self.map[index]
         if x not in self._index_cache:
-            if x not in self._index_cache and len(self._index_cache) >= self.buffer_size and not self.use_cache:
+            if len(self._index_cache) >= self.buffer_size and not self.use_cache:
                 self._unloadData(self._index_cache[0])
             self._loadData(x)
 
@@ -196,10 +195,9 @@ class DatasetIter(data.Dataset):
 
 class Subset():
     
-    def __init__(self, subset: Union[str, list[int], list[str], None] = None, shuffle: bool = True, filter: str = None) -> None:
+    def __init__(self, subset: Union[str, list[int], list[str], None] = None, shuffle: bool = True) -> None:
         self.subset = subset
         self.shuffle = shuffle
-        self.filter: list[tuple[str]] = [tuple(a.split(":")) for a in filter.split(";")] if filter is not None else None
 
     def __call__(self, names: list[str], infos: list[dict[str, tuple[np.ndarray, Attribute]]]) -> set[str]:
         inter_name = set(names[0])
@@ -207,16 +205,7 @@ class Subset():
             inter_name = inter_name.intersection(set(n))
         names = sorted(list(inter_name))
         
-        names_filtred = []
-        if self.filter is not None:
-            for name in names:
-                for info in infos:
-                    if all([info[name][1].isInfo(f[0], f[1]) for f in self.filter]):
-                        names_filtred.append(name)
-                        continue
-        else:
-            names_filtred = names
-        size = len(names_filtred)
+        size = len(names)
         index = []
         if self.subset is None:
             index = list(range(0, size))
@@ -230,7 +219,7 @@ class Subset():
                     for name in f:
                         train_names.append(name.strip())
                 index = []
-                for i, name in enumerate(names_filtred):
+                for i, name in enumerate(names):
                     if name in train_names:
                         index.append(i)
             elif self.subset.startswith("~") and os.path.exists(self.subset[1:]):
@@ -239,7 +228,7 @@ class Subset():
                     for name in f:
                         exclude_names.append(name.strip())
                 index = []
-                for i, name in enumerate(names_filtred):
+                for i, name in enumerate(names):
                     if name not in exclude_names:
                         index.append(i)
 
@@ -252,26 +241,27 @@ class Subset():
                         index = self.subset
                 if isinstance(self.subset[0], str):
                     index = []
-                    for i, name in enumerate(names_filtred):
+                    for i, name in enumerate(names):
                         if name in self.subset:
                             index.append(i)
         if self.shuffle:
             index = random.sample(index, len(index))
-        return set([names_filtred[i] for i in index])
+        return set([names[i] for i in index])
     
     def __str__(self):
-        return "Subset : " + str(self.subset) + " shuffle : "+ str(self.shuffle) + " filter : "+ str(self.filter)
+        return "Subset : " + str(self.subset) + " shuffle : "+ str(self.shuffle)
+    
 class TrainSubset(Subset):
 
     @config()
-    def __init__(self, subset: Union[str, list[int], list[str], None] = None, shuffle: bool = True, filter: str = None) -> None:
-        super().__init__(subset, shuffle, filter)
+    def __init__(self, subset: Union[str, list[int], list[str], None] = None, shuffle: bool = True) -> None:
+        super().__init__(subset, shuffle)
 
 class PredictionSubset(Subset):
 
     @config()
-    def __init__(self, subset: Union[str, list[int], list[str], None] = None, filter: str = None) -> None:
-        super().__init__(subset, False, filter)
+    def __init__(self, subset: Union[str, list[int], list[str], None] = None) -> None:
+        super().__init__(subset, False)
 
 class Data(ABC):
     
@@ -280,7 +270,6 @@ class Data(ABC):
                         patch : Union[DatasetPatch, None],
                         use_cache : bool,
                         subset : Subset,
-                        num_workers : int,
                         batch_size : int,
                         validation: Union[float, str, list[int], list[str], None] = None,
                         inlineAugmentations: bool = False,
@@ -293,7 +282,7 @@ class Data(ABC):
         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)
-        self.dataLoader_args = dict(num_workers=num_workers if use_cache else 0, pin_memory=True)
+        self.dataLoader_args = dict(num_workers=int(os.environ["KONFAI_WORKERS"]), pin_memory=True)
         self.data : list[list[dict[str, list[DatasetManager]]], dict[str, list[DatasetManager]]] = []
         self.map : list[list[list[tuple[int, int, int]]], list[tuple[int, int, int]]] = []
         self.datasets: dict[str, Dataset] = {}
@@ -547,10 +536,9 @@ class DataTrain(Data):
                         patch : Union[DatasetPatch, None] = DatasetPatch(),
                         use_cache : bool = True,
                         subset : Union[TrainSubset, dict[str, TrainSubset]] = TrainSubset(),
-                        num_workers : int = 4,
                         batch_size : int = 1,
                         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 {})
+        super().__init__(dataset_filenames, groups_src, patch, use_cache, subset, batch_size, validation, inlineAugmentations, augmentations if augmentations else {})
 
 class DataPrediction(Data):
 
@@ -560,10 +548,9 @@ class DataPrediction(Data):
                         augmentations : Union[dict[str, DataAugmentationsList], None] = {"DataAugmentation_0" : DataAugmentationsList()},
                         patch : Union[DatasetPatch, None] = DatasetPatch(),
                         subset : Union[PredictionSubset, dict[str, PredictionSubset]] = PredictionSubset(),
-                        num_workers : int = 4,
                         batch_size : int = 1) -> None:
 
-        super().__init__(dataset_filenames, groups_src, patch, False, subset, num_workers, batch_size, dataAugmentationsList=augmentations if augmentations else {})
+        super().__init__(dataset_filenames, groups_src, patch, False, subset, batch_size, dataAugmentationsList=augmentations if augmentations else {})
 
 class DataMetric(Data):
 
@@ -571,7 +558,6 @@ class DataMetric(Data):
     def __init__(self,  dataset_filenames : list[str] = ["default:./Dataset"], 
                         groups_src : dict[str, GroupMetric] = {"default" : GroupMetric()},
                         subset : Union[PredictionSubset, dict[str, PredictionSubset]] = PredictionSubset(),
-                        validation: Union[str, None] = None,
-                        num_workers : int = 4) -> None:
+                        validation: Union[str, None] = None) -> 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=validation)
+        super().__init__(dataset_filenames=dataset_filenames, groups_src=groups_src, patch=None, use_cache=False, subset=subset, batch_size=1, validation=validation)

konfai/data/patching.py

@@ -81,7 +81,6 @@ class PathCombine(ABC):
 
 class Mean(PathCombine):
 
-    @config("Mean")
     def __init__(self) -> None:
         super().__init__()
 
@@ -90,7 +89,6 @@ class Mean(PathCombine):
     
 class Cosinus(PathCombine):
 
-    @config("Cosinus")
     def __init__(self) -> None:
         super().__init__()
 
@@ -144,7 +142,7 @@ class Accumulator():
 
 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:
+    def __init__(self, patch_size: list[int], overlap: Union[int, None], padValue: float = 0, extend_slice: int = 0) -> None:
         self.patch_size = patch_size
         self.overlap = overlap 
         if isinstance(self.overlap, int):
@@ -152,15 +150,8 @@ class Patch(ABC):
                  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
-        self.mask = None
         self.padValue = padValue
         self.extend_slice = extend_slice
-        if self.path_mask is not None:
-            if os.path.exists(self.path_mask):
-                self.mask = torch.tensor(sitk.GetArrayFromImage(sitk.ReadImage(self.path_mask)))
-            else:
-                raise NameError('Mask file not found')
             
     def load(self, shape : dict[int, list[int]], a: int = 0) -> None:
         self._patch_slices[a], self._nb_patch_per_dim[a] = get_patch_slices_from_shape(self.patch_size, shape, self.overlap)
@@ -199,9 +190,6 @@ class Patch(ABC):
             padding.append(p)
 
         data_sliced = F.pad(data_sliced, tuple(padding), "constant", 0 if data_sliced.dtype == torch.uint8 and self.padValue < 0 else self.padValue)
-        if self.mask is not None:
-            outside = torch.ones_like(data_sliced)*(0 if data_sliced.dtype == torch.uint8 and self.padValue < 0 else self.padValue)
-            data_sliced = torch.where(self.mask == 0, outside, data_sliced)
         
         for d in [i for i, v in enumerate(reversed(self.patch_size)) if v == 1]:
             data_sliced = torch.squeeze(data_sliced, dim = len(data_sliced.shape)-d-1)
@@ -213,14 +201,14 @@ class Patch(ABC):
 class DatasetPatch(Patch):
 
     @config("Patch")
-    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)
+    def __init__(self, patch_size : list[int] = [128, 128, 128], overlap : Union[int, None] = None, padValue: float = 0, extend_slice: int = 0) -> None:
+        super().__init__(patch_size, overlap, padValue, extend_slice)
 
 class ModelPatch(Patch):
 
     @config("Patch")
-    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)
+    def __init__(self, patch_size : list[int] = [128, 128, 128], overlap : Union[int, None] = None, patchCombine: Union[str, None] = None, padValue: float = 0, extend_slice: int = 0) -> None:
+        super().__init__(patch_size, overlap, padValue, extend_slice)
         self.patchCombine = patchCombine
 
     def disassemble(self, *dataList: torch.Tensor) -> Iterator[list[torch.Tensor]]:
@@ -247,7 +235,7 @@ class DatasetManager():
         for transformFunction in pre_transforms:
             _shape = transformFunction.transformShape(_shape, cache_attribute)
         
-        self.patch = DatasetPatch(patch_size=patch.patch_size, overlap=patch.overlap, mask=patch.path_mask, padValue=patch.padValue, extend_slice=patch.extend_slice) if patch else DatasetPatch(_shape)
+        self.patch = DatasetPatch(patch_size=patch.patch_size, overlap=patch.overlap, padValue=patch.padValue, extend_slice=patch.extend_slice) if patch else DatasetPatch(_shape)
         self.patch.load(_shape, 0)
         self.shape = _shape
         self.dataAugmentationsList = dataAugmentationsList
@@ -314,7 +302,7 @@ class DatasetManager():
                     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

konfai/data/transform.py

@@ -36,7 +36,7 @@ class TransformLoader:
         pass
     
     def getTransform(self, classpath : str, DL_args : str) -> Transform:
-        module, name = _getModule(classpath, "data.transform")
+        module, name = _getModule(classpath, "konfai.data.transform")
         return config("{}.{}".format(DL_args, classpath))(getattr(importlib.import_module(module), name))(config = None)
 
 class Clip(Transform):
@@ -213,7 +213,7 @@ class Resample(Transform, ABC):
 
 class ResampleToResolution(Resample):
 
-    def __init__(self, spacing : list[Union[float, None]] = [1., 1., 1.]) -> None:
+    def __init__(self, spacing : list[float] = [1., 1., 1.]) -> None:
         self.spacing = torch.tensor([0 if s < 0 else s for s in spacing])
 
     def transformShape(self, shape: list[int], cache_attribute: Attribute) -> list[int]:
@@ -244,34 +244,36 @@ class ResampleToResolution(Resample):
         cache_attribute["Size"] = np.asarray(size)
         return self._resample(input, size)
 
-class ResampleToSize(Resample):
+class ResampleToShape(Resample):
 
-    def __init__(self, size : list[int] = [100,512,512]) -> None:
-        self.size = size
+    def __init__(self, shape : list[float] = [100,256,256]) -> None:
+        self.shape = torch.tensor([0 if s < 0 else s for s in shape])
 
     def transformShape(self, shape: list[int], cache_attribute: Attribute) -> list[int]:
+        print(shape)
         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):
+        if len(shape) != len(self.shape):
             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
+        new_shape = self.shape
+        for i, s in enumerate(self.shape):
+            if s == 0:
+                new_shape[i] = shape[i]
+        print(new_shape)
+        return new_shape
     
     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]
+        shape = self.shape
+        image_shape =  torch.tensor([int(x) for x in torch.tensor(input.shape[1:])])
+        for i, s in enumerate(self.shape):
+            if s == 0:
+                shape[i] = image_shape[i]
         if "Spacing" in cache_attribute:
-            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)
+            cache_attribute["Spacing"] = torch.flip(image_shape/shape*torch.flip(cache_attribute.get_tensor("Spacing"), dims=[0]), dims=[0])
+        cache_attribute["Size"] = image_shape
+        cache_attribute["Size"] = shape
+        return self._resample(input, shape)
 
 class ResampleTransform(Transform):
 
@@ -410,7 +412,7 @@ class Gradient(Transform):
     def inverse(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
         return input
 
-class ArgMax(Transform):
+class Argmax(Transform):
 
     def __init__(self, dim: int = 0) -> None:
         self.dim = dim

konfai/evaluator.py

@@ -27,7 +27,7 @@ class CriterionsLoader():
     def getCriterions(self, output_group : str, target_group : str) -> dict[torch.nn.Module, CriterionsAttr]:
         criterions = {}
         for module_classpath, criterionsAttr in self.criterionsLoader.items():
-            module, name = _getModule(module_classpath, "metric.measure")
+            module, name = _getModule(module_classpath, "konfai.metric.measure")
             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
 

konfai/main.py

@@ -9,8 +9,6 @@ import sys
 sys.path.insert(0, os.getcwd())
 
 def main():
-    import tracemalloc
-
     parser = argparse.ArgumentParser(description="KonfAI", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     try:
         with setup(parser) as distributedObject:

konfai/metric/measure.py

@@ -9,7 +9,6 @@ from scipy import linalg
 
 import torch.nn.functional as F
 import os
-import tqdm
 
 from typing import Callable, Union
 from functools import partial
@@ -19,10 +18,11 @@ import copy
 from abc import abstractmethod
 
 from konfai.utils.config import config
-from konfai.utils.utils import _getModule
+from konfai.utils.utils import _getModule, download_url
 from konfai.data.patching import ModelPatch
 from konfai.network.blocks import LatentDistribution
 from konfai.network.network import ModelLoader, Network
+from tqdm import tqdm
 
 modelsRegister = {}
 
@@ -122,7 +122,7 @@ class LPIPS(MaskedLoss):
 
         patchIterator = datasetPatch.disassemble(LPIPS.normalize(x), LPIPS.normalize(y))
         loss = 0
-        with tqdm.tqdm(iterable = enumerate(patchIterator), leave=False, total=datasetPatch.getSize(0)) as batch_iter:
+        with tqdm(iterable = enumerate(patchIterator), leave=False, total=datasetPatch.getSize(0)) as batch_iter:
             for i, patch_input in batch_iter:
                 real, fake = LPIPS.preprocessing(patch_input[0]), LPIPS.preprocessing(patch_input[1])
                 loss += loss_fn_alex(real, fake).item()
@@ -146,9 +146,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]
+        target = F.interpolate(
+            targets[0],
+            output.shape[2:],
+            mode="nearest")
         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)
@@ -242,7 +244,7 @@ class PerceptualLoss(Criterion):
     
     class Module():
         
-        @config(None)
+        @config()
         def __init__(self, losses: dict[str, float] = {"Gram": 1, "torch_nn_L1Loss": 1}) -> None:
             self.losses = losses
             self.DL_args = os.environ['KONFAI_CONFIG_PATH'] if "KONFAI_CONFIG_PATH" in os.environ else ""
@@ -250,7 +252,7 @@ class PerceptualLoss(Criterion):
         def getLoss(self) -> dict[torch.nn.Module, float]:
             result: dict[torch.nn.Module, float] = {}
             for loss, l in self.losses.items():
-                module, name = _getModule(loss, "metric.measure")
+                module, name = _getModule(loss, "konfai.metric.measure")
                 result[config(self.DL_args)(getattr(importlib.import_module(module), name))(config=None)] = l   
             return result
         
@@ -413,16 +415,19 @@ class FocalLoss(Criterion):
         self.gamma = gamma
         self.reduction = reduction
 
-    def forward(self, output: torch.Tensor, targets: list[torch.Tensor]) -> torch.Tensor:
-        target = targets[0].long()
+    def forward(self, output: torch.Tensor, *targets: list[torch.Tensor]) -> torch.Tensor:
+        target = F.interpolate(
+            targets[0],
+            output.shape[2:],
+            mode="nearest").long()
         
         logpt = F.log_softmax(output, dim=1)
         pt = torch.exp(logpt)
 
-        logpt = logpt.gather(1, target.unsqueeze(1)) 
-        pt = pt.gather(1, target.unsqueeze(1))
+        logpt = logpt.gather(1, target) 
+        pt = pt.gather(1, target)
 
-        at = self.alpha[target].unsqueeze(1)
+        at = self.alpha.to(target.device)[target].unsqueeze(1)
         loss = -at * ((1 - pt) ** self.gamma) * logpt
 
         if self.reduction == "mean":
@@ -458,13 +463,11 @@ class FID(Criterion):
         return features
 
     def calculate_fid(real_features: np.ndarray, generated_features: np.ndarray) -> float:
-        # Calculate mean and covariance statistics
         mu1 = np.mean(real_features, axis=0)
         sigma1 = np.cov(real_features, rowvar=False)
         mu2 = np.mean(generated_features, axis=0)
         sigma2 = np.cov(generated_features, rowvar=False)
 
-        # Calculate FID score
         diff = mu1 - mu2
         covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
         if np.iscomplexobj(covmean):
@@ -480,7 +483,6 @@ class FID(Criterion):
         generated_features = FID.get_features(generated_images, self.inception_model)
 
         return FID.calculate_fid(real_features, generated_features)
-        
 
 class MutualInformationLoss(torch.nn.Module):
     def __init__(self, num_bins: int = 23, sigma_ratio: float = 0.5, smooth_nr: float = 1e-7, smooth_dr: float = 1e-7) -> None:
@@ -498,7 +500,6 @@ class MutualInformationLoss(torch.nn.Module):
         target_weight, target_probability = self.parzen_windowing_gaussian(target)
         return pred_weight, pred_probability, target_weight, target_probability
 
-
     def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         img = torch.clamp(img, 0, 1)
         img = img.reshape(img.shape[0], -1, 1)  # (batch, num_sample, 1)
@@ -512,4 +513,72 @@ class MutualInformationLoss(torch.nn.Module):
         pab = torch.bmm(wa.permute(0, 2, 1), wb.to(wa)).div(wa.shape[1])  # (batch, num_bins, num_bins)
         papb = torch.bmm(pa.permute(0, 2, 1), pb.to(pa))  # (batch, num_bins, num_bins)
         mi = torch.sum(pab * torch.log((pab + self.smooth_nr) / (papb + self.smooth_dr) + self.smooth_dr), dim=(1, 2))  # (batch)
-        return torch.mean(mi).neg()  # average over the batch and channel ndims
+        return torch.mean(mi).neg()  # average over the batch and channel ndims
+
+class FOCUS(Criterion): #Feature-Oriented Comparison for Unpaired Synthesis
+
+    def __init__(self, model_name: str, shape: list[int] = [0,0], in_channels: int = 1, losses: dict[str, list[float]] = {"Gram": 1, "torch_nn_L1Loss": 1}) -> None:
+        super().__init__()
+        if model_name is None:
+            return
+        self.shape = shape
+        self.in_channels = in_channels
+        self.losses: dict[torch.nn.Module, list[float]] = {}
+        for loss, weights in losses.items():
+            module, name = _getModule(loss, "konfai.metric.measure")
+            self.losses[config(os.environ['KONFAI_CONFIG_PATH'])(getattr(importlib.import_module(module), name))(config=None)] = weights
+
+        self.model_path = download_url(model_name, "https://huggingface.co/VBoussot/impact-torchscript-models/resolve/main/")
+        self.model: torch.nn.Module = torch.jit.load(self.model_path)
+        self.dim = len(self.shape)
+        self.shape = self.shape if all([s > 0 for s in self.shape]) else None
+        self.modules_loss: dict[str, dict[torch.nn.Module, float]] = {}
+
+        try:
+            dummy_input = torch.zeros((1, self.in_channels, *(self.shape if self.shape else [224] * self.dim)))
+            out = self.model(dummy_input)
+            if not isinstance(out, (list, tuple)):
+                raise TypeError(f"Expected model output to be a list or tuple, but got {type(out)}.")
+            if len(self.weight) != len(out):
+                raise ValueError(f"Mismatch between number of weights ({len(self.weight)}) and model outputs ({len(out)}).")
+        except Exception as e:
+            msg = (
+                f"[Model Sanity Check Failed]\n"
+                f"Input shape attempted: {dummy_input.shape}\n"
+                f"Expected output length: {len(self.weight)}\n"
+                f"Error: {type(e).__name__}: {e}"
+            )
+            raise RuntimeError(msg) from e
+        self.model = None
+        
+    def preprocessing(self, input: torch.Tensor) -> torch.Tensor:
+        if self.shape is not None and not all([input.shape[-i-1] == size for i, size in enumerate(reversed(self.shape[2:]))]):
+            input = torch.nn.functional.interpolate(input, mode=self.mode, size=tuple(self.shape), align_corners=False).type(torch.float32)
+        if input.shape[1] != self.in_channels:
+            input = input.repeat(tuple([1,3] + [1 for _ in range(self.dim)]))   
+        return input
+    
+    def _compute(self, output: torch.Tensor, targets: list[torch.Tensor]) -> torch.Tensor:
+        loss = torch.zeros((1), requires_grad = True).to(output.device, non_blocking=False).type(torch.float32)
+        output = self.preprocessing(output)
+        targets = [self.preprocessing(target) for target in targets]
+        self.model.to(output.device)
+        for zipped_output in zip(self.model(output), *[self.model(target) for target in targets]):
+            output_features = zipped_output[0]
+            targets_features = zipped_output[1:]            
+            for target_features, (loss_function, weights) in zip(targets_features, self.losses.items()):
+                for output_feature, target_feature, weight in zip(output_features ,target_features, weights):
+                    loss += weights*loss_function(output_feature, target_feature)
+        return loss
+    
+    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
+        if self.model is None:
+            self.model = torch.jit.load(self.model_path)
+        loss = torch.zeros((1), requires_grad = True).to(output.device, non_blocking=False).type(torch.float32)
+        if len(output.shape) == 5 and self.dim == 2:
+            for i in range(output.shape[2]):
+                loss = loss + self._compute(output[:, :, i, ...], [t[:, :, i, ...] for t in targets])
+            loss/=output.shape[2]
+        else:
+            loss = self._compute(output, targets)
+        return loss.to(output)

konfai/metric/schedulers.py

@@ -2,10 +2,10 @@ import torch
 import numpy as np
 from abc import abstractmethod
 from konfai.utils.config import config
+from functools import partial
 
 class Scheduler():
 
-    @config("Scheduler")
     def __init__(self, start_value: float) -> None:
         self.baseValue = float(start_value)
         self.it = 0
@@ -19,7 +19,6 @@ class Scheduler():
 
 class Constant(Scheduler):
 
-    @config("Constant")
     def __init__(self, value: float = 1):
         super().__init__(value)
 
@@ -28,8 +27,7 @@ class Constant(Scheduler):
 
 class CosineAnnealing(Scheduler):
     
-    @config("CosineAnnealing")
-    def __init__(self, start_value: float, eta_min: float = 0.00001, T_max: int = 100):
+    def __init__(self, start_value: float = 1, eta_min: float = 0.00001, T_max: int = 100):
         super().__init__(start_value)
         self.eta_min = eta_min
         self.T_max = T_max
@@ -37,13 +35,10 @@ class CosineAnnealing(Scheduler):
     def get_value(self):
         return self.eta_min + (self.baseValue - self.eta_min) *(1 + np.cos(self.it * torch.pi / self.T_max)) / 2
 
-class CosineAnnealing(Scheduler):
+class Warmup(torch.optim.lr_scheduler.LambdaLR):
     
-    @config("CosineAnnealing")
-    def __init__(self, start_value: float, eta_min: float = 0.00001, T_max: int = 100):
-        super().__init__(start_value)
-        self.eta_min = eta_min
-        self.T_max = T_max
+    def warmup(warmup_steps: int, step: int) -> float:
+        return min(1.0, (step+1) / (warmup_steps+1))
 
-    def get_value(self):
-        return self.eta_min + (self.baseValue - self.eta_min) *(1 + np.cos(self.it * torch.pi / self.T_max)) / 2
+    def __init__(self, optimizer: torch.optim.Optimizer, warmup_steps: int = 10, last_epoch=-1, verbose="deprecated"):
+        super().__init__(optimizer, partial(Warmup.warmup, warmup_steps), last_epoch, verbose)

konfai/models/classification/convNeXt.py

@@ -157,10 +157,9 @@ class Head(network.ModuleArgsDict):
 
 class ConvNeXt(network.Network):
     
-    @config("ConvNeXt")
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : ModelPatch = ModelPatch(),
                     dim : int = 3,

konfai/models/classification/resnet.py

@@ -98,10 +98,9 @@ class Head(network.ModuleArgsDict):
         
 class ResNet(network.Network):
 
-    @config("ResNet")
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : ModelPatch = ModelPatch(),
                     dim : int = 3,