Update body MR/CT 6mm

body/{M299.pt → M598.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8fb6a978453e6f8d2cd8ec42500f7946b53648db8043c6f2504c042910a366c5
-size 124804837
+oid sha256:c0fbbaf9375502a08f765750374c8dc31736cf938966ed8eb8f2cf6b919e1933
+size 22427409

body/Prediction.yml

@@ -1,7 +1,7 @@
 Predictor:
   Model:
-    classpath: model:Unet_TS_CT
-    Unet_TS_CT:
+    classpath: model:Unet_TS_MR
+    Unet_TS_MR:
       outputs_criterions: None
       channels:
       - 1
@@ -9,8 +9,6 @@ Predictor:
       - 64
       - 128
       - 256
-      - 320
-      - 320
   Dataset:
     groups_src:
       Volume_0:
@@ -24,10 +22,16 @@ Predictor:
                 inverse: true
               ResampleToResolution:
                 spacing:
-                - 1.5
-                - 1.5
-                - 1.5
+                - 6
+                - 6
+                - 6
                 inverse: true
+              Standardize:
+                lazy: false
+                mean: None
+                std: None
+                mask: None
+                inverse: false
               Padding:
                 padding:
                 - 32

body/app.json

@@ -1,10 +1,10 @@
 {
-    "display_name": "Segmentation: Body CT",
-    "short_description": "<b>Description:</b><br>KonfAI adaptation of the <a href='https://github.com/wasserth/TotalSegmentator'>TotalSegmentator</a> body CT model for fast segmentation into <b>body trunk</b> and <b>extremities</b> at <b>1.5 mm isotropic</b> resolution.<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 a <b>KonfAI</b> adaptation of the <a href='https://github.com/wasserth/TotalSegmentator'>TotalSegmentator</a> body CT segmentation model.<br><br><b>Capabilities:</b><br>&bull; Segmentation of <b>2 body classes</b>: body_trunc and body_extremities<br>&bull; CT input resampled to <b>1.5 mm isotropic</b><br>&bull; Lightweight body mask generation for preprocessing and downstream workflows<br><br><b>Training data:</b><br>Trained on <b>1559 CT scans</b>, with labels for the body trunk and extremities<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>",
+    "display_name": "Segmentation: Body MRI/CT 6mm",
+    "short_description": "<b>Description:</b><br>Lightweight <b>KonfAI</b> adaptation of <a href='https://github.com/wasserth/TotalSegmentator'>TotalSegmentator MRI</a> for fast body segmentation into <b>body trunk</b> and <b>extremities</b> from <b>MRI and CT</b> scans at <b>6 mm isotropic</b> resolution.<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 body configuration of <a href='https://github.com/wasserth/TotalSegmentator'>TotalSegmentator MRI</a> into the <b>KonfAI</b> inference pipeline for efficient <b>MRI/CT</b> body mask generation.<br><br><b>Capabilities:</b><br>&bull; Segmentation of <b>2 body classes</b>: body_trunc and body_extremities<br>&bull; <b>6 mm isotropic</b> input for lower memory usage and fast preprocessing workflows<br>&bull; Robust body mask generation for both <b>MRI and CT</b> downstream pipelines<br><br><b>Training data:</b><br>Trained on a multimodal cohort with <b>MRI and CT</b> scans, with labels for the body trunk and extremities<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,
     "mc_dropout": 0,
-    "models": ["M299.pt"],
+    "models": ["M598.pt"],
     "inputs": {
       "Volume": {
         "display_name": "Input Volume",

body/model.py

@@ -21,30 +21,17 @@ class UNetBlock(network.ModuleArgsDict):
 
     def __init__(self, channels, i : int = 0) -> None:
         super().__init__()
-        self.add_module("DownConvBlock", ConvBlock(in_channels=channels[0], out_channels=channels[1], stride= 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:], 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 =  2, stride = 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 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):
+class Unet_TS_MR(network.Network):
 
     def __init__(self,
         optimizer: network.OptimizerLoader = network.OptimizerLoader(),
@@ -52,7 +39,7 @@ class Unet_TS_CT(network.Network):
             "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
         },
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
-        channels: list[int] = [1, 32, 64, 128, 320, 320]) -> None:
+        channels: list[int] = [1, 32, 64, 128, 320]) -> None:
         super().__init__(
             in_channels=channels[0],
             optimizer=optimizer,
@@ -61,9 +48,8 @@ class Unet_TS_CT(network.Network):
             patch=None,
             dim=3,
         )
-        self.add_module("ClipAndNormalize", ClipAndNormalize())
         self.add_module("UNetBlock", UNetBlock(channels))
-        self.add_module("Head", UNetHead(channels[1], 118))
+        self.add_module("Head", UNetHead(channels[1], 42))
 
     def load(
         self,
@@ -71,6 +57,6 @@ class Unet_TS_CT(network.Network):
         init: bool = True,
         ema: bool = False,
     ):
-        nb_class, in_channels = state_dict["Model"]["Unet_TS_CT"]["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)

body/requirements.txt

@@ -0,0 +1 @@
+segmentation_models_pytorch

build.py

@@ -16,10 +16,7 @@ def convert_torchScript_full(model_name: str, model: torch.nn.Module, task: str,
     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 :
@@ -105,6 +102,7 @@ models = {
          "M853" : (UnetCPP_3_MR(), False, 3, url+"v2.5.0-weights/Dataset853_TotalSegMRI_total_6mm_1088subj.zip"),
          "M299" : (UnetCPP_1_CT(), "total", 1, False, url+"v2.0.0-weights/Dataset299_body_1559subj.zip"),
          "M597" : (UnetCPP_1_MR(), "body", 1, True, url+"v2.5.0-weights/Dataset597_mri_body_139subj.zip"),
+         "M598" : (UnetCPP_3_MR(), "body", 3, True, url+"v2.5.0-weights/Dataset598_mri_body_6mm_139subj.zip"),
          }
 if __name__ == "__main__":
     for name, model in models.items():

total_mr-3mm/Prediction.yml

@@ -10,7 +10,6 @@ Predictor:
       - 128
       - 256
       - 320
-      mri: true
   Dataset:
     groups_src:
       Volume_0:
@@ -74,7 +73,7 @@ Predictor:
           TensorCast:
             dtype: uint8
             inverse: true
-        after_reduction_transforms: 
+        after_reduction_transforms:
           Sum:
             dim: 0
         final_transforms: None
@@ -82,6 +81,7 @@ Predictor:
         group: Output
         same_as_group: Volume_0:Volume
         patch_combine: Cosinus
+        inverse_transform: true
         reduction: Concat
   train_name: TotalSegmentator
   manual_seed: 32