Intervertebral disc labeling using pose estimation

ivadomed/config/config_pose.json

@@ -0,0 +1,88 @@
+{
+    "command": "train",
+    "gpu_ids": [0],
+    "path_output": "/data/example/labeling_t0test1",
+    "model_name": "pose_model",
+    "debugging": false,
+    "loader_parameters": {
+        "path_data": ["/data/data-multi-subject"],
+        "target_suffix": ["_heatmap0"],
+        "extensions": [".nii.gz"],
+        "roi_params": {
+            "suffix": null,
+            "slice_filter_roi": null
+        },
+        "contrast_params": {
+            "training_validation": ["T1w"],
+            "testing": [ "T1w"],
+            "balance": {}
+        },
+        "slice_filter_params": {
+            "filter_empty_mask": true,
+            "filter_empty_input": true
+        },
+        "slice_axis": "sagittal",
+        "multichannel": false,
+	    "soft_gt": true
+    },
+    "split_dataset": {
+        "fname_split": null,
+        "random_seed": 8,
+        "split_method" : "participant_id",
+        "data_testing": {"data_type": null, "data_value":[]},
+        "balance": null,
+        "train_fraction": 0.6,
+        "test_fraction": 0.2
+    },
+    "training_parameters": {
+        "batch_size": 1,
+        "loss": {
+            "name": "JointsMSELoss"
+        },
+        "training_time": {
+            "num_epochs": 100,
+            "early_stopping_patience": 100,
+            "early_stopping_epsilon": 0.001
+        },
+        "scheduler": {
+            "initial_lr":0.0005,
+            "lr_scheduler": {
+                "name": "CosineAnnealingLR",
+                "base_lr": 1e-6,
+                "max_lr": 1e-4
+            }
+        },
+        "balance_samples": false,
+        "mixup_alpha": null,
+        "transfer_learning": null  
+    },
+    "default_model": {
+        "name": "HourglassNet"
+    },
+    "FiLMedUnet": {
+        "applied": false,
+        "metadata": "contrasts",
+        "film_layers": [0, 1, 0, 0, 0, 0, 0, 0, 0, 0]
+    },
+    "uncertainty": {
+        "epistemic": false,
+        "aleatoric": false,
+        "n_it": 0
+    },
+    "postprocessing": {
+        "remove_noise": {"thr": -1},
+        "binarize_prediction": {"thr": 0.5},
+        "uncertainty": {"thr": -1, "suffix": "_unc-vox.nii.gz"},
+        "remove_small": {"unit": "vox", "thr": 3}
+    },
+    "evaluation_parameters": {
+        "target_size": {"unit": "vox", "thr": [20, 100]},
+        "overlap": {"unit": "vox", "thr": 3}
+    },
+    "transformation": {
+	  "Resize": {"height":256, "width":256},
+      "VertebralSplitting": {"max_joint": 11, "applied_to": ["gt"]},
+      "NumpyToTensor": {},
+      "NormalizeInstance": {"applied_to": ["im"]}
+    }
+}

ivadomed/loader/loader.py

@@ -95,7 +95,6 @@ def load_dataset(bids_df, data_list, transforms_params, model_params, target_suf
     else:
         # Task selection
         task = imed_utils.get_task(model_params["name"])
-
         dataset = BidsDataset(bids_df=bids_df,
                               subject_file_lst=data_list,
                               target_suffix=target_suffix,
@@ -418,7 +417,7 @@ def get_pair_slice(self, slice_index, gt_type="segmentation"):
         else:
             gt_slices = []
             for gt_obj in gt_dataobj:
-                if gt_type == "segmentation":
+                if gt_type in ["segmentation", "pose_estimation"]:
                     if not isinstance(gt_obj, list):  # annotation from only one rater
                         gt_slices.append(np.asarray(gt_obj[..., slice_index],
                                                     dtype=np.float32))
@@ -633,6 +632,7 @@ def __init__(self, filename_pairs, length=[], stride=[], slice_axis=2, cache=Tru
 
     def load_filenames(self):
         """Load preprocessed pair data (input and gt) in handler."""
+
         for input_filenames, gt_filenames, roi_filename, metadata in self.filename_pairs:
             roi_pair = SegmentationPair(input_filenames, roi_filename, metadata=metadata, slice_axis=self.slice_axis,
                                         cache=self.cache, prepro_transforms=self.prepro_transforms)
@@ -642,7 +642,6 @@ def load_filenames(self):
                                         soft_gt=self.soft_gt)
 
             input_data_shape, _ = seg_pair.get_pair_shapes()
-
             for idx_pair_slice in range(input_data_shape[-1]):
                 slice_seg_pair = seg_pair.get_pair_slice(idx_pair_slice, gt_type=self.task)
                 self.has_bounding_box = imed_obj_detect.verify_metadata(slice_seg_pair, self.has_bounding_box)
@@ -757,14 +756,20 @@ def __getitem__(self, index):
         # Update metadata_input with metadata_roi
         metadata_gt = imed_loader_utils.update_metadata(metadata_input, metadata_gt)
 
-        if self.task == "segmentation":
+
+
+
+        if self.task in ["segmentation", "pose_estimation"]:
             # Run transforms on images
             stack_gt, metadata_gt = self.transform(sample=seg_pair_slice["gt"],
                                                    metadata=metadata_gt,
                                                    data_type="gt")
             # Make sure stack_gt is binarized
             if stack_gt is not None and not self.soft_gt:
                 stack_gt = imed_postpro.threshold_predictions(stack_gt, thr=0.5).astype(np.uint8)
+            if self.task == "pose_estimation":
+                stack_gt = torch.moveaxis(torch.squeeze(stack_gt), -1, 0)
+
 
         else:
             # Force no transformation on labels for classification task

ivadomed/losses.py

@@ -405,6 +405,47 @@ def forward(self, input, target):
 
         return mean_loss
 
+class JointsMSELoss(nn.Module):
+    """
+    Joint MSE loss for pose estimation method.
+    .. seealso::
+        Alejandro Newell et al. "Stacked Hourglass Networks for Human Pose Estimation."
+        Proceedings of the European Conference on Computer Vision. 2016.
+    Args:
+        output (tensor): prediction mask mask
+        target_and_weights (list): list as follows:
+          -- ground truth mask (Tensor): estimated mask by the pose model
+          -- target_weight (Tensor): visibility of the intervertebral disc to control the loss value for the training process
+
+    returns:
+            tensor: sum of losses computed on (mask, target) with the params
+    """    
+    def __init__(self, use_target_weight=True):
+        super(JointsMSELoss, self).__init__()
+        self.criterion = nn.MSELoss(reduction='mean')
+        self.use_target_weight = use_target_weight
+
+    def forward(self, output, target_and_weights):
+        target, target_weight = target_and_weights
+
+        batch_size = output.size(0)
+        num_joints = output.size(1)
+        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
+        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
+        loss = 0
+
+        for idx in range(num_joints):
+            heatmap_pred = heatmaps_pred[idx].squeeze()
+            heatmap_gt = heatmaps_gt[idx].squeeze()
+            if self.use_target_weight:
+                loss += 0.5 * self.criterion(
+                    heatmap_pred.mul(target_weight[:, idx]),
+                    heatmap_gt.mul(target_weight[:, idx])
+                )
+            else:
+                loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)
+
+        return loss / num_joints
 
 class LossCombination(nn.Module):
     """

ivadomed/main.py

@@ -24,7 +24,7 @@
 cudnn.benchmark = True
 
 # List of not-default available models i.e. different from Unet
-MODEL_LIST = ['Modified3DUNet', 'HeMISUnet', 'FiLMedUnet', 'resnet18', 'densenet121', 'Countception']
+MODEL_LIST = ['Modified3DUNet', 'HeMISUnet', 'FiLMedUnet', 'resnet18', 'densenet121', 'Countception', 'HourglassNet']
 
 
 def get_parser():
@@ -113,6 +113,9 @@ def film_normalize_data(context, model_params, ds_train, ds_valid, path_output):
     return model_params, ds_train, ds_valid, train_onehotencoder
 
 
+
+
+
 def get_dataset(bids_df, loader_params, data_lst, transform_params, cuda_available, device, ds_type):
     ds = imed_loader.load_dataset(bids_df, **{**loader_params, **{'data_list': data_lst,
                                                                   'transforms_params': transform_params,
@@ -535,7 +538,6 @@ def create_dataset_and_ivadomed_version_log(context):
 
 def run_main():
     imed_utils.init_ivadomed()
-
     parser = get_parser()
     args = parser.parse_args()
 
@@ -555,4 +557,6 @@ def run_main():
 
 
 if __name__ == "__main__":
+
+
     run_main()