konfai 1.1.9__py3-none-any.whl → 1.2.0__py3-none-any.whl

konfai/data/data_manager.py

@@ -6,6 +6,7 @@ from abc import ABC, abstractmethod
 from collections.abc import Iterator, Mapping
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from functools import partial
+from typing import cast
 
 import numpy as np
 import torch
@@ -270,50 +271,55 @@ class Subset:
         self.subset = subset
         self.shuffle = shuffle
 
+    def _get_index(self, subset: str | int, names: list[str]) -> list[int]:
+        size = len(names)
+        index = []
+        if isinstance(subset, int):
+            index.append(subset)
+        elif ":" in subset:
+            r = np.clip(
+                np.asarray([int(subset.split(":")[0]), int(subset.split(":")[1])]),
+                0,
+                size,
+            )
+            index = list(range(r[0], r[1]))
+        elif os.path.exists(subset):
+            train_names = []
+            with open(subset) as f:
+                for name in f:
+                    train_names.append(name.strip())
+            index = []
+            for i, name in enumerate(names):
+                if name in train_names:
+                    index.append(i)
+        elif subset.startswith("~") and os.path.exists(subset[1:]):
+            exclude_names = []
+            with open(subset[1:]) as f:
+                for name in f:
+                    exclude_names.append(name.strip())
+            index = []
+            for i, name in enumerate(names):
+                if name not in exclude_names:
+                    index.append(i)
+        return index
+
     def __call__(self, names: list[str], infos: dict[str, tuple[list[int], Attribute]]) -> set[str]:
         names = sorted(names)
-
         size = len(names)
-        index = []
+
         if self.subset is None:
             index = list(range(0, size))
-        elif isinstance(self.subset, str):
-            if ":" in self.subset:
-                r = np.clip(
-                    np.asarray([int(self.subset.split(":")[0]), int(self.subset.split(":")[1])]),
-                    0,
-                    size,
-                )
-                index = list(range(r[0], r[1]))
-            elif os.path.exists(self.subset):
-                train_names = []
-                with open(self.subset) as f:
-                    for name in f:
-                        train_names.append(name.strip())
-                index = []
-                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:]):
-                exclude_names = []
-                with open(self.subset[1:]) as f:
-                    for name in f:
-                        exclude_names.append(name.strip())
-                index = []
-                for i, name in enumerate(names):
-                    if name not in exclude_names:
-                        index.append(i)
-
         elif isinstance(self.subset, list):
-            index = []
-            if len(self.subset) > 0:
-                for s in self.subset:
-                    if isinstance(s, int):
-                        index.append(s)
-                    elif isinstance(s, str):
-                        for i, name in enumerate(names):
-                            if name in self.subset:
-                                index.append(i)
+            index_set: set[int] = set()
+            for s in self.subset:
+                if len(index_set) == 0:
+                    index_set.update(set(self._get_index(s, names)))
+                else:
+                    index_set = index_set.intersection(set(self._get_index(s, names)))
+                index = list(index_set)
+                print(index)
+        else:
+            index = self._get_index(self.subset, names)
         if self.shuffle:
             index = random.sample(index, len(index))  # nosec B311
         return {names[i] for i in index}
@@ -456,7 +462,7 @@ class Data(ABC):
                 mappings.append(list(mapping[-offset:]) if itr + offset > len(mapping) else mapping[itr : itr + offset])
         return mappings
 
-    def get_data(self, world_size: int) -> list[list[DataLoader]]:
+    def get_data(self, world_size: int) -> tuple[list[list[DataLoader]], list[str], list[str]]:
         datasets: dict[str, list[tuple[str, bool]]] = {}
         if self.dataset_filenames is None or len(self.dataset_filenames) == 0:
             raise DatasetManagerError("No dataset filenames were provided")
@@ -502,8 +508,8 @@ class Data(ABC):
                     f"Group source '{group_src}' not found in any dataset.",
                     f"Dataset filenames provided: {self.dataset_filenames}",
                     f"Available groups across all datasets: "
-                    "{[f'{f} {d.get_group()}' for f, d in self.datasets.items()]}"
-                    f"Please check that an entry in the dataset with the name '{group_src}.{format}' exists.",
+                    f"{[f'{f} {d.get_group()}' for f, d in self.datasets.items()]}\n"
+                    f"Please check that an entry in the dataset with the name '{group_src}' exists.",
                 )
 
             for group_dest in self.groups_src[group_src]:
@@ -596,8 +602,7 @@ class Data(ABC):
 
         data, mapping = self._get_datasets(list(subset_names), dataset_name)
 
-        train_mapping = mapping
-        validate_mapping = []
+        index = []
         if isinstance(self.validation, float) or isinstance(self.validation, int):
             if self.validation <= 0 or self.validation >= 1:
                 raise DatasetManagerError(
@@ -605,24 +610,16 @@ class Data(ABC):
                     f"Received: {self.validation}",
                     "Example: validation = 0.2  # for a 20% validation split",
                 )
-
-            train_mapping, validate_mapping = (
-                mapping[: int(math.floor(len(mapping) * (1 - self.validation)))],
-                mapping[int(math.floor(len(mapping) * (1 - self.validation))) :],
-            )
+            index = [m[0] for m in mapping[int(math.floor(len(mapping) * (1 - self.validation))) :]]
         elif isinstance(self.validation, str):
             if ":" in self.validation:
                 index = list(range(int(self.validation.split(":")[0]), int(self.validation.split(":")[1])))
-                train_mapping = [m for m in mapping if m[0] not in index]
-                validate_mapping = [m for m in mapping if m[0] in index]
             elif os.path.exists(self.validation):
                 validation_names = []
                 with open(self.validation) 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]
-                train_mapping = [m for m in mapping if m[0] not in index]
-                validate_mapping = [m for m in mapping if m[0] in index]
             else:
                 raise DatasetManagerError(
                     f"Invalid string value for 'validation': '{self.validation}'",
@@ -634,25 +631,23 @@ class Data(ABC):
                     "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):
-                    train_mapping = [m for m in mapping if m[0] not in self.validation]
-                    validate_mapping = [m for m in mapping 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_mapping = [m for m in mapping if m[0] not in index]
-                    validate_mapping = [m for m in mapping if m[0] in index]
-                else:
-                    raise DatasetManagerError(
-                        "Invalid list type for 'validation': elements of type "
-                        f"'{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 isinstance(self.validation[0], int):
+                index = cast(list[int], self.validation)
+            elif isinstance(self.validation[0], str):
+                index = [i for i, n in enumerate(subset_names) if n in self.validation]
+            else:
+                raise DatasetManagerError(
+                    "Invalid list type for 'validation': elements of type "
+                    f"'{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}",
+                )
+        train_mapping = [m for m in mapping if m[0] not in index]
+        validate_mapping = [m for m in mapping if m[0] in index]
+
         if len(train_mapping) == 0:
             raise DatasetManagerError(
                 "No data left for training after applying the validation split.",
@@ -668,6 +663,9 @@ class Data(ABC):
                 f"Validation setting: {self.validation}",
                 "Please increase the validation size, increase the dataset, or disable validation.",
             )
+
+        validation_names = [name for i, name in enumerate(subset_names) if i in index]
+        train_names = [name for name in subset_names if name not in validation_names]
         train_mappings = Data._split(train_mapping, world_size)
         validate_mappings = Data._split(validate_mapping, world_size)
 
@@ -701,7 +699,7 @@ class Data(ABC):
                         **self.dataLoader_args,
                     )
                 )
-        return data_loaders
+        return data_loaders, train_names, validation_names
 
 
 class DataTrain(Data):
@@ -773,7 +771,7 @@ class DataMetric(Data):
             dataset_filenames=dataset_filenames,
             groups_src=groups_src,
             patch=None,
-            use_cache=False,
+            use_cache=True,
             subset=subset,
             batch_size=1,
             validation=validation,

konfai/data/transform.py

@@ -1,4 +1,5 @@
 import importlib
+import tempfile
 from abc import ABC, abstractmethod
 from typing import Any
 
@@ -670,3 +671,26 @@ class OneHot(Transform):
 
     def inverse(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
         return torch.argmax(tensor, dim=1).unsqueeze(1)
+
+
+class TotalSegmentator(Transform):
+
+    def __init__(self, task: str = "total"):
+        super().__init__()
+        self.task = task
+
+    def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
+        from totalsegmentator.python_api import totalsegmentator
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            image = data_to_image(tensor.numpy(), cache_attribute)
+            sitk.WriteImage(image, tmpdir + "/image.nii.gz")
+            seg = totalsegmentator(tmpdir + "/image.nii.gz", tmpdir, task=self.task, skip_saving=True, quiet=True)
+        return (
+            torch.from_numpy(np.array(np.asanyarray(seg.dataobj), copy=True).astype(np.uint8, copy=False))
+            .permute(2, 1, 0)
+            .unsqueeze(0)
+        )
+
+    def inverse(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
+        return tensor

konfai/evaluator.py

@@ -152,14 +152,14 @@ class Statistics:
                 - mean and count
         """
         return {
-            "max": np.max(values),
-            "min": np.min(values),
-            "std": np.std(values),
-            "25pc": np.percentile(values, 25),
-            "50pc": np.percentile(values, 50),
-            "75pc": np.percentile(values, 75),
-            "mean": np.mean(values),
-            "count": len(values),
+            "max": np.nanmax(values) if np.any(~np.isnan(values)) else np.nan,
+            "min": np.nanmin(values) if np.any(~np.isnan(values)) else np.nan,
+            "std": np.nanstd(values) if np.any(~np.isnan(values)) else np.nan,
+            "25pc": np.nanpercentile(values, 25) if np.any(~np.isnan(values)) else np.nan,
+            "50pc": np.nanpercentile(values, 50) if np.any(~np.isnan(values)) else np.nan,
+            "75pc": np.nanpercentile(values, 75) if np.any(~np.isnan(values)) else np.nan,
+            "mean": np.nanmean(values) if np.any(~np.isnan(values)) else np.nan,
+            "count": np.count_nonzero(~np.isnan(values)) if np.any(~np.isnan(values)) else np.nan,
         }
 
     def write(self, outputs: list[dict[str, dict[str, Any]]]) -> None:
@@ -261,10 +261,26 @@ class Evaluator(DistributedObject):
                 ]
                 name = data_dict[output_group][1][0]
                 for metric in self.metrics[output_group][target_group]:
-                    result[f"{output_group}:{target_group}:{metric.__class__.__name__}"] = metric(
+                    loss = metric(
                         (data_dict[output_group][0].to(0) if torch.cuda.is_available() else data_dict[output_group][0]),
                         *targets,
-                    ).item()
+                    )
+                    if isinstance(loss, tuple):
+                        true_loss = loss[1]
+                    else:
+                        true_loss = loss.item()
+
+                    if isinstance(true_loss, dict):
+                        loss = 0
+                        c = 0
+                        for k, v in true_loss.items():
+                            result[f"{output_group}:{target_group}:{metric.__class__.__name__}:{k}"] = v
+                            if not np.isnan(v):
+                                loss += v
+                                c += 1
+                        result[f"{output_group}:{target_group}:{metric.__class__.__name__}"] = loss / c
+                    else:
+                        result[f"{output_group}:{target_group}:{metric.__class__.__name__}"] = true_loss
         statistics.add(result, name)
         return result
 
@@ -304,7 +320,7 @@ class Evaluator(DistributedObject):
             f"{self.metric_path}{metric_namefile_src}.yml",
         )
 
-        self.dataloader = self.dataset.get_data(world_size)
+        self.dataloader, _, _ = self.dataset.get_data(world_size)
 
         groups_dest = [group for groups in self.dataset.groups_src.values() for group in groups]
 

konfai/metric/measure.py

@@ -46,7 +46,7 @@ class MaskedLoss(Criterion):
         self.loss = loss
         self.mode_image_masked = mode_image_masked
 
-    def get_mask(self, *targets: torch.Tensor) -> torch.Tensor:
+    def get_mask(self, targets: list[torch.Tensor]) -> torch.Tensor | None:
         result = None
         if len(targets) > 0:
             result = targets[0]
@@ -54,29 +54,35 @@ class MaskedLoss(Criterion):
                 result = result * mask
         return result
 
-    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> tuple[torch.Tensor, float]:
         loss = torch.tensor(0, dtype=torch.float32).to(output.device)
-        mask = self.get_mask(targets[1:])
-        for batch in range(output.shape[0]):
-            if mask is not None:
+        true_loss = 0
+        true_nb = 0
+        mask = self.get_mask(list(targets[1:]))
+        if mask is not None:
+            for batch in range(output.shape[0]):
                 if self.mode_image_masked:
-                    for i in torch.unique(mask[batch]):
-                        if i != 0:
-                            loss += self.loss(
-                                output[batch, ...] * torch.where(mask == i, 1, 0),
-                                targets[0][batch, ...] * torch.where(mask == i, 1, 0),
-                            )
+                    loss_b = self.loss(
+                        output[batch, ...] * torch.where(mask == 1, 1, 0),
+                        targets[0][batch, ...] * torch.where(mask == 1, 1, 0),
+                    )
                 else:
-                    for i in torch.unique(mask[batch]):
-                        if i != 0:
-                            index = mask[batch, ...] == i
-                            loss += self.loss(
-                                torch.masked_select(output[batch, ...], index),
-                                torch.masked_select(targets[0][batch, ...], index),
-                            )
-            else:
-                loss += self.loss(output[batch, ...], targets[0][batch, ...])
-        return loss / output.shape[0]
+                    index = mask[batch, ...] == 1
+                    loss_b = self.loss(
+                        torch.masked_select(output[batch, ...], index),
+                        torch.masked_select(targets[0][batch, ...], index),
+                    )
+
+                loss += loss_b
+                if torch.any(mask[batch] == 1):
+                    true_loss += loss_b.item()
+                    true_nb += 1
+        else:
+            loss_tmp = self.loss(output, targets[0])
+            loss += loss_tmp
+            true_loss += loss_tmp.item()
+            true_nb += 1
+        return loss / output.shape[0], np.nan if true_nb == 0 else true_loss / true_nb
 
 
 class MSE(MaskedLoss):
@@ -137,7 +143,7 @@ class LPIPS(MaskedLoss):
 
     @staticmethod
     def preprocessing(tensor: torch.Tensor) -> torch.Tensor:
-        return tensor.repeat((1, 3, 1, 1))
+        return tensor.repeat((1, 3, 1, 1)).to(0)
 
     @staticmethod
     def _loss(loss_fn_alex, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
@@ -159,14 +165,14 @@ class LPIPS(MaskedLoss):
     def __init__(self, model: str = "alex") -> None:
         import lpips
 
-        super().__init__(partial(LPIPS._loss, lpips.LPIPS(net=model)), True)
+        super().__init__(partial(LPIPS._loss, lpips.LPIPS(net=model).to(0)), True)
 
 
 class Dice(Criterion):
 
-    def __init__(self, smooth: float = 1e-6) -> None:
+    def __init__(self, labels: list[int] | None = None) -> None:
         super().__init__()
-        self.smooth = smooth
+        self.labels = labels
 
     def flatten(self, tensor: torch.Tensor) -> torch.Tensor:
         return tensor.permute((1, 0) + tuple(range(2, tensor.dim()))).contiguous().view(tensor.size(1), -1)
@@ -174,26 +180,26 @@ class Dice(Criterion):
     def dice_per_channel(self, tensor: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         tensor = self.flatten(tensor)
         target = self.flatten(target)
-        return (2.0 * (tensor * target).sum() + self.smooth) / (tensor.sum() + target.sum() + self.smooth)
+        return (2.0 * (tensor * target).sum() + 1e-6) / (tensor.sum() + target.sum() + 1e-6)
 
     def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         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)
-        )
-        return 1 - torch.mean(self.dice_per_channel(output, target))
+        result = {}
+        loss = torch.tensor(0, dtype=torch.float32).to(output.device)
+        labels = self.labels if self.labels is not None else torch.unique(target)
+        for label in labels:
+            tp = target == label
+            if tp.any().item():
+                if output.shape[1] > 1:
+                    pp = output[:, label].unsqueeze(1)
+                else:
+                    pp = output == label
+                loss_tmp = self.dice_per_channel(pp.float(), tp.float())
+                loss += loss_tmp
+                result[label] = loss_tmp.item()
+            else:
+                result[label] = np.nan
+        return 1 - loss / len(labels), result
 
 
 class GradientImages(Criterion):
@@ -621,23 +627,14 @@ class MutualInformationLoss(torch.nn.Module):
 class IMPACTSynth(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]},
+        self, model_name: str, shape: list[int] = [0, 0], in_channels: int = 1, weights: list[float] = [1]
     ) -> None:
         super().__init__()
         if model_name is None:
             return
         self.in_channels = in_channels
-        self.losses: dict[torch.nn.Module, list[float]] = {}
-        for loss, weights in losses.items():
-            module, name = get_module(loss, "konfai.metric.measure")
-            self.losses[
-                config(os.environ["KONFAI_CONFIG_PATH"])(getattr(importlib.import_module(module), name))(config=None)
-            ] = weights
-
+        self.loss = torch.nn.L1Loss()
+        self.weights = weights
         self.model_path = download_url(
             model_name,
             "https://huggingface.co/VBoussot/impact-torchscript-models/resolve/main/",
@@ -648,19 +645,17 @@ class IMPACTSynth(Criterion):  # Feature-Oriented Comparison for Unpaired Synthe
         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)))
+            dummy_input = torch.zeros((1, self.in_channels, *(self.shape if self.shape else [224] * self.dim))).to(0)
             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)})."
-                )
+            if len(weights) != len(out):
+                raise ValueError(f"Mismatch between number of weights ({len(weights)}) 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"Expected output length: {len(weights)}\n"
                 f"Error: {type(e).__name__}: {e}"
             )
             raise RuntimeError(msg) from e
@@ -682,12 +677,14 @@ class IMPACTSynth(Criterion):  # Feature-Oriented Comparison for Unpaired Synthe
         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]
+        for zipped_output in zip(self.weights, self.model(output), *[self.model(target) for target in targets]):
+            weight = zipped_output[0]
+            if weight == 0:
+                continue
+            output_feature = zipped_output[1]
             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 += weight * loss_function(output_feature, target_feature)
+            for target_feature in targets_features:
+                loss += weight * self.loss(output_feature, target_feature)
         return loss
 
     def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
@@ -696,7 +693,7 @@ class IMPACTSynth(Criterion):  # Feature-Oriented Comparison for Unpaired Synthe
         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 += self._compute(output[:, :, i, ...], [t[:, :, i, ...] for t in targets])
             loss /= output.shape[2]
         else:
             loss = self._compute(output, list(targets))

konfai/network/network.py

@@ -188,11 +188,16 @@ class Measure:
         def reset_loss(self) -> None:
             self._loss.clear()
 
-        def add(self, weight: float, value: torch.Tensor) -> None:
-            if self.is_loss or value != 0:
-                self._loss.append(value if self.is_loss else value.detach())
-                self._values.append(value.item())
-                self._weight.append(weight)
+        def add(self, weight: float, value: torch.Tensor | tuple[torch.Tensor, float]) -> None:
+            if isinstance(value, tuple):
+                loss_value, true_value = value
+            else:
+                loss_value = value
+                true_value = value.item()
+
+            self._loss.append(loss_value if self.is_loss else loss_value.detach())
+            self._values.append(true_value)
+            self._weight.append(weight)
 
         def get_last_loss(self) -> torch.Tensor:
             return self._loss[-1] * self._weight[-1] if len(self._loss) else torch.zeros((1), requires_grad=True)
@@ -239,13 +244,13 @@ class Measure:
                 for target_group_tmp in target_group.split(";"):
                     if target_group_tmp not in group_dest:
                         raise MeasureError(
-                            f"The target_group '{target_group_tmp}' defined in "
+                            f"The target_group {target_group_tmp} defined in "
                             "'outputs_criterions.{output_group}.targets_criterions'"
                             " was not found in the available destination groups.",
                             "This target_group is expected for loss or metric computation, "
                             "but was not loaded in 'group_dest'.",
-                            f"Please make sure that the group '{target_group_tmp}' is defined in "
-                            "'Dataset:groups_src:...:groups_dest:'{target_group_tmp}'' "
+                            f"Please make sure that the group {target_group_tmp} is defined in "
+                            "Dataset:groups_src:...:groups_dest: {target_group_tmp} "
                             "and correctly loaded from the dataset.",
                         )
                 for criterion in self.outputs_criterions[output_group][target_group]:
@@ -868,7 +873,10 @@ class Network(ModuleArgsDict, ABC):
                     init_gain=self.init_gain,
                 )
             )
-        name = "Model" + ("_EMA" if ema else "")
+        name = "Model"
+        if ema:
+            if name + "_EMA" in state_dict:
+                name += "_EMA"
         if name in state_dict:
             value = state_dict[name]
             model_state_dict_tmp = {}
@@ -893,7 +901,9 @@ class Network(ModuleArgsDict, ABC):
                     model_state_dict[alias] = model_state_dict_tmp[alias]
             self.load_state_dict(model_state_dict)
         if f"{self.get_name()}_optimizer_state_dict" in state_dict and self.optimizer:
+            last_lr = self.optimizer.param_groups[0]["lr"]
             self.optimizer.load_state_dict(state_dict[f"{self.get_name()}_optimizer_state_dict"])
+            self.optimizer.param_groups[0]["lr"] = last_lr
         if f"{self.get_name()}_it" in state_dict:
             _it = state_dict.get(f"{self.get_name()}_it")
             if isinstance(_it, int):

konfai/predictor.py

@@ -793,7 +793,7 @@ class Predictor(DistributedObject):
             exit(0)
 
         self.size = len(self.gpu_checkpoints) + 1 if self.gpu_checkpoints else 1
-        self.dataloader = self.dataset.get_data(world_size // self.size)
+        self.dataloader, _, _ = self.dataset.get_data(world_size // self.size)
         for name, output_dataset in self.outputs_dataset.items():
             output_dataset.load(
                 name.replace(".", ":"),

konfai/trainer.py

@@ -14,7 +14,6 @@ from konfai import (
     config_file,
     current_date,
     konfai_state,
-    models_directory,
     path_to_models,
     setups_directory,
     statistics_directory,
@@ -147,7 +146,6 @@ class _Trainer:
         self.global_rank = global_rank
         self.local_rank = local_rank
         self.size = size
-
         self.save_checkpoint_mode = save_checkpoint_mode
         self.train_name = train_name
         self.epochs = epochs
@@ -186,6 +184,9 @@ class _Trainer:
         self.dataloader_training.dataset.load("Train")
         if self.dataloader_validation is not None:
             self.dataloader_validation.dataset.load("Validation")
+            if State[konfai_state()] != State.TRAIN:
+                self._validate()
+
         with tqdm.tqdm(
             iterable=range(self.epoch, self.epochs),
             leave=False,
@@ -604,38 +605,7 @@ class Trainer(DistributedObject):
         Serializes model and EMA model (if any) in both state_dict and full model formats.
         Also saves optimizer states and YAML config snapshot.
         """
-        path_checkpoint = checkpoints_directory() + self.name + "/"
-        path_model = models_directory() + self.name + "/"
-        if os.path.exists(path_checkpoint) and os.listdir(path_checkpoint):
-            for directory in [
-                path_model,
-                f"{path_model}Serialized/",
-                f"{path_model}StateDict/",
-            ]:
-                if not os.path.exists(directory):
-                    os.makedirs(directory)
-
-            for name in sorted(os.listdir(path_checkpoint)):
-                checkpoint = torch.load(path_checkpoint + name, weights_only=True, map_location="cpu")
-                self.model.load(checkpoint, init=False, ema=False)
-
-                torch.save(self.model, f"{path_model}Serialized/{name}")
-                torch.save(
-                    {"Model": self.model.state_dict()},
-                    f"{path_model}StateDict/{name}",
-                )
-
-                if self.model_ema is not None:
-                    self.model_ema.module.load(checkpoint, init=False, ema=True)
-                    torch.save(
-                        self.model_ema.module,
-                        f"{path_model}Serialized/{current_date()}_EMA.pt",
-                    )
-                    torch.save(
-                        {"Model_EMA": self.model_ema.module.state_dict()},
-                        f"{path_model}StateDict/{current_date()}_EMA.pt",
-                    )
-
+        if os.path.exists(self.config_namefile):
             os.rename(
                 self.config_namefile,
                 self.config_namefile.replace(".yml", "") + "_" + str(self.it) + ".yml",
@@ -662,6 +632,7 @@ class Trainer(DistributedObject):
             world_size (int): Total number of distributed processes.
         """
         state = State[konfai_state()]
+        print(checkpoints_directory() + self.name + "/")
         if state != State.RESUME and os.path.exists(checkpoints_directory() + self.name + "/"):
             if os.environ["KONFAI_OVERWRITE"] != "True":
                 accept = input(f"The model {self.name} already exists ! Do you want to overwrite it (yes,no) : ")
@@ -669,7 +640,6 @@ class Trainer(DistributedObject):
                     return
             for directory_path in [
                 statistics_directory(),
-                models_directory(),
                 checkpoints_directory(),
                 setups_directory(),
             ]:
@@ -698,7 +668,13 @@ class Trainer(DistributedObject):
             os.makedirs(setups_directory() + self.name + "/")
         shutil.copyfile(self.config_namefile_src + ".yml", self.config_namefile)
 
-        self.dataloader = self.dataset.get_data(world_size // self.size)
+        self.dataloader, train_names, validation_names = self.dataset.get_data(world_size // self.size)
+        with open(setups_directory() + self.name + "/Train_" + str(self.it) + ".txt", "w") as f:
+            for name in train_names:
+                f.write(name + "\n")
+        with open(setups_directory() + self.name + "/Validation_" + str(self.it) + ".txt", "w") as f:
+            for name in validation_names:
+                f.write(name + "\n")
 
     def run_process(
         self,

konfai/utils/utils.py

@@ -41,7 +41,7 @@ def description(model, model_ema=None, show_memory: bool = True, train: bool = T
             "("
             + " ".join(
                 [
-                    f"{name}({network.optimizer.param_groups[0]['lr']:.6f} if network.optimizer else 0) : "
+                    f"{name}({(network.optimizer.param_groups[0]['lr'] if network.optimizer else 0):.6f}) : "
                     + " ".join(
                         f"{k.split(':')[-1]}({w:.2f}) : {v:.6f}"
                         for (k, v), w in zip(
@@ -58,11 +58,11 @@ def description(model, model_ema=None, show_memory: bool = True, train: bool = T
     model_loss_desc = loss_desc(model)
     result = ""
     if len(model_loss_desc) > 2:
-        f"Loss {model_loss_desc} "
+        result += f"Loss {model_loss_desc} "
     if model_ema is not None:
         model_ema_loss_desc = loss_desc(model_ema)
         if len(model_ema_loss_desc) > 2:
-            result += f"Loss EMA {loss_desc(model_ema_loss_desc)} "
+            result += f"Loss EMA {model_ema_loss_desc} "
     result += gpu_info()
     if show_memory:
         result += f" | {get_memory_info()}"
@@ -765,6 +765,8 @@ def download_url(model_name: str, url: str) -> str:
     if not isinstance(locations, list) or not locations:
         raise ImportError("No valid submodule_search_locations found")
     base_path = Path(locations[0]) / "metric" / "models"
+    os.makedirs(base_path, exist_ok=True)
+
     subdirs = Path(model_name).parent
     model_dir = base_path / subdirs
     model_dir.mkdir(exist_ok=True)

{konfai-1.1.9.dist-info → konfai-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.9
+Version: 1.2.0
 Summary: Modular and configurable Deep Learning framework with YAML and PyTorch
 Author-email: Valentin Boussot <boussot.v@gmail.com>
 License-Expression: Apache-2.0

{konfai-1.1.9.dist-info → konfai-1.2.0.dist-info}/RECORD

@@ -1,15 +1,15 @@
 konfai/__init__.py,sha256=qjE9Rqxo1sMrkqGS8I5xlGQMZnjIfU-CGgSI5Wmbmbs,1231
-konfai/evaluator.py,sha256=jsRzVSFjK-V1rZVK9kmN0Gh5-F2JhJrJv291UGNm8CM,16736
+konfai/evaluator.py,sha256=xAKWUDvdSxqYRUsKqH6ieQF06LWa785aE4zLv4I3_i4,17850
 konfai/main.py,sha256=Fc4HcJEhPmgunj_f-QYyvQNvjHrKHSUv27Okgu6V5_A,3842
-konfai/predictor.py,sha256=k9S-AH-wGvmr4YQF2IczJ2Nb9_aTZwNd9f6iu4s9v78,34591
-konfai/trainer.py,sha256=wNHgDh0LtxTi0-aWCkT90hjjpJFaX_zRWyA5esVrsLY,28072
+konfai/predictor.py,sha256=a8QbZPMKrf8bFz1XKZs7j6HT0RXC09HdiIQbvUBqKBM,34597
+konfai/trainer.py,sha256=4mc-8r-FxtX_EAn2su8qd-BLLQ0D0So8hh5rmuf6Hqs,27163
 konfai/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 konfai/data/augmentation.py,sha256=vcJE7mosvUkwbpbTN_lGP0S1uJrJYGjlLrt9VnDdJYY,27792
-konfai/data/data_manager.py,sha256=FOxpmL56-Cgqsms2rdPk1NML5OiuuIQW49-G0j3O2Os,31564
+konfai/data/data_manager.py,sha256=7-ruYS2HJMvLzTq5425loOf2OCpAm5MspdYoO74ICsw,31180
 konfai/data/patching.py,sha256=jS35OxnJagKNUnJu7TzuGZpVj9fP-6H4nc2OEYOGgt8,16494
-konfai/data/transform.py,sha256=MmA1vgXAj1V-e-8RNa1XAaHbWI5i85NoVyCp7yZ2-kg,26816
+konfai/data/transform.py,sha256=YCldsqTTBFFCqc_VdvyuNVs2kmV56CxQBN5XhEoPxho,27745
 konfai/metric/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-konfai/metric/measure.py,sha256=7CHpLWCuuNNk3cYvGiNql6nt_pq6PA5SHIVNiSpeolk,28144
+konfai/metric/measure.py,sha256=0mOIZKTa2u0UECpoDSbdJUhttAw_e1BlsROQQpi1oBk,27804
 konfai/metric/schedulers.py,sha256=TpYMA24FMpxRnqfhMGb0i_Mm-bzT9kySbBgvkYk-6wM,1327
 konfai/models/classification/convNeXt.py,sha256=Ha9QYd1-JEYUwL7zgNNAWeuJLbUT7LCNUkLcsdySAHM,10060
 konfai/models/classification/resnet.py,sha256=4-l7HtpP_OqApDT8XjTH94fXIuiSzz833SUWXP2oFJo,10813
@@ -24,15 +24,15 @@ konfai/models/segmentation/NestedUNet.py,sha256=W4uauwF0HY8ybi49iYiTlKLdJEyD7SaC
 konfai/models/segmentation/UNet.py,sha256=Pu_LiQdO4Mrzyn0HRE6rwxUjHGH4OG-JpzWB_U1K46g,5602
 konfai/network/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 konfai/network/blocks.py,sha256=l70_oOcz5Hmol2xmxruG0kke_2SVgO3rXYXVTMSdAS8,15645
-konfai/network/network.py,sha256=FPv3i3pnyThFMd1KgMk3wNz41btE5_Kby95dUwG7PsM,54458
+konfai/network/network.py,sha256=PWI6W4sz7G5Pbb-79l7mL61AoyHWlQytRtBzvwh3Ro0,54800
 konfai/utils/ITK.py,sha256=HVed4Z96X1jTaWrrQNdoBMqOtVK9InAPlDBJu-5uv3g,15476
 konfai/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 konfai/utils/config.py,sha256=a7t44CYMUT5oCDdjL94IswhCVfFbQ5FCgDWZktDDkc4,14347
 konfai/utils/dataset.py,sha256=Au22fcADKyDJMfS8Z9q8kEXLtKkoufJsH7Pwly6pALo,28288
-konfai/utils/utils.py,sha256=AtPWHJh_RdAGK2m9Cv3BXvUZTpspVHRTPXXvRfn5dZg,28366
-konfai-1.1.9.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-konfai-1.1.9.dist-info/METADATA,sha256=6PaWz831mjixz5kTN_-wBtDNGVBk0SjwayD4AGyUC5o,2451
-konfai-1.1.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-konfai-1.1.9.dist-info/entry_points.txt,sha256=fG82HRN5-g39ACSOCtij_I3N6EHxfYnMR0D7TI_8pW8,81
-konfai-1.1.9.dist-info/top_level.txt,sha256=xF470dkIlFoFqTZEOlRehKJr4WU_8OKGXrJqYm9vWKs,7
-konfai-1.1.9.dist-info/RECORD,,
+konfai/utils/utils.py,sha256=jCj3tZ8agQYceSY_tlVYp88UFPE5oUn6tXrqnZGrKiI,28410
+konfai-1.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+konfai-1.2.0.dist-info/METADATA,sha256=0i5f81WQyblGv3hF9L4l-skjcPmx1NK7NwtPmcdAQIY,2451
+konfai-1.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+konfai-1.2.0.dist-info/entry_points.txt,sha256=fG82HRN5-g39ACSOCtij_I3N6EHxfYnMR0D7TI_8pW8,81
+konfai-1.2.0.dist-info/top_level.txt,sha256=xF470dkIlFoFqTZEOlRehKJr4WU_8OKGXrJqYm9vWKs,7
+konfai-1.2.0.dist-info/RECORD,,