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

konfai/metric/measure.py

@@ -1,334 +1,434 @@
-from abc import ABC
+import copy
 import importlib
-import numpy as np
-import torch
-
-from torchvision.models import inception_v3, Inception_V3_Weights
-from torchvision.transforms import functional as F
-from scipy import linalg    
-
-import torch.nn.functional as F
 import os
-
-from typing import Callable, Union
+from abc import ABC, abstractmethod
+from collections.abc import Callable
 from functools import partial
-import torch.nn.functional as F
+
+import numpy as np
+import torch
+import torch.nn.functional as F  # noqa: N812
+from scipy import linalg
 from skimage.metrics import structural_similarity
-import copy
-from abc import abstractmethod
+from torchvision.models import Inception_V3_Weights, inception_v3
+from tqdm import tqdm
 
-from konfai.utils.config import config
-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
+from konfai.utils.config import config
+from konfai.utils.utils import download_url, get_module
+
+models_register = {}
 
-modelsRegister = {}
 
 class Criterion(torch.nn.Module, ABC):
 
     def __init__(self) -> None:
         super().__init__()
 
-    def init(self, model : torch.nn.Module, output_group : str, target_group : str) -> str:
+    def init(self, model: torch.nn.Module, output_group: str, target_group: str) -> str:
         return output_group
 
     @abstractmethod
-    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         pass
 
+
 class MaskedLoss(Criterion):
 
-    def __init__(self, loss : Callable[[torch.Tensor, torch.Tensor], torch.Tensor], mode_image_masked: bool) -> None:
+    def __init__(
+        self,
+        loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        mode_image_masked: bool,
+    ) -> None:
         super().__init__()
         self.loss = loss
         self.mode_image_masked = mode_image_masked
 
-    def getMask(self, targets: list[torch.Tensor]) -> torch.Tensor:
+    def get_mask(self, targets: list[torch.Tensor]) -> torch.Tensor | None:
         result = None
         if len(targets) > 0:
             result = targets[0]
             for mask in targets[1:]:
-                result = result*mask
+                result = result * mask
         return result
 
-    def forward(self, output: torch.Tensor, *targets : list[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.getMask(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):
 
+    @staticmethod
     def _loss(reduction: str, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         return torch.nn.MSELoss(reduction=reduction)(x, y)
 
     def __init__(self, reduction: str = "mean") -> None:
         super().__init__(partial(MSE._loss, reduction), False)
 
+
 class MAE(MaskedLoss):
 
+    @staticmethod
     def _loss(reduction: str, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         return torch.nn.L1Loss(reduction=reduction)(x, y)
-    
+
     def __init__(self, reduction: str = "mean") -> None:
         super().__init__(partial(MAE._loss, reduction), False)
 
+
 class PSNR(MaskedLoss):
 
-    def _loss(dynamic_range: float, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: 
+    @staticmethod
+    def _loss(dynamic_range: float, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         mse = torch.mean((x - y).pow(2))
-        psnr = 10 * torch.log10(dynamic_range**2 / mse) 
-        return psnr 
+        psnr = 10 * torch.log10(dynamic_range**2 / mse)
+        return psnr
 
-    def __init__(self, dynamic_range: Union[float, None] = None) -> None:
-        dynamic_range = dynamic_range if dynamic_range else 1024+3071
+    def __init__(self, dynamic_range: float | None = None) -> None:
+        dynamic_range = dynamic_range if dynamic_range else 1024 + 3071
         super().__init__(partial(PSNR._loss, dynamic_range), False)
-    
+
+
 class SSIM(MaskedLoss):
-    
+
+    @staticmethod
     def _loss(dynamic_range: float, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        return structural_similarity(x[0][0].detach().cpu().numpy(), y[0][0].cpu().numpy(), data_range=dynamic_range, gradient=False, full=False)
-    
-    def __init__(self, dynamic_range: Union[float, None] = None) -> None:
-        dynamic_range = dynamic_range if dynamic_range else 1024+3000
+        return structural_similarity(
+            x[0][0].detach().cpu().numpy(),
+            y[0][0].cpu().numpy(),
+            data_range=dynamic_range,
+            gradient=False,
+            full=False,
+        )
+
+    def __init__(self, dynamic_range: float | None = None) -> None:
+        dynamic_range = dynamic_range if dynamic_range else 1024 + 3000
         super().__init__(partial(SSIM._loss, dynamic_range), True)
 
 
 class LPIPS(MaskedLoss):
 
-    def normalize(input: torch.Tensor) -> torch.Tensor:
-        return (input-torch.min(input))/(torch.max(input)-torch.min(input))*2-1 
-        
-    def preprocessing(input: torch.Tensor) -> torch.Tensor:
-        return input.repeat((1,3,1,1))
-    
+    @staticmethod
+    def normalize(tensor: torch.Tensor) -> torch.Tensor:
+        return (tensor - torch.min(tensor)) / (torch.max(tensor) - torch.min(tensor)) * 2 - 1
+
+    @staticmethod
+    def preprocessing(tensor: torch.Tensor) -> torch.Tensor:
+        return tensor.repeat((1, 3, 1, 1)).to(0)
+
+    @staticmethod
     def _loss(loss_fn_alex, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        datasetPatch = ModelPatch([1, 64, 64])
-        datasetPatch.load(x.shape[2:])
+        dataset_patch = ModelPatch([1, 64, 64])
+        dataset_patch.load(x.shape[2:])
 
-        patchIterator = datasetPatch.disassemble(LPIPS.normalize(x), LPIPS.normalize(y))
+        patch_iterator = dataset_patch.disassemble(LPIPS.normalize(x), LPIPS.normalize(y))
         loss = 0
-        with tqdm(iterable = enumerate(patchIterator), leave=False, total=datasetPatch.getSize(0)) as batch_iter:
-            for i, patch_input in batch_iter:
+        with tqdm(
+            iterable=enumerate(patch_iterator),
+            leave=False,
+            total=dataset_patch.get_size(0),
+        ) as batch_iter:
+            for _, patch_input in batch_iter:
                 real, fake = LPIPS.preprocessing(patch_input[0]), LPIPS.preprocessing(patch_input[1])
                 loss += loss_fn_alex(real, fake).item()
-        return loss/datasetPatch.getSize(0)
+        return loss / dataset_patch.get_size(0)
 
     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
-    
-    def flatten(self, tensor : torch.Tensor) -> torch.Tensor:
+        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)
 
-    def dice_per_channel(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        input = self.flatten(input)
+    def dice_per_channel(self, tensor: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        tensor = self.flatten(tensor)
         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 = 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))
+        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")
+        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):
 
     def __init__(self):
         super().__init__()
-    
+
     @staticmethod
-    def _image_gradient2D(image : torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def _image_gradient_2d(image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         dx = image[:, :, 1:, :] - image[:, :, :-1, :]
         dy = image[:, :, :, 1:] - image[:, :, :, :-1]
         return dx, dy
 
     @staticmethod
-    def _image_gradient3D(image : torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def _image_gradient_3d(
+        image: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         dx = image[:, :, 1:, :, :] - image[:, :, :-1, :, :]
         dy = image[:, :, :, 1:, :] - image[:, :, :, :-1, :]
         dz = image[:, :, :, :, 1:] - image[:, :, :, :, :-1]
         return dx, dy, dz
-        
-    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
+
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         target_0 = targets[0]
         if len(output.shape) == 5:
-            dx, dy, dz = GradientImages._image_gradient3D(output)
+            dx, dy, dz = GradientImages._image_gradient_3d(output)
             if target_0 is not None:
-                dx_tmp, dy_tmp, dz_tmp = GradientImages._image_gradient3D(target_0)
+                dx_tmp, dy_tmp, dz_tmp = GradientImages._image_gradient_3d(target_0)
                 dx -= dx_tmp
                 dy -= dy_tmp
                 dz -= dz_tmp
             return dx.norm() + dy.norm() + dz.norm()
         else:
-            dx, dy = GradientImages._image_gradient2D(output)
+            dx, dy = GradientImages._image_gradient_2d(output)
             if target_0 is not None:
-                dx_tmp, dy_tmp = GradientImages._image_gradient2D(target_0)
+                dx_tmp, dy_tmp = GradientImages._image_gradient_2d(target_0)
                 dx -= dx_tmp
                 dy -= dy_tmp
             return dx.norm() + dy.norm()
-        
+
+
 class BCE(Criterion):
 
-    def __init__(self, target : float = 0) -> None:
+    def __init__(self, target: float = 0) -> None:
         super().__init__()
         self.loss = torch.nn.BCEWithLogitsLoss()
-        self.register_buffer('target', torch.tensor(target).type(torch.float32))
+        self.register_buffer("target", torch.tensor(target).type(torch.float32))
 
-    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         target = self._buffers["target"]
         return self.loss(output, target.to(output.device).expand_as(output))
 
+
 class PatchGanLoss(Criterion):
 
-    def __init__(self, target : float = 0) -> None:
+    def __init__(self, target: float = 0) -> None:
         super().__init__()
         self.loss = torch.nn.MSELoss()
-        self.register_buffer('target', torch.tensor(target).type(torch.float32))
+        self.register_buffer("target", torch.tensor(target).type(torch.float32))
 
-    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         target = self._buffers["target"]
-        return self.loss(output, (torch.ones_like(output)*target).to(output.device))
+        return self.loss(output, (torch.ones_like(output) * target).to(output.device))
+
 
 class WGP(Criterion):
 
     def __init__(self) -> None:
         super().__init__()
-    
-    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
-        return torch.mean((output - 1)**2)
+
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
+        return torch.mean((output - 1) ** 2)
+
 
 class Gram(Criterion):
 
-    def computeGram(input : torch.Tensor):
-        (b, ch, w) = input.size()
-        with torch.amp.autocast('cuda', enabled=False):
-            return input.bmm(input.transpose(1, 2)).div(ch*w)
+    @staticmethod
+    def compute_gram(tensor: torch.Tensor):
+        (b, ch, w) = tensor.size()
+        with torch.amp.autocast("cuda", enabled=False):
+            return tensor.bmm(tensor.transpose(1, 2)).div(ch * w)
 
     def __init__(self) -> None:
         super().__init__()
-        self.loss = torch.nn.L1Loss(reduction='sum')
+        self.loss = torch.nn.L1Loss(reduction="sum")
 
-    def forward(self, output: torch.Tensor, target : torch.Tensor) -> torch.Tensor:
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
+        target = targets[0]
         if len(output.shape) > 3:
             output = output.view(output.shape[0], output.shape[1], int(np.prod(output.shape[2:])))
         if len(target.shape) > 3:
             target = target.view(target.shape[0], target.shape[1], int(np.prod(target.shape[2:])))
-        return self.loss(Gram.computeGram(output), Gram.computeGram(target))
+        return self.loss(Gram.compute_gram(output), Gram.compute_gram(target))
+
 
 class PerceptualLoss(Criterion):
-    
-    class Module():
-        
+
+    class Module:
+
         @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 ""
+            self.konfai_args = os.environ["KONFAI_CONFIG_PATH"] if "KONFAI_CONFIG_PATH" in os.environ else ""
 
-        def getLoss(self) -> dict[torch.nn.Module, float]:
+        def get_loss(self) -> dict[torch.nn.Module, float]:
             result: dict[torch.nn.Module, float] = {}
-            for loss, l in self.losses.items():
-                module, name = _getModule(loss, "konfai.metric.measure")
-                result[config(self.DL_args)(getattr(importlib.import_module(module), name))(config=None)] = l   
+            for loss, loss_value in self.losses.items():
+                module, name = get_module(loss, "konfai.metric.measure")
+                result[config(self.konfai_args)(getattr(importlib.import_module(module), name))(config=None)] = (
+                    loss_value
+                )
             return result
-        
-    def __init__(self, modelLoader : ModelLoader = ModelLoader(), path_model : str = "name", modules : dict[str, Module] = {"UNetBlock_0.DownConvBlock.Activation_1": Module({"Gram": 1, "torch_nn_L1Loss": 1})}, shape: list[int] = [128, 128, 128]) -> None:
+
+    def __init__(
+        self,
+        model_loader: ModelLoader = ModelLoader(),
+        path_model: str = "name",
+        modules: dict[str, Module] = {
+            "UNetBlock_0.DownConvBlock.Activation_1": Module({"Gram": 1, "torch_nn_L1Loss": 1})
+        },
+        shape: list[int] = [128, 128, 128],
+    ) -> None:
         super().__init__()
         self.path_model = path_model
-        if self.path_model not in modelsRegister:
-            self.model = modelLoader.getModel(train=False, DL_args=os.environ['KONFAI_CONFIG_PATH'].split("PerceptualLoss")[0]+"PerceptualLoss.Model", DL_without=["optimizer", "schedulers", "nb_batch_per_step", "init_type", "init_gain", "outputsCriterions", "drop_p"])
+        if self.path_model not in models_register:
+            self.model = model_loader.get_model(
+                train=False,
+                konfai_args=os.environ["KONFAI_CONFIG_PATH"].split("PerceptualLoss")[0] + "PerceptualLoss.Model",
+                konfai_without=[
+                    "optimizer",
+                    "schedulers",
+                    "nb_batch_per_step",
+                    "init_type",
+                    "init_gain",
+                    "outputs_criterions",
+                    "drop_p",
+                ],
+            )
             if path_model.startswith("https"):
                 state_dict = torch.hub.load_state_dict_from_url(path_model)
-                state_dict = {"Model": {self.model.getName() : state_dict["model"]}}
+                state_dict = {"Model": {self.model.get_name(): state_dict["model"]}}
             else:
                 state_dict = torch.load(path_model, weights_only=True)
             self.model.load(state_dict)
-            modelsRegister[self.path_model] = self.model
+            models_register[self.path_model] = self.model
         else:
-            self.model = modelsRegister[self.path_model]
+            self.model = models_register[self.path_model]
 
         self.shape = shape
-        self.mode = "trilinear" if  len(shape) == 3 else "bilinear"
+        self.mode = "trilinear" if len(shape) == 3 else "bilinear"
         self.modules_loss: dict[str, dict[torch.nn.Module, float]] = {}
         for name, losses in modules.items():
-            self.modules_loss[name.replace(":", ".")] = losses.getLoss()
+            self.modules_loss[name.replace(":", ".")] = losses.get_loss()
 
         self.model.eval()
         self.model.requires_grad_(False)
         self.models: dict[int, torch.nn.Module] = {}
 
-    def preprocessing(self, input: torch.Tensor) -> torch.Tensor:
-        #if not all([input.shape[-i-1] == size for i, size in enumerate(reversed(self.shape[2:]))]):
-        #    input = F.interpolate(input, mode=self.mode, size=tuple(self.shape), align_corners=False).type(torch.float32)
-        #if input.shape[1] != self.model.in_channels:
-        #    input = input.repeat(tuple([1,self.model.in_channels] + [1 for _ in range(len(self.shape))]))   
-        #input = (input - torch.min(input))/(torch.max(input)-torch.min(input))
-        #input = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(input)
-        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]
-        for zipped_output in zip([output], *[[target] for target in targets]):
+    def preprocessing(self, tensor: torch.Tensor) -> torch.Tensor:
+        # if not all([tensor.shape[-i-1] == size for i, size in enumerate(reversed(self.shape[2:]))]):
+        #    tensor = F.interpolate(tensor, mode=self.mode,
+        # size=tuple(self.shape), align_corners=False).type(torch.float32)
+        # if tensor.shape[1] != self.model.in_channels:
+        #    tensor = tensor.repeat(tuple([1,self.model.in_channels] + [1 for _ in range(len(self.shape))]))
+        # tensor = (tensor - torch.min(tensor))/(torch.max(tensor)-torch.min(tensor))
+        # tensor = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(tensor)
+        return tensor
+
+    def _compute(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
+        loss = torch.zeros((1), requires_grad=True).to(output.device, non_blocking=False).type(torch.float32)
+        output_preprocessing = self.preprocessing(output)
+        targets_preprocessing = [self.preprocessing(target) for target in targets]
+        for zipped_output in zip([output_preprocessing], *[[target] for target in targets_preprocessing]):
             output = zipped_output[0]
             targets = zipped_output[1:]
 
-            for zipped_layers in list(zip(self.models[output.device.index].get_layers([output], set(self.modules_loss.keys()).copy()), *[self.models[output.device.index].get_layers([target], set(self.modules_loss.keys()).copy()) for target in targets])):
-                output_layer = zipped_layers[0][1].view(zipped_layers[0][1].shape[0], zipped_layers[0][1].shape[1], int(np.prod(zipped_layers[0][1].shape[2:])))
-                for (loss_function, l), target_layer in zip(self.modules_loss[zipped_layers[0][0]].items(), zipped_layers[1:]):
-                    target_layer = target_layer[1].view(target_layer[1].shape[0], target_layer[1].shape[1], int(np.prod(target_layer[1].shape[2:])))
-                    loss = loss+l*loss_function(output_layer.float(), target_layer.float())/output_layer.shape[0]
+            for zipped_layers in list(
+                zip(
+                    self.models[output.device.index].get_layers([output], set(self.modules_loss.keys()).copy()),
+                    *[
+                        self.models[output.device.index].get_layers([target], set(self.modules_loss.keys()).copy())
+                        for target in targets
+                    ],
+                )
+            ):
+                output_layer = zipped_layers[0][1].view(
+                    zipped_layers[0][1].shape[0],
+                    zipped_layers[0][1].shape[1],
+                    int(np.prod(zipped_layers[0][1].shape[2:])),
+                )
+                for (loss_function, loss_value), target_layer in zip(
+                    self.modules_loss[zipped_layers[0][0]].items(), zipped_layers[1:]
+                ):
+                    target_layer = target_layer[1].view(
+                        target_layer[1].shape[0],
+                        target_layer[1].shape[1],
+                        int(np.prod(target_layer[1].shape[2:])),
+                    )
+                    loss = (
+                        loss
+                        + loss_value * loss_function(output_layer.float(), target_layer.float()) / output_layer.shape[0]
+                    )
         return loss
-    
-    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
+
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         if output.device.index not in self.models:
             del os.environ["device"]
             self.models[output.device.index] = Network.to(copy.deepcopy(self.model).eval(), output.device.index).eval()
-        loss = torch.zeros((1), requires_grad = True).to(output.device, non_blocking=False).type(torch.float32)
+        loss = torch.zeros((1), requires_grad=True).to(output.device, non_blocking=False).type(torch.float32)
         if len(output.shape) == 5 and len(self.shape) == 2:
             for i in range(output.shape[2]):
-                loss = loss + self._compute(output[:, :, i, ...], [t[:, :, i, ...] for t in targets])/output.shape[2]
+                loss = loss + self._compute(output[:, :, i, ...], [t[:, :, i, ...] for t in targets]) / output.shape[2]
         else:
             loss = self._compute(output, targets)
         return loss.to(output)
 
+
 class KLDivergence(Criterion):
-    
-    def __init__(self, shape: list[int], dim : int = 100, mu : float = 0, std : float = 1) -> None:
+
+    def __init__(self, shape: list[int], dim: int = 100, mu: float = 0, std: float = 1) -> None:
         super().__init__()
-        self.latentDim = dim
+        self.latent_dim = dim
         self.mu = torch.Tensor([mu])
         self.std = torch.Tensor([std])
         self.modelDim = 3
         self.shape = shape
         self.loss = torch.nn.KLDivLoss()
-        
-    def init(self, model : Network, output_group : str, target_group : str) -> str:
+
+    def init(self, model: Network, output_group: str, target_group: str) -> str:
         super().init(model, output_group, target_group)
         model._compute_channels_trace(model, model.in_channels, None, None)
 
@@ -338,52 +438,35 @@ class KLDivergence(Criterion):
 
         modules = last_module._modules.copy()
         last_module._modules.clear()
-        
+
         for name, value in modules.items():
             last_module._modules[name] = value
             if name == output_group.split(".")[-1]:
-                last_module.add_module("LatentDistribution", LatentDistribution(shape = self.shape, latentDim=self.latentDim))
-        return ".".join(output_group.split(".")[:-1])+".LatentDistribution.Concat"
+                last_module.add_module(
+                    "LatentDistribution",
+                    LatentDistribution(shape=self.shape, latent_dim=self.latent_dim),
+                )
+        return ".".join(output_group.split(".")[:-1]) + ".LatentDistribution.Concat"
 
-    def forward(self, output: torch.Tensor, targets : list[torch.Tensor]) -> torch.Tensor:
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         mu = output[:, 0, :]
         log_std = output[:, 1, :]
-        return torch.mean(-0.5 * torch.sum(1 + log_std - mu**2 - torch.exp(log_std), dim = 1), dim = 0)
-
-    """
-    def forward(self, input : torch.Tensor) -> torch.Tensor:
-        mu = input[:, 0, :]
-        log_std = input[:, 1, :]
+        return torch.mean(-0.5 * torch.sum(1 + log_std - mu**2 - torch.exp(log_std), dim=1), dim=0)
 
-        z = input[:, 2, :]
 
-        q = torch.distributions.Normal(mu, log_std)
-
-        target_mu = torch.ones((self.latentDim)).to(input.device)*self.mu.to(input.device)
-        target_std = torch.ones((self.latentDim)).to(input.device)*self.std.to(input.device)
-
-        p = torch.distributions.Normal(target_mu, target_std)
-        
-        log_pz = p.log_prob(z)
-        log_qzx = q.log_prob(z)
-
-        kl = (log_pz - log_qzx)
-        kl = kl.sum(-1)
-        return kl
-    """
-    
 class Accuracy(Criterion):
 
     def __init__(self) -> None:
         super().__init__()
-        self.n : int = 0
-        self.corrects = torch.zeros((1))
+        self.n: int = 0
+        self.corrects = torch.zeros(1)
 
-    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         target_0 = targets[0]
         self.n += output.shape[0]
         self.corrects += (torch.argmax(torch.softmax(output, dim=1), dim=1) == target_0).sum().float().cpu()
-        return self.corrects/self.n
+        return self.corrects / self.n
+
 
 class TripletLoss(Criterion):
 
@@ -391,9 +474,10 @@ class TripletLoss(Criterion):
         super().__init__()
         self.triplet_loss = torch.nn.TripletMarginLoss(margin=1.0, p=2, eps=1e-7)
 
-    def forward(self, output: torch.Tensor) -> torch.Tensor:
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         return self.triplet_loss(output[0], output[1], output[2])
 
+
 class L1LossRepresentation(Criterion):
 
     def __init__(self) -> None:
@@ -403,28 +487,31 @@ class L1LossRepresentation(Criterion):
     def _variance(self, features: torch.Tensor) -> torch.Tensor:
         return torch.mean(torch.clamp(1 - torch.var(features, dim=0), min=0))
 
-    def forward(self, output: torch.Tensor) -> torch.Tensor:
-        return self.loss(output[0], output[1])+ self._variance(output[0]) + self._variance(output[1])
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
+        return self.loss(output[0], output[1]) + self._variance(output[0]) + self._variance(output[1])
+
 
 class FocalLoss(Criterion):
 
-    def __init__(self, gamma: float = 2.0, alpha: list[float] = [0.5, 2.0, 0.5, 0.5, 1], reduction: str = "mean"):
+    def __init__(
+        self,
+        gamma: float = 2.0,
+        alpha: list[float] = [0.5, 2.0, 0.5, 0.5, 1],
+        reduction: str = "mean",
+    ):
         super().__init__()
         raw_alpha = torch.tensor(alpha, dtype=torch.float32)
         self.alpha = raw_alpha / raw_alpha.sum() * len(raw_alpha)
         self.gamma = gamma
         self.reduction = reduction
 
-    def forward(self, output: torch.Tensor, *targets: list[torch.Tensor]) -> torch.Tensor:
-        target = F.interpolate(
-            targets[0],
-            output.shape[2:],
-            mode="nearest").long()
-        
+    def forward(self, output: torch.Tensor, *targets: 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) 
+        logpt = logpt.gather(1, target)
         pt = pt.gather(1, target)
 
         at = self.alpha.to(target.device)[target].unsqueeze(1)
@@ -435,7 +522,7 @@ class FocalLoss(Criterion):
         elif self.reduction == "sum":
             return loss.sum()
         return loss
-    
+
 
 class FID(Criterion):
 
@@ -449,19 +536,26 @@ class FID(Criterion):
 
         def forward(self, x: torch.Tensor) -> torch.Tensor:
             return self.model(x)
-        
+
     def __init__(self) -> None:
         super().__init__()
         self.inception_model = FID.InceptionV3().cuda()
-        
+
+    @staticmethod
     def preprocess_images(image: torch.Tensor) -> torch.Tensor:
-        return F.normalize(F.resize(image, (299, 299)).repeat((1,3,1,1)), mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]).cuda()
+        return F.normalize(
+            F.resize(image, (299, 299)).repeat((1, 3, 1, 1)),
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225],
+        ).cuda()
 
+    @staticmethod
     def get_features(images: torch.Tensor, model: torch.nn.Module) -> np.ndarray:
         with torch.no_grad():
             features = model(images).cpu().numpy()
         return features
 
+    @staticmethod
     def calculate_fid(real_features: np.ndarray, generated_features: np.ndarray) -> float:
         mu1 = np.mean(real_features, axis=0)
         sigma1 = np.cov(real_features, rowvar=False)
@@ -475,7 +569,7 @@ class FID(Criterion):
 
         return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * np.trace(covmean)
 
-    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
+    def forward(self, output: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
         real_images = FID.preprocess_images(targets[0].squeeze(0).permute([1, 0, 2, 3]))
         generated_images = FID.preprocess_images(output.squeeze(0).permute([1, 0, 2, 3]))
 
@@ -484,8 +578,15 @@ class FID(Criterion):
 
         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:
+    def __init__(
+        self,
+        num_bins: int = 23,
+        sigma_ratio: float = 0.5,
+        smooth_nr: float = 1e-7,
+        smooth_dr: float = 1e-7,
+    ) -> None:
         super().__init__()
         bin_centers = torch.linspace(0.0, 1.0, num_bins)
         sigma = torch.mean(bin_centers[1:] - bin_centers[:-1]) * sigma_ratio
@@ -495,7 +596,9 @@ class MutualInformationLoss(torch.nn.Module):
         self.smooth_nr = float(smooth_nr)
         self.smooth_dr = float(smooth_dr)
 
-    def parzen_windowing(self, pred: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    def parzen_windowing(
+        self, pred: torch.Tensor, target: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         pred_weight, pred_probability = self.parzen_windowing_gaussian(pred)
         target_weight, target_probability = self.parzen_windowing_gaussian(target)
         return pred_weight, pred_probability, target_weight, target_probability
@@ -503,82 +606,95 @@ class MutualInformationLoss(torch.nn.Module):
     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)
-        weight = torch.exp(-self.preterm.to(img) * (img - self.bin_centers.to(img)) ** 2)  # (batch, num_sample, num_bin)
+        weight = torch.exp(
+            -self.preterm.to(img) * (img - self.bin_centers.to(img)) ** 2
+        )  # (batch, num_sample, num_bin)
         weight = weight / torch.sum(weight, dim=-1, keepdim=True)  # (batch, num_sample, num_bin)
         probability = torch.mean(weight, dim=-2, keepdim=True)  # (batch, 1, num_bin)
         return weight, probability
 
-    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        wa, pa, wb, pb = self.parzen_windowing(pred, target)  # (batch, num_sample, num_bin), (batch, 1, num_bin)
+    def forward(self, pred: torch.Tensor, *targets: torch.Tensor) -> torch.Tensor:
+        wa, pa, wb, pb = self.parzen_windowing(pred, targets[0])  # (batch, num_sample, num_bin), (batch, 1, num_bin)
         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)
+        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
 
-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:
+class IMPACTSynth(Criterion):  # Feature-Oriented Comparison for Unpaired Synthesis
+
+    def __init__(
+        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.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.loss = torch.nn.L1Loss()
+        self.weights = 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)  # nosec B614
+        self.dim = len(shape)
+        self.shape = shape if all(s > 0 for s in 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)))
+            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
         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 preprocessing(self, tensor: torch.Tensor) -> torch.Tensor:
+        if self.shape is not None and not all(
+            tensor.shape[-i - 1] == size for i, size in enumerate(reversed(self.shape[2:]))
+        ):
+            tensor = torch.nn.functional.interpolate(
+                tensor, mode=self.mode, size=tuple(self.shape), align_corners=False
+            ).type(torch.float32)
+        if tensor.shape[1] != self.in_channels:
+            tensor = tensor.repeat(tuple([1, 3] + [1 for _ in range(self.dim)]))
+        return tensor
+
     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)
+        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)
+        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_feature in targets_features:
+                loss += weight * self.loss(output_feature, target_feature)
         return loss
-    
-    def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
+
+    def forward(self, output: torch.Tensor, *targets: 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)
+            self.model = torch.jit.load(self.model_path)  # nosec B614
+        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]
+                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)
+            loss = self._compute(output, list(targets))
+        return loss.to(output)