Add fast inference

Build.py

@@ -34,8 +34,6 @@ def convert_torchScript_full(model_name: str, model: torch.nn.Module, type: int,
     model.load_state_dict(tmp)
     torch.save({"Model" : {"Unet_TS" : tmp}}, f"{model_name}.pt")
 
-
-
 def download(url: str) -> dict[str, torch.Tensor]:
     with open(url.split("/")[-1], 'wb') as f:
         with requests.get(url, stream=True) as r:

Model.py

@@ -1,5 +1,6 @@
 import torch
 from konfai.network import network, blocks
+from konfai.predictor import Reduction
 
 class ConvBlock(network.ModuleArgsDict):
     def __init__(self, in_channels : int, out_channels : int, stride: int = 1 ) -> None:
@@ -15,15 +16,16 @@ class UNetHead(network.ModuleArgsDict):
     def __init__(self, in_channels: int, nb_class: int) -> None:
         super().__init__()
         self.add_module("Conv", torch.nn.Conv3d(in_channels = in_channels, out_channels = nb_class, kernel_size = 1, stride = 1, padding = 0))
+        self.add_module("Softmax", torch.nn.Softmax(dim=1))
     
 class UNetBlock(network.ModuleArgsDict):
 
-    def __init__(self, channels, nb_class: int, mri: bool, i : int = 0) -> None:
+    def __init__(self, channels, mri: bool, i : int = 0) -> None:
         super().__init__()
         self.add_module("DownConvBlock", ConvBlock(in_channels=channels[0], out_channels=channels[1], stride= ((1,2,2) if mri and i > 4 else 2) if i>0 else 1))
      
         if len(channels) > 2:
-            self.add_module("UNetBlock", UNetBlock(channels[1:], nb_class, mri, i+1))
+            self.add_module("UNetBlock", UNetBlock(channels[1:], mri, i+1))
             self.add_module("UpConvBlock", ConvBlock(in_channels=channels[1]*2, out_channels=channels[1]))
             
         if i > 0:
@@ -39,7 +41,6 @@ class Unet_TS(network.Network):
         },
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         channels = [1, 32, 64, 128, 320, 320], 
-        nb_class: int = 41,
         mri: bool = False) -> None:
         super().__init__(
             in_channels=channels[0],
@@ -49,5 +50,29 @@ class Unet_TS(network.Network):
             patch=None,
             dim=3,
         )
-        self.add_module("UNetBlock", UNetBlock(channels, nb_class, mri))
-        self.add_module("Head", UNetHead(channels[1], nb_class))
+        self.add_module("UNetBlock", UNetBlock(channels, mri))
+        self.add_module("Head", UNetHead(channels[1], 42))
+
+    def load(
+        self,
+        state_dict: dict[str, dict[str, torch.Tensor] | int],
+        init: bool = True,
+        ema: bool = False,
+    ):
+        nb_class, in_channels = state_dict["Model"]["Unet_TS"]["Head.Conv.weight"].shape[:2]
+        self["Head"].add_module("Conv", torch.nn.Conv3d(in_channels = in_channels, out_channels = nb_class, kernel_size = 1, stride = 1, padding = 0))
+        super().load(state_dict, init, ema)
+
+class Combine(Reduction):
+
+    def __init__(self):
+        pass
+
+    def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor:
+        fg_all = torch.cat([p[:, 1:, ...] for p in tensors], dim=1)
+        sum_fg = fg_all.sum(dim=1, keepdim=True)
+        bg = (1.0 - sum_fg).clamp(min=1e-6)
+
+        probs = torch.cat([bg, fg_all], dim=1)
+
+        return probs / probs.sum(dim=1, keepdim=True)

Prediction_CT.yml

@@ -11,7 +11,6 @@ Predictor:
       - 256
       - 320
       - 320
-      nb_class: 25
       mri: false
   Dataset:
     groups_src:
@@ -85,15 +84,15 @@ Predictor:
             dtype: uint8
             inverse: true
         dataset_filename: Dataset:mha
-        group: MASK
+        group: Seg
         same_as_group: Volume:Volume
         patch_combine: Cosinus
         inverse_transform: true
         reduction: Mean
-  train_name: Curvas
+  train_name: TotalSegmentator
   manual_seed: 32
   gpu_checkpoints: None
   images_log: None
-  combine: Mean
+  combine: Model:Combine
   autocast: false
   data_log: None

Prediction_CT_Fast.yml

@@ -0,0 +1,97 @@
+Predictor:
+  Model:
+    classpath: Model:Unet_TS
+    Unet_TS:
+      outputs_criterions: None
+      channels:
+      - 1
+      - 32
+      - 64
+      - 128
+      - 256
+      - 320
+      mri: false
+  Dataset:
+    groups_src:
+      Volume:
+        groups_dest:
+          Volume:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+              Canonical:
+                inverse: true
+              Clip:
+                min_value: -1024
+                max_value: 276
+                save_clip_min: false
+                save_clip_max: false
+                mask: None
+              Standardize:
+                lazy: false
+                mean: -370.00039267657144
+                std: 436.5998675471528
+                mask: None
+                inverse: true
+              ResampleToResolution:
+                spacing:
+                - 3
+                - 3
+                - 3
+                inverse: true
+              Padding:
+                padding:
+                - 32
+                - 32
+                - 32
+                - 32
+                - 32
+                - 32
+                mode: constant
+                inverse: true
+            patch_transforms: None
+            is_input: true
+    augmentations: None
+    Patch:
+      patch_size:
+      - 96
+      - 128
+      - 160
+      overlap: 32
+      mask: None
+      pad_value: 0
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:nii.gz
+    use_cache: false
+    batch_size: 1
+  outputs_dataset:
+    Head:Conv:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          Softmax:
+            dim: 0
+          Argmax:
+            dim: 0
+          TensorCast:
+            dtype: uint8
+            inverse: true
+        dataset_filename: Dataset:mha
+        group: Seg
+        same_as_group: Volume:Volume
+        patch_combine: Cosinus
+        inverse_transform: true
+        reduction: Mean
+  train_name: TotalSegmentator
+  manual_seed: 32
+  gpu_checkpoints: None
+  images_log: None
+  combine: Model:Combine
+  autocast: false
+  data_log: None

Prediction_MR.yml

@@ -11,7 +11,7 @@ Predictor:
       - 256
       - 320
       - 320
-      nb_class: 30
+      mri: true
   Dataset:
     groups_src:
       Volume:
@@ -78,15 +78,15 @@ Predictor:
             dtype: uint8
             inverse: true
         dataset_filename: Dataset:mha
-        group: MASK
+        group: Seg
         same_as_group: Volume:Volume
         patch_combine: Cosinus
         inverse_transform: true
         reduction: Mean
-  train_name: Curvas
+  train_name: TotalSegmentator
   manual_seed: 32
   gpu_checkpoints: None
   images_log: None
-  combine: Mean
+  combine: Model:Combine
   autocast: false
   data_log: None

Prediction_MR_Fast.yml

@@ -0,0 +1,91 @@
+Predictor:
+  Model:
+    classpath: Model:Unet_TS
+    Unet_TS:
+      outputs_criterions: None
+      channels:
+      - 1
+      - 32
+      - 64
+      - 128
+      - 256
+      - 320
+      mri: true
+  Dataset:
+    groups_src:
+      Volume:
+        groups_dest:
+          Volume:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+              Canonical:
+                inverse: true
+              Standardize:
+                lazy: false
+                mean: None
+                std: None
+                mask: None
+                inverse: false
+              ResampleToResolution:
+                spacing:
+                - 3
+                - 3
+                - 3
+                inverse: true
+              Padding:
+                padding:
+                - 32
+                - 32
+                - 32
+                - 32
+                - 32
+                - 32
+                mode: constant
+                inverse: true
+            patch_transforms: None
+            is_input: true
+    augmentations: None
+    Patch:
+      patch_size:
+      - 96
+      - 128
+      - 160
+      overlap: 32
+      mask: None
+      pad_value: 0
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:nii.gz
+    use_cache: false
+    batch_size: 1
+  outputs_dataset:
+    Head:Conv:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          Softmax:
+            dim: 0
+          Argmax:
+            dim: 0
+          TensorCast:
+            dtype: uint8
+            inverse: true
+        dataset_filename: Dataset:mha
+        group: Seg
+        same_as_group: Volume:Volume
+        patch_combine: Cosinus
+        inverse_transform: true
+        reduction: Mean
+  train_name: TotalSegmentator
+  manual_seed: 32
+  gpu_checkpoints: None
+  images_log: None
+  combine: Model:Combine
+  autocast: false
+  data_log: None