Add body CT

body/Evaluation.yml

@@ -0,0 +1,29 @@
+Evaluator:
+  metrics:
+    Output:
+      targets_criterions:
+        Reference;Mask:
+          criterions_loader:
+            DiceSaveMap:
+              labels: None
+              dataset: ./Evaluations/TotalSegmentator/Output:mha
+              group: Seg_MAE_map
+  Dataset:
+    groups_src:
+      Mask_0:
+        groups_dest:
+          Mask:
+            transforms: None
+      Volume_0:
+        groups_dest:
+          Output:
+            transforms: None
+      Reference_0:
+        groups_dest:
+          Reference:
+            transforms: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:a:mha
+    validation: None
+  train_name: TotalSegmentator

body/M299.pt

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

body/Prediction.yml

@@ -0,0 +1,87 @@
+Predictor:
+  Model:
+    classpath: model:Unet_TS_CT
+    Unet_TS_CT:
+      outputs_criterions: None
+      channels:
+      - 1
+      - 32
+      - 64
+      - 128
+      - 256
+      - 320
+      - 320
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Volume:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+              Canonical:
+                inverse: true
+              ResampleToResolution:
+                spacing:
+                - 1.5
+                - 1.5
+                - 1.5
+                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/:mha
+    use_cache: false
+    batch_size: 1
+  outputs_dataset:
+    Head:Softmax:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: 
+          Softmax:
+            dim: 0
+          Argmax:
+            dim: 0
+          TensorCast:
+            dtype: uint8
+            inverse: true
+        after_reduction_transforms:
+          Sum:
+            dim: 0
+        final_transforms: None
+        dataset_filename: Output:mha
+        group: Output
+        same_as_group: Volume_0:Volume
+        patch_combine: Cosinus
+        inverse_transform: true
+        reduction: Concat
+  train_name: TotalSegmentator
+  manual_seed: 32
+  gpu_checkpoints: None
+  images_log: None
+  combine: Concat
+  autocast: true
+  data_log: None

body/app.json

@@ -0,0 +1,42 @@
+{
+    "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>",
+    "tta": 0,
+    "mc_dropout": 0,
+    "models": ["M299.pt"],
+    "inputs": {
+      "Volume": {
+        "display_name": "Input Volume",
+        "volume_type": "VOLUME",
+        "required": true
+      }
+    },
+    "outputs": {
+      "Segmentation": {
+        "display_name": "Segmentation",
+        "volume_type": "SEGMENTATION",
+        "required": true
+      }
+    },
+    "inputs_evaluations": {
+      "Image": {
+        "Evaluation.yml": {
+          "Segmentation": {
+            "display_name": "Output Segmentation",
+            "volume_type": "VOLUME",
+            "required": true
+          },
+          "GT_Segmentation": {
+            "display_name": "GT Segmentation",
+            "volume_type": "VOLUME",
+            "required": true
+          }
+        }
+      }
+    },
+    "terminology": {
+      "1": { "name": "body_trunc", "color": "#4B79EA" },
+      "2": { "name": "body_extremities", "color": "#A8EA4B" }
+    }
+}

body/model.py

@@ -0,0 +1,76 @@
+import torch
+from konfai.network import network, blocks
+
+class ConvBlock(network.ModuleArgsDict):
+    def __init__(self, in_channels : int, out_channels : int, stride: int = 1 ) -> None:
+        super().__init__()
+        self.add_module("Conv_0", torch.nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=1, bias=True))
+        self.add_module("Norm_0", torch.nn.InstanceNorm3d(num_features=out_channels, affine=True))
+        self.add_module("Activation_0", torch.nn.LeakyReLU(negative_slope=0.01, inplace=True))
+        self.add_module("Conv_1", torch.nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=True))
+        self.add_module("Norm_1", torch.nn.InstanceNorm3d(num_features=out_channels, affine=True))
+        self.add_module("Activation_1", torch.nn.LeakyReLU(negative_slope=0.01, inplace=True))
+        
+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, 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))
+     
+        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("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):
+
+    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()},
+        channels: list[int] = [1, 32, 64, 128, 320, 320]) -> None:
+        super().__init__(
+            in_channels=channels[0],
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            patch=None,
+            dim=3,
+        )
+        self.add_module("ClipAndNormalize", ClipAndNormalize())
+        self.add_module("UNetBlock", UNetBlock(channels))
+        self.add_module("Head", UNetHead(channels[1], 118))
+
+    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_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)

build.py

@@ -94,15 +94,17 @@ models = {
          "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"),
+         "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")
+         "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"),
          }
 if __name__ == "__main__":
     for name, model in models.items():