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

konfai/models/registration/registration.py

@@ -1,99 +1,168 @@
 import torch
+import torch.nn.functional as F  # noqa: N812
 from torch.nn.parameter import Parameter
-import torch.nn.functional as F
 
-from konfai.network import network, blocks
-from konfai.utils.config import config
 from konfai.models.segmentation import UNet
+from konfai.network import blocks, network
+
 
 class VoxelMorph(network.Network):
-                
-    def __init__(   self,
-                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
-                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
-                    dim : int = 3,
-                    channels : list[int] = [4, 16,32,32,32],
-                    blockConfig: blocks.BlockConfig = blocks.BlockConfig(),
-                    nb_conv_per_stage: int = 2,
-                    downSampleMode: str = "MAXPOOL",
-                    upSampleMode: str = "CONV_TRANSPOSE",
-                    attention : bool = False,
-                    shape : list[int] = [192, 192, 192],
-                    int_steps : int = 7,
-                    int_downsize : int = 2,
-                    nb_batch_per_step : int = 1,
-                    rigid: bool = False):
-        super().__init__(in_channels = channels[0], optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim = dim, nb_batch_per_step=nb_batch_per_step)
-        self.add_module("Concat", blocks.Concat(), in_branch=[0,1], out_branch=["input_concat"])
-        self.add_module("UNetBlock_0", UNet.UNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, block = blocks.ConvBlock, nb_class=0, dim=dim), in_branch=["input_concat"], out_branch=["unet"])
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        dim: int = 3,
+        channels: list[int] = [4, 16, 32, 32, 32],
+        block_config: blocks.BlockConfig = blocks.BlockConfig(),
+        nb_conv_per_stage: int = 2,
+        downsample_mode: str = "MAXPOOL",
+        upsample_mode: str = "CONV_TRANSPOSE",
+        attention: bool = False,
+        shape: list[int] = [192, 192, 192],
+        int_steps: int = 7,
+        int_downsize: int = 2,
+        nb_batch_per_step: int = 1,
+        rigid: bool = False,
+    ):
+        super().__init__(
+            in_channels=channels[0],
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=dim,
+            nb_batch_per_step=nb_batch_per_step,
+        )
+        self.add_module("Concat", blocks.Concat(), in_branch=[0, 1], out_branch=["input_concat"])
+        self.add_module(
+            "UNetBlock_0",
+            UNet.UNetBlock(
+                channels,
+                nb_conv_per_stage,
+                block_config,
+                downsample_mode=blocks.DownsampleMode[downsample_mode],
+                upsample_mode=blocks.UpsampleMode[upsample_mode],
+                attention=attention,
+                block=blocks.ConvBlock,
+                nb_class=0,
+                dim=dim,
+            ),
+            in_branch=["input_concat"],
+            out_branch=["unet"],
+        )
 
         if rigid:
-            self.add_module("Flow", Rigid(channels[1], dim), in_branch=["unet"], out_branch=["pos_flow"])
+            self.add_module(
+                "Flow",
+                Rigid(channels[1], dim),
+                in_branch=["unet"],
+                out_branch=["pos_flow"],
+            )
         else:
-            self.add_module("Flow", Flow(channels[1], int_steps, int_downsize, shape, dim), in_branch=["unet"], out_branch=["pos_flow"])
-        self.add_module("MovingImageResample", SpatialTransformer(shape, rigid=rigid), in_branch=[1, "pos_flow"], out_branch=["moving_image_resample"])
+            self.add_module(
+                "Flow",
+                Flow(channels[1], int_steps, int_downsize, shape, dim),
+                in_branch=["unet"],
+                out_branch=["pos_flow"],
+            )
+        self.add_module(
+            "MovingImageResample",
+            SpatialTransformer(shape, rigid=rigid),
+            in_branch=[1, "pos_flow"],
+            out_branch=["moving_image_resample"],
+        )
+
 
 class Flow(network.ModuleArgsDict):
 
-    def __init__(self, in_channels: int, int_steps: int, int_downsize: int, shape: list[int], dim: int) -> None:
+    def __init__(
+        self,
+        in_channels: int,
+        int_steps: int,
+        int_downsize: int,
+        shape: list[int],
+        dim: int,
+    ) -> None:
         super().__init__()
-        self.add_module("Head", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = dim, kernel_size = 3, stride = 1, padding = 1))
+        self.add_module(
+            "Head",
+            blocks.get_torch_module("Conv", dim)(
+                in_channels=in_channels,
+                out_channels=dim,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            ),
+        )
         self["Head"].weight = Parameter(torch.distributions.Normal(0, 1e-5).sample(self["Head"].weight.shape))
         self["Head"].bias = Parameter(torch.zeros(self["Head"].bias.shape))
-    
+
         if int_steps > 0 and int_downsize > 1:
             self.add_module("DownSample", ResizeTransform(int_downsize))
 
         if int_steps > 0:
-            self.add_module("Integrate_pos_flow", VecInt([int(dim / int_downsize) for dim in shape], int_steps))
-        
+            self.add_module(
+                "Integrate_pos_flow",
+                VecInt([int(dim / int_downsize) for dim in shape], int_steps),
+            )
+
         if int_steps > 0 and int_downsize > 1:
             self.add_module("Upsample_pos_flow", ResizeTransform(1 / int_downsize))
 
+
 class Rigid(network.ModuleArgsDict):
 
     def __init__(self, in_channels: int, dim: int) -> None:
         super().__init__()
         self.add_module("ToFeatures", torch.nn.Flatten(1))
-        self.add_module("Head", torch.nn.Linear(in_channels*512*512, 2))
+        self.add_module("Head", torch.nn.Linear(in_channels * 512 * 512, 2))
 
     def init(self, init_type: str, init_gain: float):
         self["Head"].weight.data.fill_(0)
         self["Head"].bias.data.copy_(torch.tensor([0, 0], dtype=torch.float))
-        
+
+
 class MaskFlow(torch.nn.Module):
 
     def __init__(self):
         super().__init__()
-    
+
     def forward(self, mask: torch.Tensor, *flows: torch.Tensor):
         result = torch.zeros_like(flows[0])
         for i, flow in enumerate(flows):
-            result = result+torch.where(mask == i+1, flow, torch.tensor(0))
+            result = result + torch.where(mask == i + 1, flow, torch.tensor(0))
         return result
 
+
 class SpatialTransformer(torch.nn.Module):
-    
-    def __init__(self, size : list[int], rigid: bool = False):
+
+    def __init__(self, size: list[int], rigid: bool = False):
         super().__init__()
         self.rigid = rigid
         if not rigid:
             vectors = [torch.arange(0, s) for s in size]
-            grids = torch.meshgrid(vectors, indexing='ij')
+            grids = torch.meshgrid(vectors, indexing="ij")
             grid = torch.stack(grids)
             grid = torch.unsqueeze(grid, 0)
             grid = grid.type(torch.float)
-            self.register_buffer('grid', grid)
+            self.register_buffer("grid", grid)
 
     def forward(self, src: torch.Tensor, flow: torch.Tensor):
         if self.rigid:
             new_locs = torch.zeros((flow.shape[0], 2, 3)).to(flow.device)
-            new_locs[:, 0,0] = 1
-            new_locs[:, 1,1] = 1
-            new_locs[:, 0,2] = flow[:, 0]
-            new_locs[:, 1,2] = flow[:, 1]
-            return F.grid_sample(src, F.affine_grid(new_locs, src.size()), align_corners=True, mode="bilinear")
+            new_locs[:, 0, 0] = 1
+            new_locs[:, 1, 1] = 1
+            new_locs[:, 0, 2] = flow[:, 0]
+            new_locs[:, 1, 2] = flow[:, 1]
+            return F.grid_sample(
+                src,
+                F.affine_grid(new_locs, src.size()),
+                align_corners=True,
+                mode="bilinear",
+            )
         else:
             new_locs = self.grid + flow
             shape = flow.shape[2:]
@@ -102,13 +171,15 @@ class SpatialTransformer(torch.nn.Module):
             new_locs = new_locs.permute(0, 2, 3, 1)
             return F.grid_sample(src, new_locs[..., [1, 0]], align_corners=True, mode="bilinear")
 
+
 class VecInt(torch.nn.Module):
 
-    def __init__(self, inshape, nsteps):
+    def __init__(self, inshape: list[int], nsteps: int):
         super().__init__()
-        assert nsteps >= 0, 'nsteps should be >= 0, found: %d' % nsteps
+        if nsteps < 0:
+            raise NameError(f"nsteps should be >= 0, found: {nsteps}")
         self.nsteps = nsteps
-        self.scale = 1.0 / (2 ** self.nsteps)
+        self.scale = 1.0 / (2**self.nsteps)
         self.transformer = SpatialTransformer(inshape)
 
     def forward(self, vec: torch.Tensor):
@@ -119,16 +190,28 @@ class VecInt(torch.nn.Module):
 
 
 class ResizeTransform(torch.nn.Module):
-   
-    def __init__(self, size):
+
+    def __init__(self, size: float):
         super().__init__()
         self.factor = 1.0 / size
 
     def forward(self, x: torch.Tensor):
         if self.factor < 1:
-            x = F.interpolate(x, align_corners=True, scale_factor=self.factor, mode="bilinear", recompute_scale_factor = True)
+            x = F.interpolate(
+                x,
+                align_corners=True,
+                scale_factor=self.factor,
+                mode="bilinear",
+                recompute_scale_factor=True,
+            )
             x = self.factor * x
         elif self.factor > 1:
             x = self.factor * x
-            x = F.interpolate(x, align_corners=True, scale_factor=self.factor, mode="bilinear", recompute_scale_factor = True)
-        return x
+            x = F.interpolate(
+                x,
+                align_corners=True,
+                scale_factor=self.factor,
+                mode="bilinear",
+                recompute_scale_factor=True,
+            )
+        return x

konfai/models/representation/representation.py

@@ -1,21 +1,35 @@
 import torch
 
-from konfai.network import network, blocks
-from konfai.utils.config import config
+from konfai.network import blocks, network
+
 
 class ConvBlock(torch.nn.Module):
-    
-    def __init__(self, in_channels : int, out_channels : int) -> None:
+
+    def __init__(self, in_channels: int, out_channels: int) -> None:
         super().__init__()
-        self.Conv_0 = torch.nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=True)
+        self.Conv_0 = torch.nn.Conv3d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=True,
+        )
         self.Norm_0 = torch.nn.InstanceNorm3d(num_features=out_channels, affine=True)
         self.Activation_0 = torch.nn.LeakyReLU(negative_slope=0.01)
-        self.Conv_1 = torch.nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=True)
+        self.Conv_1 = torch.nn.Conv3d(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=True,
+        )
         self.Norm_1 = torch.nn.InstanceNorm3d(num_features=out_channels, affine=True)
         self.Activation_1 = torch.nn.LeakyReLU(negative_slope=0.01)
-        
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        output = self.Conv_0(input)
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        output = self.Conv_0(tensor)
         output = self.Norm_0(output)
         output = self.Activation_0(output)
         output = self.Conv_1(output)
@@ -23,34 +37,54 @@ class ConvBlock(torch.nn.Module):
         output = self.Activation_1(output)
         return output
 
-class UnetCPP_1_Layers(torch.nn.Module):
+
+class UnetCPP1Layers(torch.nn.Module):
 
     def __init__(self) -> None:
         super().__init__()
         self.DownConvBlock_0 = ConvBlock(in_channels=1, out_channels=32)
-        
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        return self.DownConvBlock_0(input)
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return self.DownConvBlock_0(tensor)
+
 
 class Adaptation(torch.nn.Module):
 
     def __init__(self) -> None:
         super().__init__()
-        self.Encoder_1 = UnetCPP_1_Layers()
+        self.Encoder_1 = UnetCPP1Layers()
         self.ToFeatures = blocks.ToFeatures(3)
         self.FCT_1 = torch.nn.Linear(32, 32, bias=True)
 
-    def forward(self, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> list[torch.Tensor]:
+    def forward(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         self.Encoder_1.requires_grad_(False)
         self.FCT_1.requires_grad_(True)
-        return self.FCT_1(self.ToFeatures(self.Encoder_1(A))), self.FCT_1(self.ToFeatures(self.Encoder_1(B))), self.FCT_1(self.ToFeatures(self.Encoder_1(C)))
+        return (
+            self.FCT_1(self.ToFeatures(self.Encoder_1(a))),
+            self.FCT_1(self.ToFeatures(self.Encoder_1(b))),
+            self.FCT_1(self.ToFeatures(self.Encoder_1(c))),
+        )
+
 
 class Representation(network.Network):
 
-    def __init__(   self,
-                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
-                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
-                    dim : int = 3):
-        super().__init__(in_channels = 1, optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim=dim, init_type="kaiming")
-        self.add_module("Model", Adaptation(), in_branch=[0,1,2])
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        dim: int = 3,
+    ):
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=dim,
+            init_type="kaiming",
+        )
+        self.add_module("Model", Adaptation(), in_branch=[0, 1, 2])