Integrate uncertainty handling and embed uncertainty code directly into KonfAI as a builtin component Konfai==1.3.6

CBCT/Evaluation.yml

@@ -0,0 +1,58 @@
+Evaluator:
+  metrics:
+    sCT:
+      targets_criterions:
+        reference:
+          criterions_loader:
+            MAE:
+              reduction: mean
+            PSNR:
+              dynamic_range: None
+            SSIM:
+              dynamic_range: None
+    sCT_seg:
+      targets_criterions:
+        reference_seg:
+          criterions_loader:
+            Dice:
+              labels: None
+  Dataset:
+    groups_src:
+      sCT:
+        groups_dest:
+          sCT:
+            transforms: None
+          sCT_seg:
+            transforms:
+              KonfAIInference:
+                repo_id: VBoussot/MRSegmentator-KonfAI
+                model_name: MRSegmentator
+                number_of_ensemble: 1
+                number_of_tta: 0
+                number_of_mc_dropout: 0
+                per_channel: false
+              Save:
+                dataset: ./Dataset:nii.gz
+                group: None
+      Volume:
+        groups_dest:
+          reference:
+            transforms: None
+          reference_seg:
+            transforms:
+              KonfAIInference:
+                repo_id: VBoussot/MRSegmentator-KonfAI
+                model_name: MRSegmentator
+                number_of_ensemble: 1
+                number_of_tta: 0
+                number_of_mc_dropout: 0
+                per_channel: false
+              Save:
+                dataset: ./Dataset:nii.gz
+                group: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:a:nii.gz
+    - ./Predictions/ImpactSynth/Output:i:nii.gz
+    validation: None
+  train_name: ImpactSynth

CBCT/Model.py

@@ -28,41 +28,4 @@ class UNetpp(network.Network):
             classes=1,  
             activation=None  
         ))
-        self.add_module("Head", Head())
-
-class Concat(Reduction):
-
-    def __init__(self):
-        pass
-
-    def __call__(self, tensor: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
-        if isinstance(tensor, list):
-            return torch.stack(tensor, dim=2).squeeze(1)
-        else:
-            return  tensor.view(tensor.shape[0]*tensor.shape[1], -1, *tensor.shape[3:])
-        
-class Uncertainty(Transform):
-
-    def __init__(self):
-        pass
-
-    def __call__(self, name: str, tensors: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        dataset = Dataset("./Predictions/ImpactSynth/Dataset", "mha")
-        for i, tensor in enumerate(tensors):
-            dataset.write(f"sCT_{i}", name, data_to_image(tensor.unsqueeze(0).numpy(), cache_attribute))
-        data_var =  tensors.var(0) if tensors.shape[0] > 1 else torch.zeros_like(tensors[0])
-        dataset.write(f"sCT_var", name, data_to_image(data_var.unsqueeze(0).numpy(), cache_attribute))
-        return tensors.mean(0).unsqueeze(0)
-    
-class UnNormalize(Transform):
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.v_min = -1024
-        self.v_max = 3071
-
-    def __call__(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        return (input + 1)/2*(self.v_max-self.v_min) + self.v_min
-    
-    def inverse(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        pass
+        self.add_module("Head", Head())

CBCT/Prediction.yml

@@ -30,7 +30,7 @@ Predictor:
                 max_value: 1
                 inverse: true
             is_input: true
-    augmentations: 
+    augmentations:
       DataAugmentation_0:
         data_augmentations:
           Flip:
@@ -62,23 +62,28 @@ Predictor:
         before_reduction_transforms: None
         after_reduction_transforms: None
         final_transforms:
-          Model:UnNormalize: {}
-          Model:Uncertainty: {}
+          UnNormalize:
+            min_value: -1024
+            max_value: 3071
           TensorCast:
             dtype: int16
-            inverse: true
-        dataset_filename: Dataset:mha
+            inverse: false
+          InferenceStack:
+            dataset: Uncertainty:nii.gz
+            name: InferenceStack
+            mode: mean
+        dataset_filename: Output:nii.gz
         group: sCT
         same_as_group: Volume:Volume
         patch_combine: None
         inverse_transform: false
-        reduction: Model:Concat
-        Model:Concat: {}
+        reduction: Concat
+        Concat: {}
   train_name: ImpactSynth
   manual_seed: 32
   gpu_checkpoints: None
   images_log: None
-  combine: Model:Concat
+  combine: Concat
   autocast: false
   data_log: None
-  Model:Concat: {}
+  Concat: {}

CBCT/Uncertainty.yml

@@ -0,0 +1,43 @@
+Evaluator:
+  metrics:
+    Uncertainty:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean: {}
+    Comformity:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean: {}
+  Dataset:
+    groups_src:
+      InferenceStack:
+        groups_dest:
+          Uncertainty:
+            transforms:
+              Variance: {}
+              Save:
+                dataset: ./Uncertainty:nii.gz
+                group: None
+          Comformity:
+            transforms:
+              KonfAIInference:
+                repo_id: VBoussot/MRSegmentator-KonfAI
+                model_name: MRSegmentator
+                number_of_ensemble: 1
+                number_of_tta: 0
+                number_of_mc_dropout: 0
+                per_channel: true
+              Save/1:
+                dataset: ./Uncertainty:nii.gz
+                group: None
+              Variance: {}
+              Save/2:
+                dataset: ./Uncertainty:nii.gz
+                group: Comformity_var
+    subset: None
+    dataset_filenames:
+    - ./Uncertainty:nii.gz
+    validation: None
+  train_name: ImpactSynth

CBCT/requirements.txt

@@ -1 +1,2 @@
-segmentation_models_pytorch
+segmentation_models_pytorch
+konfai==1.3.6

MR/Evaluation.yml

@@ -0,0 +1,58 @@
+Evaluator:
+  metrics:
+    sCT:
+      targets_criterions:
+        reference:
+          criterions_loader:
+            MAE:
+              reduction: mean
+            PSNR:
+              dynamic_range: None
+            SSIM:
+              dynamic_range: None
+    sCT_seg:
+      targets_criterions:
+        reference_seg:
+          criterions_loader:
+            Dice:
+              labels: None
+  Dataset:
+    groups_src:
+      sCT:
+        groups_dest:
+          sCT:
+            transforms: None
+          sCT_seg:
+            transforms:
+              KonfAIInference:
+                repo_id: VBoussot/MRSegmentator-KonfAI
+                model_name: MRSegmentator
+                number_of_ensemble: 1
+                number_of_tta: 0
+                number_of_mc_dropout: 0
+                per_channel: false
+              Save:
+                dataset: ./Dataset:nii.gz
+                group: None
+      Volume:
+        groups_dest:
+          reference:
+            transforms: None
+          reference_seg:
+            transforms:
+              KonfAIInference:
+                repo_id: VBoussot/MRSegmentator-KonfAI
+                model_name: MRSegmentator
+                number_of_ensemble: 1
+                number_of_tta: 0
+                number_of_mc_dropout: 0
+                per_channel: false
+              Save:
+                dataset: ./Dataset:nii.gz
+                group: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:a:nii.gz
+    - ./Predictions/ImpactSynth/Output:i:nii.gz
+    validation: None
+  train_name: ImpactSynth

MR/Model.py

@@ -28,43 +28,4 @@ class UNetpp(network.Network):
             classes=1,  
             activation=None  
         ))
-        self.add_module("Head", Head())
-
-class Concat(Reduction):
-
-    def __init__(self):
-        pass
-
-    def __call__(self, tensor: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
-        if isinstance(tensor, list):
-            return torch.stack(tensor, dim=2).squeeze(1)
-        else:
-            return  tensor.view(tensor.shape[0]*tensor.shape[1], -1, *tensor.shape[3:])
-        
-        
-class Uncertainty(Transform):
-
-    def __init__(self):
-        pass
-
-    def __call__(self, name: str, tensors: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        dataset = Dataset("./Predictions/ImpactSynth/Dataset", "mha")
-        print("end", tensors.shape)
-        for i, tensor in enumerate(tensors):
-            dataset.write(f"sCT_{i}", name, data_to_image(tensor.unsqueeze(0).numpy(), cache_attribute))
-        data_var =  tensors.var(0) if tensors.shape[0] > 1 else torch.zeros_like(tensors[0])
-        dataset.write(f"sCT_var", name, data_to_image(data_var.unsqueeze(0).numpy(), cache_attribute))
-        return tensors.mean(0).unsqueeze(0)
-    
-class UnNormalize(Transform):
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.v_min = -1024
-        self.v_max = 3071
-
-    def __call__(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        return (input + 1)/2*(self.v_max-self.v_min) + self.v_min
-    
-    def inverse(self, name: str, input : torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        pass
+        self.add_module("Head", Head())

MR/Prediction.yml

@@ -10,13 +10,25 @@ Predictor:
         groups_dest:
           Volume:
             transforms:
-              Standardize:
-                lazy: false
-                mean: None
-                std: None
+              Clip:
+                min_value: min
+                max_value: percentile:99.5
+                save_clip_min: false
+                save_clip_max: false
                 mask: None
+              Normalize:
+                lazy: true
+                channels: None
+                min_value: -1
+                max_value: 1
+                inverse: true
+            patch_transforms:
+              Normalize:
+                lazy: false
+                channels: None
+                min_value: -1
+                max_value: 1
                 inverse: true
-            patch_transforms: None
             is_input: true
     augmentations:
       DataAugmentation_0:
@@ -50,23 +62,28 @@ Predictor:
         before_reduction_transforms: None
         after_reduction_transforms: None
         final_transforms:
-          Model:UnNormalize: {}
-          Model:Uncertainty: {}
+          UnNormalize:
+            min_value: -1024
+            max_value: 3071
           TensorCast:
             dtype: int16
-            inverse: true
-        dataset_filename: Dataset:mha
+            inverse: false
+          InferenceStack:
+            dataset: Uncertainty:nii.gz
+            name: InferenceStack
+            mode: mean
+        dataset_filename: Output:nii.gz
         group: sCT
         same_as_group: Volume:Volume
         patch_combine: None
         inverse_transform: false
-        reduction: Model:Concat
-        Model:Concat: {}
+        reduction: Concat
+        Concat: {}
   train_name: ImpactSynth
   manual_seed: 32
   gpu_checkpoints: None
   images_log: None
-  combine: Model:Concat
+  combine: Concat
   autocast: false
   data_log: None
-  Model:Concat: {}
+  Concat: {}

MR/Uncertainty.yml

@@ -0,0 +1,43 @@
+Evaluator:
+  metrics:
+    Uncertainty:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean: {}
+    Comformity:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean: {}
+  Dataset:
+    groups_src:
+      InferenceStack:
+        groups_dest:
+          Uncertainty:
+            transforms:
+              Variance: {}
+              Save:
+                dataset: ./Uncertainty:nii.gz
+                group: None
+          Comformity:
+            transforms:
+              KonfAIInference:
+                repo_id: VBoussot/MRSegmentator-KonfAI
+                model_name: MRSegmentator
+                number_of_ensemble: 1
+                number_of_tta: 0
+                number_of_mc_dropout: 0
+                per_channel: true
+              Save/1:
+                dataset: ./Uncertainty:nii.gz
+                group: None
+              Variance: {}
+              Save/2:
+                dataset: ./Uncertainty:nii.gz
+                group: Comformity_var
+    subset: None
+    dataset_filenames:
+    - ./Uncertainty:nii.gz
+    validation: None
+  train_name: ImpactSynth

MR/requirements.txt

@@ -1 +1,2 @@
-segmentation_models_pytorch
+segmentation_models_pytorch
+konfai==1.3.6