Fix clipping and normalization in CT-specific model preprocessing

Build.py → build.py

@@ -6,15 +6,20 @@ import shutil
 from pathlib import Path
 import os
 from functools import partial
-from Model import Unet_TS
+from total.model import Unet_TS_CT
+from total_mr.model import Unet_TS_MR
+import json
 
-def convert_torchScript_full(model_name: str, model: torch.nn.Module, type: int, url: str) -> None:
-    state_dict = download(url)
+def convert_torchScript_full(model_name: str, model: torch.nn.Module, task: str, type: int, mri: bool, url: str) -> None:
+    state_dict = download(url, model_name, mri)
     tmp = {}
     with open("Destination_Unet_{}.txt".format(type)) as f2:
         it = iter(state_dict.keys())
         for l1 in f2:
+            print(l1)
+
             key = next(it)
+            print(key)
             while "decoder.seg_layers" in key:
                 if type == 1:
                     if "decoder.seg_layers.4" in key :
@@ -30,17 +35,24 @@ def convert_torchScript_full(model_name: str, model: torch.nn.Module, type: int,
             while "all_modules" in key or "decoder.encoder" in key:
                 key = next(it)
             tmp[l1.replace("\n", "")] = state_dict[key]
+    if not mri:
+        tmp["ClipAndNormalize.mean"] = state_dict["mean"]
+        tmp["ClipAndNormalize.std"] = state_dict["std"]
+        tmp["ClipAndNormalize.clip_min"] = state_dict["percentile_00_5"]
+        tmp["ClipAndNormalize.clip_max"] = state_dict["percentile_99_5"]
+    state_dict = {"Model" : {model.name : tmp}}
+    model.load(state_dict)
+    dest_path = Path(f"./{task}")
+    dest_path.mkdir(exist_ok=True)
+    torch.save(state_dict, str(dest_path/f"{model_name}.pt"))
 
-    model.load_state_dict(tmp)
-    torch.save({"Model" : {"Unet_TS" : tmp}}, f"{model_name}.pt")
-
-def download(url: str) -> dict[str, torch.Tensor]:
+def download(url: str, model_name: str, mri: bool) -> dict[str, torch.Tensor]:
     with open(url.split("/")[-1], 'wb') as f:
         with requests.get(url, stream=True) as r:
             r.raise_for_status()
 
             total_size = int(r.headers.get('content-length', 0))
-            progress_bar = tqdm(total=total_size, unit='B', unit_scale=True, desc="Downloading")
+            progress_bar = tqdm(total=total_size, unit='B', unit_scale=True, desc=f"Downloading {model_name}")
             for chunk in r.iter_content(chunk_size=8192 * 16):
                 progress_bar.update(len(chunk))
                 f.write(chunk)
@@ -48,33 +60,50 @@ def download(url: str) -> dict[str, torch.Tensor]:
     with zipfile.ZipFile(url.split("/")[-1], 'r') as zip_f:
         zip_f.extractall(url.split("/")[-1].replace(".zip", ""))
     os.remove(url.split("/")[-1])
-    state_dict = torch.load(next(Path(url.split("/")[-1].replace(".zip", "")).rglob("checkpoint_final.pth"), None), weights_only=False)["network_weights"]
-    shutil.rmtree(url.split("/")[-1].replace(".zip", ""))
+    zip_path = Path(url.split("/")[-1].replace(".zip", ""))
+
+    state_dict = torch.load(next(zip_path.rglob("checkpoint_final.pth"), None), map_location="cpu", weights_only=False)["network_weights"]
+    if not mri:
+        dataset_fingerprint_path = next(zip_path.rglob("dataset_fingerprint.json"), None)
+        with open(dataset_fingerprint_path, "r") as f:
+            data = json.load(f)
+
+        ch0 = data["foreground_intensity_properties_per_channel"]["0"]
+
+        state_dict["mean"] = torch.tensor([ch0["mean"]])
+        state_dict["std"] = torch.tensor([ch0["std"]])
+        state_dict["percentile_00_5"] = torch.tensor([ch0["percentile_00_5"]])
+        state_dict["percentile_99_5"] = torch.tensor([ch0["percentile_99_5"]])
+    shutil.rmtree(zip_path)
     return state_dict
 
 url = "https://github.com/wasserth/TotalSegmentator/releases/download/"
 
-UnetCPP_1 = partial(Unet_TS, channels = [1,32,64,128,256,320,320])
-UnetCPP_2 = partial(Unet_TS, channels = [1,32,64,128,256,320])
-UnetCPP_3 = partial(Unet_TS, channels = [1,32,64,128,256])
+UnetCPP_1_CT = partial(Unet_TS_CT, channels = [1,32,64,128,256,320,320])
+UnetCPP_2_CT = partial(Unet_TS_CT, channels = [1,32,64,128,256,320])
+UnetCPP_3_CT = partial(Unet_TS_CT, channels = [1,32,64,128,256])
 
-models = {
-         "M291" : (UnetCPP_1(nb_class=25), 1, url+"v2.0.0-weights/Dataset291_TotalSegmentator_part1_organs_1559subj.zip"), 
-         "M292" : (UnetCPP_1(nb_class=27), 1, url+"v2.0.0-weights/Dataset292_TotalSegmentator_part2_vertebrae_1532subj.zip"),
-         "M293" : (UnetCPP_1(nb_class=19), 1, url+"v2.0.0-weights/Dataset293_TotalSegmentator_part3_cardiac_1559subj.zip"),
-         "M294" : (UnetCPP_1(nb_class=24), 1, url+"v2.0.0-weights/Dataset294_TotalSegmentator_part4_muscles_1559subj.zip"),
-         "M295" : (UnetCPP_1(nb_class=27), 1, url+"v2.0.0-weights/Dataset295_TotalSegmentator_part5_ribs_1559subj.zip"),
-         "M297" : (UnetCPP_2(nb_class=118), 2, url+"v2.0.4-weights/Dataset297_TotalSegmentator_total_3mm_1559subj_v204.zip"),
-         "M298" : (UnetCPP_2(nb_class=118), 2, url+"v2.0.0-weights/Dataset298_TotalSegmentator_total_6mm_1559subj.zip"),
-         "M730" : (UnetCPP_1(nb_class=30, mri = True), 1, url+"v2.2.0-weights/Dataset730_TotalSegmentatorMRI_part1_organs_495subj.zip"),
-         "M731" : (UnetCPP_1(nb_class=28, mri = True), 1, url+"v2.2.0-weights/Dataset731_TotalSegmentatorMRI_part2_muscles_495subj.zip"),
-         "M732" : (UnetCPP_2(nb_class=57), 2, url+"v2.2.0-weights/Dataset732_TotalSegmentatorMRI_total_3mm_495subj.zip"),
-         "M733" : (UnetCPP_3(nb_class=57), 3, url+"v2.2.0-weights/Dataset733_TotalSegmentatorMRI_total_6mm_495subj.zip"),
-         "M850" : (UnetCPP_1(nb_class=30, mri = True), 1, url+"v2.5.0-weights/Dataset850_TotalSegMRI_part1_organs_1088subj.zip"),
-         "M851" : (UnetCPP_1(nb_class=22, mri = True), 1, url+"v2.5.0-weights/Dataset851_TotalSegMRI_part2_muscles_1088subj.zip"),
-         "M852" : (UnetCPP_2(nb_class=51), 2, url+"v2.5.0-weights/Dataset852_TotalSegMRI_total_3mm_1088subj.zip"),
-         "M853" : (UnetCPP_3(nb_class=51), 3, url+"v2.5.0-weights/Dataset853_TotalSegMRI_total_6mm_1088subj.zip")}
+UnetCPP_1_MR = partial(Unet_TS_MR, channels = [1,32,64,128,256,320,320])
+UnetCPP_2_MR = partial(Unet_TS_MR, channels = [1,32,64,128,256,320])
+UnetCPP_3_MR = partial(Unet_TS_MR, channels = [1,32,64,128,256])
 
+models = {
+         #"M291" : (UnetCPP_1_CT(), "total", 1, False, url+"v2.0.0-weights/Dataset291_TotalSegmentator_part1_organs_1559subj.zip"), 
+         #"M292" : (UnetCPP_1_CT(), "total", 1, False, url+"v2.0.0-weights/Dataset292_TotalSegmentator_part2_vertebrae_1532subj.zip"),
+         #"M293" : (UnetCPP_1_CT(), "total", 1, False, url+"v2.0.0-weights/Dataset293_TotalSegmentator_part3_cardiac_1559subj.zip"),
+         #"M294" : (UnetCPP_1_CT(), "total", 1, False, url+"v2.0.0-weights/Dataset294_TotalSegmentator_part4_muscles_1559subj.zip"),
+         #"M295" : (UnetCPP_1_CT(), "total", 1, False, url+"v2.0.0-weights/Dataset295_TotalSegmentator_part5_ribs_1559subj.zip"),
+         #"M297" : (UnetCPP_2_CT(), "total-3mm", 2, False, url+"v2.0.4-weights/Dataset297_TotalSegmentator_total_3mm_1559subj_v204.zip"),
+         #"M298" : (UnetCPP_2_CT(), 2, False, url+"v2.0.0-weights/Dataset298_TotalSegmentator_total_6mm_1559subj.zip"),
+         #"M730" : (UnetCPP_1_MR(), True, 1, url+"v2.2.0-weights/Dataset730_TotalSegmentatorMRI_part1_organs_495subj.zip"),
+         #"M731" : (UnetCPP_1_MR(), True, 1, url+"v2.2.0-weights/Dataset731_TotalSegmentatorMRI_part2_muscles_495subj.zip"),
+         #"M732" : (UnetCPP_2_MR(), False, 2, url+"v2.2.0-weights/Dataset732_TotalSegmentatorMRI_total_3mm_495subj.zip"),
+         #"M733" : (UnetCPP_3_MR(), False, 3, url+"v2.2.0-weights/Dataset733_TotalSegmentatorMRI_total_6mm_495subj.zip"),
+         #"M850" : (UnetCPP_1_MR(), "total_mr", 1, True, url+"v2.5.0-weights/Dataset850_TotalSegMRI_part1_organs_1088subj.zip"),
+         #"M851" : (UnetCPP_1_MR(), "total_mr", 1, True, url+"v2.5.0-weights/Dataset851_TotalSegMRI_part2_muscles_1088subj.zip"),
+         "M852" : (UnetCPP_2_MR(), "total_mr-3mm", 2, True, url+"v2.5.0-weights/Dataset852_TotalSegMRI_total_3mm_1088subj.zip"),
+         #"M853" : (UnetCPP_3_MR(), False, 3, url+"v2.5.0-weights/Dataset853_TotalSegMRI_total_6mm_1088subj.zip")
+         }
 if __name__ == "__main__":
     for name, model in models.items():
-        convert_torchScript_full(name, model[0], model[1], model[2])
+        convert_torchScript_full(name, *model)

total-3mm/M297.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6490a06ac9d242757af99fc674ba8f74ecb5c62009ecda6c62cda217b9cbbcb7
-size 66225317
+oid sha256:d50dab1ce25e6d55f9e7b2f6a5a6282293e6c9add63278558149e8a1146f5f5b
+size 66225357

total-3mm/Prediction.yml

@@ -1,7 +1,7 @@
 Predictor:
   Model:
-    classpath: Model:Unet_TS
-    Unet_TS:
+    classpath: model:Unet_TS_CT
+    Unet_TS_CT:
       outputs_criterions: None
       channels:
       - 1
@@ -10,7 +10,6 @@ Predictor:
       - 128
       - 256
       - 320
-      mri: false
   Dataset:
     groups_src:
       Volume_0:
@@ -22,24 +21,18 @@ Predictor:
                 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: false
               ResampleToResolution:
                 spacing:
                 - 3
                 - 3
                 - 3
                 inverse: true
+              Standardize:
+                lazy: false
+                mean: None
+                std: None
+                mask: None
+                inverse: false
               Padding:
                 padding:
                 - 32

total-3mm/app.json

@@ -1,5 +1,5 @@
 {
-    "display_name": "TotalSegmentator 3mm",
+    "display_name": "Segmentation: TotalSegmentator 3mm",
     "short_description": "<b>Description:</b><br>Lightweight KonfAI adaptation of <a href='https://github.com/wasserth/TotalSegmentator'>TotalSegmentator</a> trained at <b>3 mm resolution</b>, reducing GPU/RAM requirements while segmenting <b>118 anatomical structures</b> in whole-body CT.<br><br><b>How to cite:</b><br><cite>J. Wasserthal et al., <i>TotalSegmentator: Robust Segmentation of 104 Anatomical Structures in CT Images</i>, Radiology: AI, 2023.</cite>",
     "description": "<b>Description:</b><br>KonfAI-optimized version of the original nnU-Net-based TotalSegmentator 3 mm model.<br><br><b>Capabilities:</b><br>• Whole-body CT segmentation of <b>118 structures</b> (organs, bones, muscles, vessels)<br>• Reduced computational footprint for lower memory and faster throughput<br>• <b>3 mm isotropic</b> inference for easier deployment on large datasets<br><br><b>Training data:</b><br>Trained on <b>1204 clinically-derived CT scans</b> with strong diversity in contrast phases, scanner types and pathologies, with expert-reviewed manual annotations<br><br><b>How to cite:</b><br><cite>J. Wasserthal et al., <i>TotalSegmentator: Robust Segmentation of 104 Anatomical Structures in CT Images</i>, Radiology: AI, 2023.</cite>",
     "tta": 0,

total/Model.py → total-3mm/model.py

@@ -1,6 +1,5 @@
 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:
@@ -20,19 +19,32 @@ class UNetHead(network.ModuleArgsDict):
     
 class UNetBlock(network.ModuleArgsDict):
 
-    def __init__(self, channels, mri: bool, i : int = 0) -> None:
+    def __init__(self, channels, 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))
+        self.add_module("DownConvBlock", ConvBlock(in_channels=channels[0], out_channels=channels[1], stride= 2 if i>0 else 1))
      
         if len(channels) > 2:
-            self.add_module("UNetBlock", UNetBlock(channels[1:], mri, i+1))
+            self.add_module("UNetBlock", UNetBlock(channels[1:], i+1))
             self.add_module("UpConvBlock", ConvBlock(in_channels=channels[1]*2, out_channels=channels[1]))
             
         if i > 0:
-            self.add_module("CONV_TRANSPOSE", torch.nn.ConvTranspose3d(in_channels = channels[1], out_channels = channels[0], kernel_size =  (1,2,2) if mri and i > 4 else 2, stride = (1,2,2) if mri and i > 4 else 2, padding = 0))
+            self.add_module("CONV_TRANSPOSE", torch.nn.ConvTranspose3d(in_channels = channels[1], out_channels = channels[0], kernel_size =  2, stride = 2, padding = 0))
             self.add_module("SkipConnection", blocks.Concat(), in_branch=[0, 1])
 
-class Unet_TS(network.Network):
+class ClipAndNormalize(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.register_buffer("clip_min", torch.empty(1))
+        self.register_buffer("clip_max", torch.empty(1))
+        self.register_buffer("mean", torch.empty(1))
+        self.register_buffer("std", torch.empty(1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.clamp(x, self.clip_min, self.clip_max)
+        return (x - self.mean) / (self.std)
+
+class Unet_TS_CT(network.Network):
 
     def __init__(self,
         optimizer: network.OptimizerLoader = network.OptimizerLoader(),
@@ -40,8 +52,7 @@ class Unet_TS(network.Network):
             "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
         },
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        channels = [1, 32, 64, 128, 320, 320], 
-        mri: bool = False) -> None:
+        channels: list[int] = [1, 32, 64, 128, 320]) -> None:
         super().__init__(
             in_channels=channels[0],
             optimizer=optimizer,
@@ -50,8 +61,9 @@ class Unet_TS(network.Network):
             patch=None,
             dim=3,
         )
-        self.add_module("UNetBlock", UNetBlock(channels, mri))
-        self.add_module("Head", UNetHead(channels[1], 42))
+        self.add_module("ClipAndNormalize", ClipAndNormalize())
+        self.add_module("UNetBlock", UNetBlock(channels))
+        self.add_module("Head", UNetHead(channels[1], 118))
 
     def load(
         self,
@@ -59,6 +71,6 @@ class Unet_TS(network.Network):
         init: bool = True,
         ema: bool = False,
     ):
-        nb_class, in_channels = state_dict["Model"]["Unet_TS"]["Head.Conv.weight"].shape[:2]
+        nb_class, in_channels = state_dict["Model"]["Unet_TS_CT"]["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)
+        super().load(state_dict, init, ema)

total/M291.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:78065b11af6e339feaa49c72c3aa45d78ae272be1432b27bc17285428a851e6a
-size 124807929
+oid sha256:a28e7fe8660329f225f93dbb61feefa0037a5734de87f2a35ab18d2ddfb7601c
+size 124807717

total/M292.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1e0a8d8c17572b392bfdbc480edd169ac46fd8962fdf30b81de2406a6f32f275
-size 124808185
+oid sha256:374ffdb916c267ab5f96ca9f2b754ff17eff9b48762de686e9524e0ade5909f3
+size 124807973

total/M293.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9638deb528c1dafc8edfe5287404eefc764770c913810628cd62b23b5950a4c0
-size 124807161
+oid sha256:3104f4cc3c5796468a9d0d68fee602786eeb414a2a28ff2a8bd7f2210604daed
+size 124806949

total/M294.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:11879030dad493ff88d49cbcbbbbee1f5671c79df743ae931b1bd7bbf7302b5e
-size 124807801
+oid sha256:8293c863503c8fa858f2eee7f599717b0cc4c16a50b0500107af1e8d905fe669
+size 124807589

total/M295.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4e586a1082dcdb3a054e2ccd670d2506c7436e1f0f6f616cdc771bfe5c41d948
-size 124808185
+oid sha256:01bc78eee961a71a675c7335c88622b1524151982b546113fe1e9a4e3a2f8032
+size 124807973

total/Prediction.yml

@@ -1,7 +1,7 @@
 Predictor:
   Model:
-    classpath: Model:Unet_TS
-    Unet_TS:
+    classpath: model:Unet_TS_CT
+    Unet_TS_CT:
       outputs_criterions: None
       channels:
       - 1
@@ -11,7 +11,6 @@ Predictor:
       - 256
       - 320
       - 320
-      mri: false
   Dataset:
     groups_src:
       Volume_0:
@@ -23,18 +22,6 @@ Predictor:
                 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: false
               ResampleToResolution:
                 spacing:
                 - 1.5

total/app.json

@@ -1,5 +1,5 @@
 {
-    "display_name": "Total Segmentator",
+    "display_name": "Segmentation: Total Segmentator",
     "short_description": "<b>Description:</b><br>KonfAI-accelerated adaptation of <a href='https://github.com/wasserth/TotalSegmentator'>TotalSegmentator</a>, delivering fast whole-body CT segmentation of <b>118 anatomical structures</b> (1.5 mm resolution) with reduced inference cost compared to the original nnU-Net implementation.<br><br><b>How to cite:</b><br><cite>J. Wasserthal et al., <i>TotalSegmentator: Robust Segmentation of 104 Anatomical Structures in CT Images</i>, Radiology: AI, 2023.</cite>",
     "description": "<b>Description:</b><br>This model is an optimized adaptation of the original <a href='https://github.com/wasserth/TotalSegmentator'>TotalSegmentator</a> for the <b>KonfAI</b> deep learning framework.<br><br><b>Capabilities:</b><br>• Segmentation of <b>118 anatomical classes</b> covering organs, bones, muscles, and vessels<br>• Enhanced runtime and memory efficiency vs. the original nnU-Net implementation<br>• High-resolution input: <b>1.5 mm isotropic</b><br><br><b>Training data:</b><br>Trained on a diverse dataset of <b>1204 whole-body CT examinations</b> including different scanners, acquisition settings, contrast phases, and major pathologies (27 organs, 59 bones, 10 muscles, 8 vessels), with manual expert-reviewed annotations<br><br><b>How to cite:</b><br><cite>J. Wasserthal et al., <i>TotalSegmentator: Robust Segmentation of 104 Anatomical Structures in CT Images</i>, Radiology: AI, 2023.</cite>",
     "tta": 0,

total_mr-3mm/Model.py → total/model.py

@@ -1,6 +1,5 @@
 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:
@@ -20,19 +19,32 @@ class UNetHead(network.ModuleArgsDict):
     
 class UNetBlock(network.ModuleArgsDict):
 
-    def __init__(self, channels, mri: bool, i : int = 0) -> None:
+    def __init__(self, channels, 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))
+        self.add_module("DownConvBlock", ConvBlock(in_channels=channels[0], out_channels=channels[1], stride= 2 if i>0 else 1))
      
         if len(channels) > 2:
-            self.add_module("UNetBlock", UNetBlock(channels[1:], mri, i+1))
+            self.add_module("UNetBlock", UNetBlock(channels[1:], i+1))
             self.add_module("UpConvBlock", ConvBlock(in_channels=channels[1]*2, out_channels=channels[1]))
             
         if i > 0:
-            self.add_module("CONV_TRANSPOSE", torch.nn.ConvTranspose3d(in_channels = channels[1], out_channels = channels[0], kernel_size =  (1,2,2) if mri and i > 4 else 2, stride = (1,2,2) if mri and i > 4 else 2, padding = 0))
+            self.add_module("CONV_TRANSPOSE", torch.nn.ConvTranspose3d(in_channels = channels[1], out_channels = channels[0], kernel_size =  2, stride = 2, padding = 0))
             self.add_module("SkipConnection", blocks.Concat(), in_branch=[0, 1])
 
-class Unet_TS(network.Network):
+class ClipAndNormalize(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.register_buffer("clip_min", torch.empty(1))
+        self.register_buffer("clip_max", torch.empty(1))
+        self.register_buffer("mean", torch.empty(1))
+        self.register_buffer("std", torch.empty(1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.clamp(x, self.clip_min, self.clip_max)
+        return (x - self.mean) / (self.std)
+
+class Unet_TS_CT(network.Network):
 
     def __init__(self,
         optimizer: network.OptimizerLoader = network.OptimizerLoader(),
@@ -40,8 +52,7 @@ class Unet_TS(network.Network):
             "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
         },
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        channels = [1, 32, 64, 128, 320, 320], 
-        mri: bool = False) -> None:
+        channels: list[int] = [1, 32, 64, 128, 320, 320]) -> None:
         super().__init__(
             in_channels=channels[0],
             optimizer=optimizer,
@@ -50,8 +61,9 @@ class Unet_TS(network.Network):
             patch=None,
             dim=3,
         )
-        self.add_module("UNetBlock", UNetBlock(channels, mri))
-        self.add_module("Head", UNetHead(channels[1], 42))
+        self.add_module("ClipAndNormalize", ClipAndNormalize())
+        self.add_module("UNetBlock", UNetBlock(channels))
+        self.add_module("Head", UNetHead(channels[1], 118))
 
     def load(
         self,
@@ -59,6 +71,6 @@ class Unet_TS(network.Network):
         init: bool = True,
         ema: bool = False,
     ):
-        nb_class, in_channels = state_dict["Model"]["Unet_TS"]["Head.Conv.weight"].shape[:2]
+        nb_class, in_channels = state_dict["Model"]["Unet_TS_CT"]["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)
+        super().load(state_dict, init, ema)

total_mr-3mm/M852.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:22f810c3c77079ca3a196f577d9ab21733aa1e94cb006896c9d35d7d7f726157
-size 66216485
+oid sha256:4d74547add966128f251df5d76f61f151800990b9c4ee55aaa60b4646a9a3ea8
+size 66215333

total_mr-3mm/Prediction.yml

@@ -22,18 +22,18 @@ Predictor:
                 inverse: false
               Canonical:
                 inverse: true
-              Standardize:
-                lazy: false
-                mean: None
-                std: None
-                mask: None
-                inverse: false
               ResampleToResolution:
                 spacing:
                 - 3
                 - 3
                 - 3
                 inverse: true
+              Standardize:
+                lazy: false
+                mean: None
+                std: None
+                mask: None
+                inverse: false
               Padding:
                 padding:
                 - 32

total_mr-3mm/app.json

@@ -1,5 +1,5 @@
 {
-    "display_name": "TotalSegmentator MRI 3mm",
+    "display_name": "Segmentation: TotalSegmentator MRI 3mm",
     "short_description": "<b>Description:</b><br>Lightweight KonfAI adaptation of <a href='https://github.com/wasserth/TotalSegmentator'>TotalSegmentator MRI</a>, enabling fast multimodal segmentation of <b>50 key anatomical structures</b> from <b>MRI and CT</b> scans at <b>3 mm</b> resolution, greatly reducing memory usage and inference time compared to the original nnU-Net workflow.<br><br><b>How to cite:</b><br><cite>T. Akinci D’Antonoli et al., <i>TotalSegmentator MRI: Robust Sequence-Independent Segmentation of Multiple Anatomic Structures in MRI</i>, Radiology, 2025.</cite>",
     "description": "<b>Description:</b><br>This model integrates the reduced-resolution MRI-3mm configuration of TotalSegmentator into the <b>KonfAI</b> accelerated inference pipeline for efficient MRI/CT deployment.<br><br><b>Capabilities:</b><br>• Segmentation of <b>50 essential anatomical structures</b> (major organs, key bones, large vessels)<br>• <b>3 mm</b> isotropic input for high-throughput processing and lower GPU requirements<br>• Robust to acquisition variability including scanner type, contrast, and sequence parameters<br><br><b>Training data:</b><br>Trained on a diverse cohort of <b>1143 clinical scans</b> including <b>616 MRI</b> (multi-site, multi-scanner, multi-sequence) and <b>527 CT</b>, with expert-validated reference masks<br><br>><b>How to cite:</b><br><cite>T. Akinci D’Antonoli et al., <i>TotalSegmentator MRI: Robust Sequence-Independent Segmentation of Multiple Anatomic Structures in MRI</i>, Radiology, 2025.</cite>",
     "tta": 0,

total_mr/Model.py → total_mr-3mm/model.py

@@ -1,6 +1,5 @@
 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:
@@ -20,19 +19,19 @@ class UNetHead(network.ModuleArgsDict):
     
 class UNetBlock(network.ModuleArgsDict):
 
-    def __init__(self, channels, mri: bool, i : int = 0) -> None:
+    def __init__(self, channels, 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))
+        self.add_module("DownConvBlock", ConvBlock(in_channels=channels[0], out_channels=channels[1], stride= ((1,2,2) if i > 4 else 2) if i>0 else 1))
      
         if len(channels) > 2:
-            self.add_module("UNetBlock", UNetBlock(channels[1:], mri, i+1))
+            self.add_module("UNetBlock", UNetBlock(channels[1:], i+1))
             self.add_module("UpConvBlock", ConvBlock(in_channels=channels[1]*2, out_channels=channels[1]))
             
         if i > 0:
-            self.add_module("CONV_TRANSPOSE", torch.nn.ConvTranspose3d(in_channels = channels[1], out_channels = channels[0], kernel_size =  (1,2,2) if mri and i > 4 else 2, stride = (1,2,2) if mri and i > 4 else 2, padding = 0))
+            self.add_module("CONV_TRANSPOSE", torch.nn.ConvTranspose3d(in_channels = channels[1], out_channels = channels[0], kernel_size =  (1,2,2) if i > 4 else 2, stride = (1,2,2) if i > 4 else 2, padding = 0))
             self.add_module("SkipConnection", blocks.Concat(), in_branch=[0, 1])
 
-class Unet_TS(network.Network):
+class Unet_TS_MR(network.Network):
 
     def __init__(self,
         optimizer: network.OptimizerLoader = network.OptimizerLoader(),
@@ -40,8 +39,7 @@ class Unet_TS(network.Network):
             "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
         },
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        channels = [1, 32, 64, 128, 320, 320], 
-        mri: bool = False) -> None:
+        channels: list[int] = [1, 32, 64, 128, 320]) -> None:
         super().__init__(
             in_channels=channels[0],
             optimizer=optimizer,
@@ -50,7 +48,7 @@ class Unet_TS(network.Network):
             patch=None,
             dim=3,
         )
-        self.add_module("UNetBlock", UNetBlock(channels, mri))
+        self.add_module("UNetBlock", UNetBlock(channels))
         self.add_module("Head", UNetHead(channels[1], 42))
 
     def load(
@@ -59,6 +57,6 @@ class Unet_TS(network.Network):
         init: bool = True,
         ema: bool = False,
     ):
-        nb_class, in_channels = state_dict["Model"]["Unet_TS"]["Head.Conv.weight"].shape[:2]
+        nb_class, in_channels = state_dict["Model"]["Unet_TS_MR"]["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)
+        super().load(state_dict, init, ema)

total_mr/M850.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dbcf0d153ff695b696748e6192a912b305f0b8deba3d04fec1d685ea39df4c37
-size 123170169
+oid sha256:bc61844cc7053882deae8ca0b2deb89d369528c44f73c935687068968ed39f87
+size 123168825

total_mr/M851.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:570f5d944fd07c26c89dd41cd105dbe470286bc91d01d92e3ba63d9904f03903
-size 123169145
+oid sha256:8ced0be1dfb7e48d6e21261ac3969d5d421d6a07b5228af6b2df55cfb6972b73
+size 123167801

total_mr/Prediction.yml

@@ -1,7 +1,7 @@
 Predictor:
   Model:
-    classpath: Model:Unet_TS
-    Unet_TS:
+    classpath: model:Unet_TS_MR
+    Unet_TS_MR:
       outputs_criterions: None
       channels:
       - 1
@@ -11,7 +11,6 @@ Predictor:
       - 256
       - 320
       - 320
-      mri: true
   Dataset:
     groups_src:
       Volume_0:

total_mr/app.json

@@ -1,5 +1,5 @@
 {
-    "display_name": "TotalSegmentator MRI",
+    "display_name": "Segmentation: TotalSegmentator MRI",
     "short_description": "<b>Description:</b><br>KonfAI-accelerated adaptation of <a href='https://github.com/wasserth/TotalSegmentator'>TotalSegmentator MRI</a>, delivering fast multimodal segmentation of <b>80 major anatomical structures</b> from <b>MRI and CT</b> scans, with significantly reduced inference overhead vs. the original nnU-Net workflow.<br><br><b>How to cite:</b><br><cite>T. Akinci D’Antonoli et al., <i>TotalSegmentator MRI: Robust Sequence-Independent Segmentation of Multiple Anatomic Structures in MRI</i>, Radiology, 2025.</cite>",
     "description": "<b>Description:</b><br>This model integrates TotalSegmentator MRI into the <b>KonfAI</b> inference framework to accelerate deployment in MRI/CT multimodal workflows.<br><br><b>Capabilities:</b><br>• Automatic segmentation of <b>80 major anatomical structures</b> (organs, vessels, skeleton, digestive system)<br>• Robust to <b>sequence variations</b> across scanners, contrasts, acquisition planes, and sites<br>• High-resolution input: <b>1.5 mm isotropic</b><br><br><b>Training data:</b><br>Trained on a highly diverse clinical dataset of <b>1143 scans</b> including <b>616 MRI</b> (30 scanners, 4 sites, many contrast types) and <b>527 CT</b> scans, with expert-validated manual segmentations <br><br><b>How to cite:</b><br><cite>T. Akinci D’Antonoli et al., <i>TotalSegmentator MRI: Robust Sequence-Independent Segmentation of Multiple Anatomic Structures in MRI</i>, Radiology, 2025.</cite>",
     "tta": 0,

total-3mm/Model.py → total_mr/model.py

@@ -1,6 +1,5 @@
 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:
@@ -20,19 +19,19 @@ class UNetHead(network.ModuleArgsDict):
     
 class UNetBlock(network.ModuleArgsDict):
 
-    def __init__(self, channels, mri: bool, i : int = 0) -> None:
+    def __init__(self, channels, 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))
+        self.add_module("DownConvBlock", ConvBlock(in_channels=channels[0], out_channels=channels[1], stride= ((1,2,2) if i > 4 else 2) if i>0 else 1))
      
         if len(channels) > 2:
-            self.add_module("UNetBlock", UNetBlock(channels[1:], mri, i+1))
+            self.add_module("UNetBlock", UNetBlock(channels[1:], i+1))
             self.add_module("UpConvBlock", ConvBlock(in_channels=channels[1]*2, out_channels=channels[1]))
             
         if i > 0:
-            self.add_module("CONV_TRANSPOSE", torch.nn.ConvTranspose3d(in_channels = channels[1], out_channels = channels[0], kernel_size =  (1,2,2) if mri and i > 4 else 2, stride = (1,2,2) if mri and i > 4 else 2, padding = 0))
+            self.add_module("CONV_TRANSPOSE", torch.nn.ConvTranspose3d(in_channels = channels[1], out_channels = channels[0], kernel_size =  (1,2,2) if i > 4 else 2, stride = (1,2,2) if i > 4 else 2, padding = 0))
             self.add_module("SkipConnection", blocks.Concat(), in_branch=[0, 1])
 
-class Unet_TS(network.Network):
+class Unet_TS_MR(network.Network):
 
     def __init__(self,
         optimizer: network.OptimizerLoader = network.OptimizerLoader(),
@@ -40,8 +39,7 @@ class Unet_TS(network.Network):
             "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
         },
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        channels = [1, 32, 64, 128, 320, 320], 
-        mri: bool = False) -> None:
+        channels: list[int] = [1, 32, 64, 128, 320, 320]) -> None:
         super().__init__(
             in_channels=channels[0],
             optimizer=optimizer,
@@ -50,7 +48,7 @@ class Unet_TS(network.Network):
             patch=None,
             dim=3,
         )
-        self.add_module("UNetBlock", UNetBlock(channels, mri))
+        self.add_module("UNetBlock", UNetBlock(channels))
         self.add_module("Head", UNetHead(channels[1], 42))
 
     def load(
@@ -59,6 +57,6 @@ class Unet_TS(network.Network):
         init: bool = True,
         ema: bool = False,
     ):
-        nb_class, in_channels = state_dict["Model"]["Unet_TS"]["Head.Conv.weight"].shape[:2]
+        nb_class, in_channels = state_dict["Model"]["Unet_TS_MR"]["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)
+        super().load(state_dict, init, ema)