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

konfai/utils/utils.py

@@ -1,119 +1,163 @@
+import argparse
+import importlib.util
 import itertools
-import pynvml
-import psutil
-
-import numpy as np
 import os
-import torch
-
+import random
+import re
+import shutil
+import socket
+import subprocess  # nosec B404
+import sys
 from abc import ABC, abstractmethod
+from contextlib import closing
 from enum import Enum
-from typing import Any, Union
+from pathlib import Path
+from typing import Any, TextIO, cast
 
-from konfai import CONFIG_FILE, EVALUATIONS_DIRECTORY, STATISTICS_DIRECTORY, PREDICTIONS_DIRECTORY, KONFAI_STATE, CUDA_VISIBLE_DEVICES
+import numpy as np
+import psutil
+import pynvml
+import requests
+import torch
 import torch.distributed as dist
-import argparse
-import subprocess
-import random
+import torch.nn.functional as F  # noqa: N812
 from torch.utils.data import DataLoader
-import torch.nn.functional as F
-import sys
-import re
-
-import requests
+from torch.utils.tensorboard.writer import SummaryWriter
 from tqdm import tqdm
-import importlib
-from pathlib import Path
-import shutil
 
-def description(model, modelEMA = None, showMemory: bool = True, train: bool = True) -> str:
-    values_desc = lambda weights, values: " ".join(["{}({:.2f}) : {:.6f}".format(name.split(":")[-1], weight, value) for (name, value), weight in zip(values.items(), weights.values())])
-    model_desc = lambda model : "("+" ".join(["{}({:.6f}) : {}".format(name, network.optimizer.param_groups[0]['lr'] if network.optimizer is not None else 0, values_desc(network.measure.getLastWeights(), network.measure.getLastValues())) for name, network in model.module.getNetworks().items() if network.measure is not None])+")"
-    result = "Loss {}".format(model_desc(model))
-    if modelEMA is not None:
-        result += " Loss EMA {}".format(model_desc(modelEMA))
-    result += " "+gpuInfo()
-    if showMemory:
-        result +=" | {}".format(memoryInfo())
+from konfai import (
+    config_file,
+    cuda_visible_devices,
+    evaluations_directory,
+    konfai_state,
+    predictions_directory,
+    statistics_directory,
+)
+
+
+def description(model, model_ema=None, show_memory: bool = True, train: bool = True) -> str:
+    def loss_desc(model):
+        return (
+            "("
+            + " ".join(
+                [
+                    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(
+                            network.measure.get_last_values().items(), network.measure.get_last_weights().values()
+                        )
+                    )
+                    for name, network in model.module.get_networks().items()
+                    if network.measure is not None
+                ]
+            )
+            + ")"
+        )
+
+    model_loss_desc = loss_desc(model)
+    result = ""
+    if len(model_loss_desc) > 2:
+        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 {model_ema_loss_desc} "
+    result += gpu_info()
+    if show_memory:
+        result += f" | {get_memory_info()}"
     return result
 
-def _getModule(classpath : str, type : str) -> tuple[str, str]:
+
+def get_module(classpath: str, default_classpath: str) -> tuple[str, str]:
     if len(classpath.split(":")) > 1:
         module = ".".join(classpath.split(":")[:-1])
         name = classpath.split(":")[-1]
     else:
-        module = type+("." if len(classpath.split(".")) > 2 else "")+".".join(classpath.split(".")[:-1])
+        module = (
+            default_classpath + ("." if len(classpath.split(".")) > 2 else "") + ".".join(classpath.split(".")[:-1])
+        )
         name = classpath.split(".")[-1]
     return module, name.split("/")[0]
 
-def cpuInfo() -> str:
-    return "CPU ({:.2f} %)".format(psutil.cpu_percent(interval=0.5))
 
-def memoryInfo() -> str:
-    return "Memory ({:.2f}G ({:.2f} %))".format(psutil.virtual_memory()[3]/2**30, psutil.virtual_memory()[2])
+def get_cpu_info() -> str:
+    return f"CPU ({psutil.cpu_percent(interval=0.5):.2f} %)"
+
+
+def get_memory_info() -> str:
+    return f"Memory ({psutil.virtual_memory()[3] / 2**30:.2f}G ({psutil.virtual_memory()[2]:.2f} %))"
+
 
-def getMemory() -> float:
-    return psutil.virtual_memory()[3]/2**30
+def get_memory() -> float:
+    return psutil.virtual_memory()[3] / 2**30
 
-def memoryForecast(memory_init : float, i : float, size : float) -> str:
-    current_memory = getMemory()
-    forecast = memory_init + ((current_memory-memory_init)*size/i) if i > 0 else 0
-    return "Memory forecast ({:.2f}G ({:.2f} %))".format(forecast, forecast/(psutil.virtual_memory()[0]/2**30)*100)
 
-def gpuInfo() -> str:
-    if CUDA_VISIBLE_DEVICES() == "":
+def memory_forecast(memory_init: float, i: float, size: float) -> str:
+    current_memory = get_memory()
+    forecast = memory_init + ((current_memory - memory_init) * size / i) if i > 0 else 0
+    return f"Memory forecast ({forecast:.2f}G ({forecast / (psutil.virtual_memory()[0] / 2**30) * 100:.2f} %))"
+
+
+def gpu_info() -> str:
+    if cuda_visible_devices() == "":
         return ""
-    
-    devices = [int(i) for i in CUDA_VISIBLE_DEVICES().split(",")]
+
+    devices = [int(i) for i in cuda_visible_devices().split(",")]
     device = devices[0]
-    
+
     if device < pynvml.nvmlDeviceGetCount():
         handle = pynvml.nvmlDeviceGetHandleByIndex(device)
         memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
     else:
         return ""
-    node_name = "Node: {} " +os.environ["SLURMD_NODENAME"] if "SLURMD_NODENAME" in os.environ else ""
-    return  "{}GPU({}) Memory GPU ({:.2f}G ({:.2f} %))".format(node_name, devices, float(memory.used)/(10**9), float(memory.used)/float(memory.total)*100)
+    node_name = "Node: {} " + os.environ["SLURMD_NODENAME"] if "SLURMD_NODENAME" in os.environ else ""
+    return f"{node_name}GPU({devices}) Memory GPU ({memory.used / 1e9:.2f}G ({memory.used / memory.total * 100:.2f} %))"
+
 
-def getMaxGPUMemory(device : Union[int, torch.device]) -> float:
+def get_max_gpu_memory(device: int | torch.device) -> float:
     if isinstance(device, torch.device):
         if str(device).startswith("cuda:"):
             device = int(str(device).replace("cuda:", ""))
         else:
             return 0
-    device = [int(i) for i in CUDA_VISIBLE_DEVICES().split(",")][device]
+    device = [int(i) for i in cuda_visible_devices().split(",")][device]
     if device < pynvml.nvmlDeviceGetCount():
         handle = pynvml.nvmlDeviceGetHandleByIndex(device)
         memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
     else:
         return 0
-    return float(memory.total)/(10**9)
+    return float(memory.total) / (10**9)
 
-def getGPUMemory(device : Union[int, torch.device]) -> float:
+
+def get_gpu_memory(device: int | torch.device) -> float:
     if isinstance(device, torch.device):
         if str(device).startswith("cuda:"):
             device = int(str(device).replace("cuda:", ""))
         else:
             return 0
-    device = [int(i) for i in CUDA_VISIBLE_DEVICES().split(",")][device]
+    device = [int(i) for i in cuda_visible_devices().split(",")][device]
     if device < pynvml.nvmlDeviceGetCount():
         handle = pynvml.nvmlDeviceGetHandleByIndex(device)
         memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
     else:
         return 0
-    return float(memory.used)/(10**9)
+    return float(memory.used) / (10**9)
+
 
-class NeedDevice(ABC):
+class NeedDevice:
 
     def __init__(self) -> None:
         super().__init__()
-        self.device : torch.device
-    
-    def setDevice(self, device : int):
-        self.device = getDevice(device)
+        self.device: torch.device
+
+    def to(self, device: int):
+        self.device = get_device(device)
+
+
+def get_device(device: int):
+    return device if torch.cuda.is_available() and device >= 0 else torch.device("cpu")
 
-def getDevice(device : int):
-    return device if torch.cuda.is_available() and device >=0 else torch.device("cpu")
 
 class State(Enum):
     TRAIN = "TRAIN"
@@ -122,13 +166,18 @@ class State(Enum):
     FINE_TUNING = "FINE_TUNING"
     PREDICTION = "PREDICTION"
     EVALUATION = "EVALUATION"
-    
+
     def __str__(self) -> str:
         return self.value
 
-def get_patch_slices_from_nb_patch_per_dim(patch_size_tmp: list[int], nb_patch_per_dim : list[tuple[int, bool]], overlap: Union[int, None]) -> list[tuple[slice]]:
+
+def get_patch_slices_from_nb_patch_per_dim(
+    patch_size_tmp: list[int],
+    nb_patch_per_dim: list[tuple[int, bool]],
+    overlap: int | None,
+) -> list[tuple[slice, ...]]:
     patch_slices = []
-    slices : list[list[slice]] = []
+    slices: list[list[slice]] = []
     if overlap is None:
         overlap = 0
     patch_size = []
@@ -138,47 +187,54 @@ def get_patch_slices_from_nb_patch_per_dim(patch_size_tmp: list[int], nb_patch_p
             patch_size.append(1)
         else:
             patch_size.append(patch_size_tmp[i])
-            i+=1
+            i += 1
 
     for dim, nb in enumerate(nb_patch_per_dim):
         slices.append([])
         for index in range(nb[0]):
-            start = (patch_size[dim]-overlap)*index
+            start = (patch_size[dim] - overlap) * index
             end = start + patch_size[dim]
-            slices[dim].append(slice(start,end))
+            slices[dim].append(slice(start, end))
     for chunk in itertools.product(*slices):
         patch_slices.append(tuple(chunk))
     return patch_slices
 
-def get_patch_slices_from_shape(patch_size: list[int], shape : list[int], overlap: Union[int, None]) -> tuple[list[tuple[slice]], list[tuple[int, bool]]]:
+
+def get_patch_slices_from_shape(
+    patch_size: list[int], shape: list[int], overlap_tmp: int | None
+) -> tuple[list[tuple[slice, ...]], list[tuple[int, bool]]]:
     if len(shape) != len(patch_size):
-         raise DatasetManagerError(
-                f"Dimension mismatch: 'patch_size' has {len(patch_size)} dimensions, but 'shape' has {len(shape)}.",
-                f"patch_size: {patch_size}",
-                f"shape: {shape}",
-                "Both must have the same number of dimensions (e.g., 3D patch for 3D volume)."
-            )
+        raise DatasetManagerError(
+            f"Dimension mismatch: 'patch_size' has {len(patch_size)} dimensions, but 'shape' has {len(shape)}.",
+            f"patch_size: {patch_size}",
+            f"shape: {shape}",
+            "Both must have the same number of dimensions (e.g., 3D patch for 3D volume).",
+        )
     patch_slices = []
     nb_patch_per_dim = []
-    slices : list[list[slice]] = []
-    if overlap is None:
-        size = [np.ceil(a/b) for a, b in zip(shape, patch_size)]
+    slices: list[list[slice]] = []
+    if overlap_tmp is None:
+        size = [np.ceil(a / b) for a, b in zip(shape, patch_size)]
         tmp = np.zeros(len(size), dtype=np.int_)
         for i, s in enumerate(size):
             if s > 1:
-                tmp[i] = np.mod(patch_size[i]-np.mod(shape[i], patch_size[i]), patch_size[i])//(size[i]-1)
+                tmp[i] = np.mod(patch_size[i] - np.mod(shape[i], patch_size[i]), patch_size[i]) // (size[i] - 1)
         overlap = tmp
     else:
-        overlap = [overlap if size > 1 else 0 for size in patch_size]
-    
+        overlap = [overlap_tmp if size > 1 else 0 for size in patch_size]
+
     for dim in range(len(shape)):
-        assert overlap[dim] < patch_size[dim],  "Overlap must be less than patch size"
+        if overlap[dim] >= patch_size[dim]:
+            raise ValueError(
+                f"Overlap must be less than patch size, got overlap={overlap[dim]}",
+                f" ≥ patch_size={patch_size[dim]} at dim={dim}",
+            )
 
     for dim in range(len(shape)):
         slices.append([])
         index = 0
         while True:
-            start = (patch_size[dim]-overlap[dim])*index
+            start = (patch_size[dim] - overlap[dim]) * index
 
             end = start + patch_size[dim]
             if end >= shape[dim]:
@@ -187,84 +243,107 @@ def get_patch_slices_from_shape(patch_size: list[int], shape : list[int], overla
                 break
             slices[dim].append(slice(start, end))
             index += 1
-        nb_patch_per_dim.append((index+1, patch_size[dim] == 1))
+        nb_patch_per_dim.append((index + 1, patch_size[dim] == 1))
 
     for chunk in itertools.product(*slices):
         patch_slices.append(tuple(chunk))
-    
+
     return patch_slices, nb_patch_per_dim
 
-def _logSignalFormat(input : np.ndarray):
-    return {str(i): channel for i, channel in enumerate(input)}
-
-def _logImageFormat(input : np.ndarray):
-    if len(input.shape) == 2:
-        input = np.expand_dims(input, axis=0)
-
-    if len(input.shape) == 3 and input.shape[0] != 1:
-        input = np.expand_dims(input, axis=0)
-    if len(input.shape) == 4:
-        input = input[:, input.shape[1]//2]
-        
-    input = input.astype(float)
-    b = -np.min(input)
-    if (np.max(input)+b) > 0:
-        return (input+b)/(np.max(input)+b)
+
+def _log_signal_format(array: np.ndarray) -> dict[str, np.ndarray]:
+    return {str(i): channel for i, channel in enumerate(array)}
+
+
+def _log_image_format(array: np.ndarray) -> np.ndarray:
+    if len(array.shape) == 2:
+        array = np.expand_dims(array, axis=0)
+
+    if len(array.shape) == 3 and array.shape[0] != 1:
+        array = np.expand_dims(array, axis=0)
+    if len(array.shape) == 4:
+        array = array[:, array.shape[1] // 2]
+
+    array = array.astype(float)
+    b = -np.min(array)
+    if (np.max(array) + b) > 0:
+        return (array + b) / (np.max(array) + b)
     else:
-        return 0*input
+        return 0 * array
+
 
-def _logImagesFormat(input : np.ndarray):
+def _log_images_format(array: np.ndarray) -> np.ndarray:
     result = []
-    for n in range(input.shape[0]):
-        result.append(_logImageFormat(input[n]))
+    for n in range(array.shape[0]):
+        result.append(_log_image_format(array[n]))
     result = np.stack(result, axis=0)
     return result
 
-def _logVideoFormat(input : np.ndarray):
-    result = []
-    for t in range(input.shape[1]):
-        result.append( _logImagesFormat(input[:, t,...]))
-    result = np.stack(result, axis=1)
+
+def _log_video_format(array: np.ndarray) -> np.ndarray:
+    result_list = []
+    for t in range(array.shape[1]):
+        result_list.append(_log_images_format(array[:, t, ...]))
+    result = np.stack(result_list, axis=1)
 
     nb_channel = result.shape[2]
     if nb_channel < 3:
         channel_split = [result[:, :, 0, ...] for i in range(3)]
     else:
         channel_split = np.split(result, 3, axis=0)
-    input = np.zeros((result.shape[0], result.shape[1], 3, *list(result.shape[3:])))
+    array = np.zeros((result.shape[0], result.shape[1], 3, *list(result.shape[3:])))
     for i, channels in enumerate(channel_split):
-        input[:,:,i] = np.mean(channels, axis=0)
-    return input
+        array[:, :, i] = np.mean(channels, axis=0)
+    return array
+
 
 class DataLog(Enum):
-    SIGNAL   = lambda tb, name, layer, it : [tb.add_scalars(name, _logSignalFormat(layer[b, :, 0]), layer.shape[0]*it+b) for b in range(layer.shape[0])],
-    IMAGE   = lambda tb, name, layer, it : tb.add_image(name, _logImageFormat(layer[0]), it),
-    IMAGES  = lambda tb, name, layer, it : tb.add_images(name, _logImagesFormat(layer), it),
-    VIDEO   = lambda tb, name, layer, it : tb.add_video(name, _logVideoFormat(layer), it),
-    AUDIO   = lambda tb, name, layer, it : tb.add_audio(name, _logImageFormat(layer), it)
+    SIGNAL = "SIGNAL"
+    IMAGE = "IMAGE"
+    IMAGES = "IMAGES"
+    VIDEO = "VIDEO"
+    AUDIO = "AUDIO"
+
+    def __call__(self, tb: SummaryWriter, name: str, layer: torch.Tensor, it: int):
+        if self == DataLog.SIGNAL:
+            return [
+                tb.add_scalars(name, _log_signal_format(layer[b, :, 0]), layer.shape[0] * it + b)
+                for b in range(layer.shape[0])
+            ]
+        elif self == DataLog.IMAGE:
+            return tb.add_image(name, _log_image_format(layer[0]), it)
+        elif self == DataLog.IMAGES:
+            return tb.add_images(name, _log_images_format(layer), it)
+        elif self == DataLog.VIDEO:
+            return tb.add_video(name, _log_video_format(layer), it)
+        elif self == DataLog.AUDIO:
+            return tb.add_audio(name, _log_image_format(layer), it)
+        else:
+            raise ValueError(f"Unsupported DataLog type: {self}")
+
 
 class Log:
     def __init__(self, name: str, rank: int) -> None:
-        if KONFAI_STATE() == "PREDICTION":
-            path = PREDICTIONS_DIRECTORY()
-        elif KONFAI_STATE() == "EVALUATION":
-            path = EVALUATIONS_DIRECTORY()
+        if konfai_state() == "PREDICTION":
+            path = predictions_directory()
+        elif konfai_state() == "EVALUATION":
+            path = evaluations_directory()
         else:
-            path = STATISTICS_DIRECTORY()
-        
+            path = statistics_directory()
+
         self.verbose = os.environ.get("KONFAI_VERBOSE", "True") == "True"
         self.log_path = os.path.join(path, name)
         os.makedirs(self.log_path, exist_ok=True)
         self.rank = rank
-        self.file = open(os.path.join(self.log_path, "log_{}.txt".format(rank)), "w", buffering=1) 
+        self.file = open(os.path.join(self.log_path, f"log_{rank}.txt"), "w", buffering=1)
         self.stdout_bak = sys.stdout
         self.stderr_bak = sys.stderr
         self._buffered_line = ""
-    
+
     def __enter__(self):
         self.file.__enter__()
-        sys.stdout = self
-        sys.stderr = self
+        sys.stdout = cast(TextIO, self)
+        sys.stderr = cast(TextIO, self)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -275,25 +354,22 @@ class Log:
     def write(self, msg: str):
         if not msg:
             return
-        
-        
-        ANSI_ESCAPE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]')
-        CARRIAGE_RETURN = re.compile(r'(?:\r|\x1b\[A).*')
-        msg_clean = ANSI_ESCAPE.sub('', msg)
-        if '\r' in msg_clean or '' in msg:
-            # On garde seulement le contenu après le dernier retour chariot
-            msg_clean = msg_clean.split('\r')[-1].strip()
+
+        ansi_escape = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]")
+        msg_clean = ansi_escape.sub("", msg)
+        if "\r" in msg_clean or "" in msg:
+            msg_clean = msg_clean.split("\r")[-1].strip()
             self._buffered_line = msg_clean
         else:
             self._buffered_line = msg_clean.strip()
 
         if self._buffered_line:
-            # Écrit dans le fichier
             self.file.write(self._buffered_line + "\n")
             self.file.flush()
         if self.verbose and (self.rank == 0 or "KONFAI_CLUSTER" in os.environ):
-            sys.__stdout__.write(msg)
-            sys.__stdout__.flush()
+            if sys.__stdout__ is not None:
+                sys.__stdout__.write(msg)
+                sys.__stdout__.flush()
 
     def flush(self):
         self.file.flush()
@@ -302,65 +378,104 @@ class Log:
         return False
 
     def fileno(self):
+        if sys.__stdout__ is None:
+            raise RuntimeError("sys.__stdout__ is None, cannot get fileno")
         return sys.__stdout__.fileno()
-    
-class TensorBoard():
-    
+
+
+class TensorBoard:
+
     def __init__(self, name: str) -> None:
-        self.process = None
+        self.process: subprocess.Popen | None = None
         self.name = name
 
     def __enter__(self):
         if "KONFAI_TENSORBOARD_PORT" in os.environ:
-            command = ["tensorboard", "--logdir", PREDICTIONS_DIRECTORY() if KONFAI_STATE() == "PREDICTION" else STATISTICS_DIRECTORY() + self.name + "/", "--port", os.environ["KONFAI_TENSORBOARD_PORT"], "--bind_all"]
-            self.process = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+            tensorboard_exe = shutil.which("tensorboard")
+            if tensorboard_exe is None:
+                raise RuntimeError("TensorBoard executable not found in PATH.")
+
+            logdir = (
+                predictions_directory() if konfai_state() == "PREDICTION" else statistics_directory() + self.name + "/"
+            )
+
+            port = os.environ.get("KONFAI_TENSORBOARD_PORT")
+            if not port or not port.isdigit():
+                raise ValueError("Invalid or missing KONFAI_TENSORBOARD_PORT.")
+
+            command = [
+                tensorboard_exe,
+                "--logdir",
+                logdir,
+                "--port",
+                port,
+                "--bind_all",
+            ]
+            self.process = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)  # nosec B603
             try:
                 s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-                s.connect(('10.255.255.255', 1))
-                IP = s.getsockname()[0]
+                s.connect(("10.255.255.255", 1))
+                ip = s.getsockname()[0]
             except Exception:
-                IP = '127.0.0.1'
+                ip = "127.0.0.1"
             finally:
                 s.close()
-            print("Tensorboard : http://{}:{}/".format(IP, os.environ["KONFAI_TENSORBOARD_PORT"]))
+            print(f"Tensorboard : http://{ip}:{os.environ['KONFAI_TENSORBOARD_PORT']}/")
         return self
-    
-    def __exit__(self, type, value, traceback):
+
+    def __exit__(self, exc_type, value, traceback):
         if self.process is not None:
             self.process.terminate()
             self.process.wait()
 
-class DistributedObject():
+
+class DistributedObject(ABC):
 
     def __init__(self, name: str) -> None:
         self.port = find_free_port()
-        self.dataloader : list[list[DataLoader]]
-        self.manual_seed: bool = None
+        self.dataloader: list[list[DataLoader]]
+        self.manual_seed: int | None = None
         self.name = name
         self.size = 1
-    
+
     @abstractmethod
     def setup(self, world_size: int):
         pass
-    
+
     def __enter__(self):
         return self
 
-    def __exit__(self, type, value, traceback):
+    def __exit__(self, exc_type, value, traceback):
         cleanup()
 
     @abstractmethod
-    def run_process(self, world_size: int, global_rank: int, local_rank: int, dataloaders: list[DataLoader]):
-        pass 
-    
-    def getMeasure(world_size: int, global_rank: int, gpu: int, models: dict[str, torch.nn.Module], n: int) -> dict[str, tuple[dict[str, tuple[float, float]], dict[str, tuple[float, float]]]]:
+    def run_process(
+        self,
+        world_size: int,
+        global_rank: int,
+        local_rank: int,
+        dataloaders: list[DataLoader],
+    ):
+        pass
+
+    @staticmethod
+    def get_measure(
+        world_size: int,
+        global_rank: int,
+        gpu: int,
+        models: dict[str, torch.nn.Module],
+        n: int,
+    ) -> dict[str, tuple[dict[str, tuple[float, float]], dict[str, tuple[float, float]]]]:
         data = {}
         for label, model in models.items():
-            for name, network in model.getNetworks().items():
+            for name, network in model.get_networks().items():
                 if network.measure is not None:
-                    data["{}{}".format(name, label)] = (network.measure.format(True, n), network.measure.format(False, n))
+                    data[f"{name}{label}"] = (
+                        network.measure.format_loss(True, n),
+                        network.measure.format_loss(False, n),
+                    )
         outputs = synchronize_data(world_size, gpu, data)
-        result = {}
+        result: dict[str, tuple[dict[str, tuple[float, float]], dict[str, tuple[float, float]]]] = {}
         if global_rank == 0:
             for output in outputs:
                 for k, v in output.items():
@@ -369,16 +484,19 @@ class DistributedObject():
                             if k not in result:
                                 result[k] = ({}, {})
                             if u not in result[k][t]:
-                                result[k][t][u] = (n[0], 0)
-                            result[k][t][u] = (result[k][t][u][0], result[k][t][u][1]+n[1]/world_size)
+                                result[k][t][u] = (n[0], 0)  # type: ignore[index]
+                            result[k][t][u] = (
+                                result[k][t][u][0],
+                                result[k][t][u][1] + n[1] / world_size,  # type: ignore[index]
+                            )
         return result
 
-    def __call__(self, rank: Union[int, None] = None) -> None:
-        with Log(self.name, rank):
-            world_size = len(self.dataloader)
-            global_rank, local_rank = setupGPU(world_size, self.port, rank)
-            if global_rank is None:
-                return
+    def __call__(self, rank: int | None = None) -> None:
+        world_size = len(self.dataloader)
+        global_rank, local_rank = setup_gpu(world_size, self.port, rank)
+        if global_rank is None or local_rank is None:
+            return
+        with Log(self.name, global_rank):
             if torch.cuda.is_available():
                 pynvml.nvmlInit()
             if self.manual_seed is not None:
@@ -390,7 +508,7 @@ class DistributedObject():
             torch.backends.cuda.matmul.allow_tf32 = True
             torch.backends.cudnn.allow_tf32 = True
             dataloaders = self.dataloader[global_rank]
-            if torch.cuda.is_available():    
+            if torch.cuda.is_available():
                 torch.cuda.set_device(local_rank)
             try:
                 self.run_process(world_size, global_rank, local_rank, dataloaders)
@@ -399,27 +517,75 @@ class DistributedObject():
                 if torch.cuda.is_available():
                     pynvml.nvmlShutdown()
 
+
 def setup(parser: argparse.ArgumentParser) -> DistributedObject:
     # KONFAI arguments
-    KONFAI_args = parser.add_argument_group('KONFAI arguments')
-    KONFAI_args.add_argument("type", type=State, choices=list(State))
-    KONFAI_args.add_argument('-y', action='store_true', help="Accept overwrite")
-    KONFAI_args.add_argument('-tb', action='store_true', help='Start TensorBoard')
-    KONFAI_args.add_argument("-c", "--config", type=str, default="None", help="Configuration file location")
-    KONFAI_args.add_argument("-g", "--gpu", type=str, default=os.environ["CUDA_VISIBLE_DEVICES"] if "CUDA_VISIBLE_DEVICES" in os.environ else "", help="List of GPU")
-    KONFAI_args.add_argument("-cpu", "--cpu", type=str, default="1" , help="List of GPU")
-    KONFAI_args.add_argument('--num-workers', '--num_workers', default=4, type=int, help='No. of workers per DataLoader & GPU')
-    KONFAI_args.add_argument("-models_dir", "--MODELS_DIRECTORY", type=str, default="./Models/", help="Models location")
-    KONFAI_args.add_argument("-checkpoints_dir", "--CHECKPOINTS_DIRECTORY", type=str, default="./Checkpoints/", help="Checkpoints location")
-    KONFAI_args.add_argument("-model", "--MODEL", type=str, default="", help="URL Model")
-    KONFAI_args.add_argument("-predictions_dir", "--PREDICTIONS_DIRECTORY", type=str, default="./Predictions/", help="Predictions location")
-    KONFAI_args.add_argument("-evaluation_dir", "--EVALUATIONS_DIRECTORY", type=str, default="./Evaluations/", help="Evaluations location")
-    KONFAI_args.add_argument("-statistics_dir", "--STATISTICS_DIRECTORY", type=str, default="./Statistics/", help="Statistics location")
-    KONFAI_args.add_argument("-setups_dir", "--SETUPS_DIRECTORY", type=str, default="./Setups/", help="Setups location")
-    KONFAI_args.add_argument('-log', action='store_true', help='Save log')
-    KONFAI_args.add_argument('-quiet', action='store_false', help='')
-    
-    
+    konfai = parser.add_argument_group("KONFAI arguments")
+    konfai.add_argument("type", type=State, choices=list(State))
+    konfai.add_argument("-y", action="store_true", help="Accept overwrite")
+    konfai.add_argument("-tb", action="store_true", help="Start TensorBoard")
+    konfai.add_argument("-c", "--config", type=str, default="None", help="Configuration file location")
+    konfai.add_argument(
+        "-g",
+        "--gpu",
+        type=str,
+        default=(os.environ["CUDA_VISIBLE_DEVICES"] if "CUDA_VISIBLE_DEVICES" in os.environ else ""),
+        help="List of GPU",
+    )
+    konfai.add_argument("-cpu", "--cpu", type=str, default="1", help="List of GPU")
+    konfai.add_argument(
+        "--num-workers",
+        "--num_workers",
+        default=4,
+        type=int,
+        help="No. of workers per DataLoader & GPU",
+    )
+    konfai.add_argument(
+        "-models_dir",
+        "--MODELS_DIRECTORY",
+        type=str,
+        default="./Models/",
+        help="Models location",
+    )
+    konfai.add_argument(
+        "-checkpoints_dir",
+        "--CHECKPOINTS_DIRECTORY",
+        type=str,
+        default="./Checkpoints/",
+        help="Checkpoints location",
+    )
+    konfai.add_argument("-model", "--MODEL", type=str, default="", help="URL Model")
+    konfai.add_argument(
+        "-predictions_dir",
+        "--PREDICTIONS_DIRECTORY",
+        type=str,
+        default="./Predictions/",
+        help="Predictions location",
+    )
+    konfai.add_argument(
+        "-evaluation_dir",
+        "--EVALUATIONS_DIRECTORY",
+        type=str,
+        default="./Evaluations/",
+        help="Evaluations location",
+    )
+    konfai.add_argument(
+        "-statistics_dir",
+        "--STATISTICS_DIRECTORY",
+        type=str,
+        default="./Statistics/",
+        help="Statistics location",
+    )
+    konfai.add_argument(
+        "-setups_dir",
+        "--SETUPS_DIRECTORY",
+        type=str,
+        default="./Setups/",
+        help="Setups location",
+    )
+    konfai.add_argument("-log", action="store_true", help="Save log")
+    konfai.add_argument("-quiet", action="store_false", help="")
+
     args = parser.parse_args()
     config = vars(args)
 
@@ -432,14 +598,14 @@ def setup(parser: argparse.ArgumentParser) -> DistributedObject:
     os.environ["KONFAI_PREDICTIONS_DIRECTORY"] = config["PREDICTIONS_DIRECTORY"]
     os.environ["KONFAI_EVALUATIONS_DIRECTORY"] = config["EVALUATIONS_DIRECTORY"]
     os.environ["KONFAI_STATISTICS_DIRECTORY"] = config["STATISTICS_DIRECTORY"]
-    
+
     os.environ["KONFAI_STATE"] = str(config["type"])
-    
+
     os.environ["KONFAI_MODEL"] = config["MODEL"]
 
     os.environ["KONFAI_SETUPS_DIRECTORY"] = config["SETUPS_DIRECTORY"]
 
-    os.environ["KONFAI_OVERWRITE"] = "{}".format(config["y"])
+    os.environ["KONFAI_OVERWRITE"] = str(config["y"])
     os.environ["KONFAI_CONFIG_MODE"] = "Done"
     if config["tb"]:
         os.environ["KONFAI_TENSORBOARD_PORT"] = str(find_free_port())
@@ -448,37 +614,49 @@ def setup(parser: argparse.ArgumentParser) -> DistributedObject:
 
     if config["config"] == "None":
         if config["type"] is State.PREDICTION:
-             os.environ["KONFAI_CONFIG_FILE"] = "Prediction.yml"
+            os.environ["KONFAI_config_file"] = "Prediction.yml"
         elif config["type"] is State.EVALUATION:
-            os.environ["KONFAI_CONFIG_FILE"] = "Evaluation.yml"
+            os.environ["KONFAI_config_file"] = "Evaluation.yml"
         else:
-            os.environ["KONFAI_CONFIG_FILE"] = "Config.yml"
+            os.environ["KONFAI_config_file"] = "Config.yml"
     else:
-        os.environ["KONFAI_CONFIG_FILE"] = config["config"]
+        os.environ["KONFAI_config_file"] = config["config"]
     torch.autograd.set_detect_anomaly(True)
     os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
-    
-    if config["type"] is State.PREDICTION:    
+
+    if config["type"] is State.PREDICTION:
         from konfai.predictor import Predictor
-        os.environ["KONFAI_ROOT"] = "Predictor"
-        return Predictor(config=CONFIG_FILE())    
+
+        os.environ["konfai_root"] = "Predictor"
+        return Predictor(config=config_file())
     elif config["type"] is State.EVALUATION:
         from konfai.evaluator import Evaluator
-        os.environ["KONFAI_ROOT"] = "Evaluator"
-        return Evaluator(config=CONFIG_FILE())
+
+        os.environ["konfai_root"] = "Evaluator"
+        return Evaluator(config=config_file())
     else:
         from konfai.trainer import Trainer
-        os.environ["KONFAI_ROOT"] = "Trainer"
-        return Trainer(config=CONFIG_FILE())
 
+        os.environ["konfai_root"] = "Trainer"
+        return Trainer(config=config_file())
 
-def setupGPU(world_size: int, port: int, rank: Union[int, None] = None) -> tuple[int , int]:
+
+def setup_gpu(world_size: int, port: int, rank: int | None = None) -> tuple[int | None, int | None]:
     try:
-        host_name = subprocess.check_output("scontrol show hostnames {}".format(os.getenv('SLURM_JOB_NODELIST')).split()).decode().splitlines()[0]
-    except:
+        nodelist = os.getenv("SLURM_JOB_NODELIST")
+        if nodelist is None:
+            raise RuntimeError("SLURM_JOB_NODELIST is not set.")
+        scontrol_path = shutil.which("scontrol")
+        if scontrol_path is None:
+            raise FileNotFoundError("scontrol not found in PATH")
+        host_name = subprocess.check_output(
+            [scontrol_path, "show", "hostnames", nodelist], text=True, stderr=subprocess.DEVNULL
+        ).strip()  # nosec B603
+    except Exception:
         host_name = "localhost"
     if rank is None:
         import submitit
+
         job_env = submitit.JobEnvironment()
         global_rank = job_env.global_rank
         local_rank = job_env.local_rank
@@ -487,42 +665,48 @@ def setupGPU(world_size: int, port: int, rank: Union[int, None] = None) -> tuple
         local_rank = rank
     if global_rank >= world_size:
         return None, None
-    #print("tcp://{}:{}".format(host_name, port))
+    # print("tcp://{}:{}".format(host_name, port))
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-        dist.init_process_group("nccl", rank=global_rank, init_method="tcp://{}:{}".format(host_name, port), world_size=world_size)
+        dist.init_process_group(
+            "nccl",
+            rank=global_rank,
+            init_method=f"tcp://{host_name}:{port}",
+            world_size=world_size,
+        )
     else:
         if not dist.is_initialized():
             dist.init_process_group(
                 backend="gloo",
                 init_method=f"tcp://{host_name}:{port}",
                 rank=global_rank,
-                world_size=world_size
+                world_size=world_size,
             )
     return global_rank, local_rank
 
-import socket
-from contextlib import closing
 
 def find_free_port():
     with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
-        s.bind(('', 0))
+        s.bind(("", 0))
         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         return s.getsockname()[1]
-    
+
+
 def cleanup():
     if dist.is_initialized():
         dist.destroy_process_group()
 
-def synchronize_data(world_size: int, gpu: int, data: any) -> list[Any]:
+
+def synchronize_data(world_size: int, gpu: int, data: Any) -> list[Any]:
     if torch.cuda.is_available():
-        outputs: list[dict[str, tuple[dict[str, float], dict[str, float]]]] = [None for _ in range(world_size)]
+        outputs: list[dict[str, tuple[dict[str, float], dict[str, float]]] | None] = [None] * world_size
         torch.cuda.set_device(gpu)
         dist.all_gather_object(outputs, data)
     else:
         outputs = [data]
     return outputs
 
+
 def _resample(data: torch.Tensor, size: list[int]) -> torch.Tensor:
     if data.dtype == torch.uint8:
         mode = "nearest"
@@ -530,38 +714,79 @@ def _resample(data: torch.Tensor, size: list[int]) -> torch.Tensor:
         mode = "bilinear"
     else:
         mode = "trilinear"
-    return F.interpolate(data.type(torch.float32).unsqueeze(0), size=tuple([s for s in reversed(size)]), mode=mode).squeeze(0).type(data.dtype)
+    return (
+        F.interpolate(
+            data.type(torch.float32).unsqueeze(0),
+            size=tuple(reversed(size)),
+            mode=mode,
+        )
+        .squeeze(0)
+        .type(data.dtype)
+    )
+
 
 def _affine_matrix(matrix: torch.Tensor, translation: torch.Tensor) -> torch.Tensor:
-    return torch.cat((torch.cat((matrix, translation.unsqueeze(0).T), dim=1), torch.tensor([[0, 0, 0, 1]])), dim=0)
+    return torch.cat(
+        (
+            torch.cat((matrix, translation.unsqueeze(0).T), dim=1),
+            torch.tensor([[0, 0, 0, 1]]),
+        ),
+        dim=0,
+    )
+
 
 def _resample_affine(data: torch.Tensor, matrix: torch.Tensor):
     if data.dtype == torch.uint8:
         mode = "nearest"
     else:
         mode = "bilinear"
-    return F.grid_sample(data.unsqueeze(0).type(torch.float32), F.affine_grid(matrix[:, :-1,...].type(torch.float32), [1]+list(data.shape), align_corners=True), align_corners=True, mode=mode, padding_mode="reflection").squeeze(0).type(data.dtype)
-
+    return (
+        F.grid_sample(
+            data.unsqueeze(0).type(torch.float32),
+            F.affine_grid(
+                matrix[:, :-1, ...].type(torch.float32),
+                [1] + list(data.shape),
+                align_corners=True,
+            ),
+            align_corners=True,
+            mode=mode,
+            padding_mode="reflection",
+        )
+        .squeeze(0)
+        .type(data.dtype)
+    )
 
 
 def download_url(model_name: str, url: str) -> str:
     spec = importlib.util.find_spec("konfai")
-    base_path = Path(spec.submodule_search_locations[0]) / "metric" / "models"
+    if spec is None or spec.submodule_search_locations is None:
+        raise ImportError("Could not locate 'konfai' package")
+    locations = spec.submodule_search_locations
+    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)
-    filetmp = model_dir / ("tmp_"+str(Path(model_name).name))
+    filetmp = model_dir / ("tmp_" + str(Path(model_name).name))
     file = model_dir / Path(model_name).name
     if file.exists():
         return str(file)
-    
+
     try:
         print(f"[FOCUS] Downloading {model_name} to {file}")
-        with requests.get(url+model_name, stream=True) as r:
+        with requests.get(url + model_name, stream=True, timeout=10) as r:
             r.raise_for_status()
-            total = int(r.headers.get('content-length', 0))
-            with open(filetmp, 'wb') as f:
-                with tqdm(total=total, unit='B', unit_scale=True, desc=f"Downloading {model_name}") as pbar:
+            total = int(r.headers.get("content-length", 0))
+            with open(filetmp, "wb") as f:
+                with tqdm(
+                    total=total,
+                    unit="B",
+                    unit_scale=True,
+                    desc=f"Downloading {model_name}",
+                ) as pbar:
                     for chunk in r.iter_content(chunk_size=8192):
                         f.write(chunk)
                         pbar.update(len(chunk))
@@ -573,63 +798,86 @@ def download_url(model_name: str, url: str) -> str:
         if filetmp.exists():
             os.remove(filetmp)
     return str(file)
-    
-SUPPORTED_EXTENSIONS = [
-    "mha", "mhd",         # MetaImage
-    "nii", "nii.gz",      # NIfTI
-    "nrrd", "nrrd.gz",    # NRRD
-    "gipl", "gipl.gz",    # GIPL
-    "hdr", "img",         # Analyze
-    "dcm",                 # DICOM (si GDCM activé)
-    "tif", "tiff",        # TIFF
-    "png", "jpg", "jpeg", "bmp",  # 2D formats
-    "h5", "itk.txt", ".fcsv", ".xml", ".vtk", ".npy"
 
+
+SUPPORTED_EXTENSIONS = [
+    "mha",
+    "mhd",  # MetaImage
+    "nii",
+    "nii.gz",  # NIfTI
+    "nrrd",
+    "nrrd.gz",  # NRRD
+    "gipl",
+    "gipl.gz",  # GIPL
+    "hdr",
+    "img",  # Analyze
+    "dcm",  # DICOM (si GDCM activé)
+    "tif",
+    "tiff",  # TIFF
+    "png",
+    "jpg",
+    "jpeg",
+    "bmp",  # 2D formats
+    "h5",
+    "itk.txt",
+    ".fcsv",
+    ".xml",
+    ".vtk",
+    ".npy",
 ]
 
+
 class KonfAIError(Exception):
 
-    def __init__(self, typeError: str, messages: list[str]) -> None:
-        super().__init__("\n[{}] {}".format(typeError, messages[0])+("\n" if len(messages) > 0 else "")+"\n→\t".join(messages[1:]))
+    def __init__(self, type_error: str, *messages: str) -> None:
+        super().__init__(
+            f"\n[{type_error}] {messages[0]}" + ("\n" if len(messages) > 0 else "") + "\n→\t".join(messages[1:])
+        )
 
 
 class ConfigError(KonfAIError):
 
     def __init__(self, *message) -> None:
-        super().__init__("Config", message)
+        super().__init__("Config", *message)
 
 
 class DatasetManagerError(KonfAIError):
 
     def __init__(self, *message) -> None:
-        super().__init__("DatasetManager", message)
+        super().__init__("DatasetManager", *message)
+
 
 class MeasureError(KonfAIError):
 
     def __init__(self, *message) -> None:
-        super().__init__("Measure", message)
+        super().__init__("Measure", *message)
+
 
 class TrainerError(KonfAIError):
 
     def __init__(self, *message) -> None:
-        super().__init__("Trainer", message)
+        super().__init__("Trainer", *message)
+
 
 class AugmentationError(KonfAIError):
 
     def __init__(self, *message) -> None:
-        super().__init__("Augmentation", message)
+        super().__init__("Augmentation", *message)
+
 
 class EvaluatorError(KonfAIError):
 
     def __init__(self, *message) -> None:
-        super().__init__("Evaluator", message)
+        super().__init__("Evaluator", *message)
+
 
 class PredictorError(KonfAIError):
 
     def __init__(self, *message) -> None:
-        super().__init__("Predictor", message)
+        super().__init__("Predictor", *message)
+
 
 class TransformError(KonfAIError):
 
     def __init__(self, *message) -> None:
-        super().__init__("Transform", message)
+        super().__init__("Transform", *message)