Pre-training on large-scale T2w healthy/pathology data

configs/train_supervised.yaml

@@ -0,0 +1,50 @@
+seed: 30
+save_test_preds: True
+
+preprocessing:
+  # Center crop/pad images to the specified size. (NOTE: done after resampling to 1.0x1.0x1.0)
+  # values correspond to R-L, A-P, I-S axes of the image after 1mm isotropic resampling.
+  # NOTE for swinunetr: spatial dimensions of input image must be divisible by 32, i.e., we cannot you, for example, 48x48x48
+  crop_pad_size: [64, 64, -1]
+  patch_size: [64, 64, 64]
+
+opt:
+  name: adam
+  lr: 0.001
+  max_epochs: 200
+  warmup_epochs: 20
+  batch_size: 2
+  # Interval between validation checks in epochs
+  check_val_every_n_epochs: 5
+  # Early stopping patience (this is until patience * check_val_every_n_epochs)
+  early_stopping_patience: 20     
+
+
+model:
+  # Model architecture to be used for training (also to be specified as args in the command line)
+  nnunet:
+    in_channels: 1
+    out_channels: 1
+    # NOTE: these info are typically taken from nnUNetPlans.json (if an nnUNet model is trained)
+    base_num_features: 32
+    max_num_features: 320
+    n_conv_per_stage_encoder: [2, 2, 2, 2, 2, 2]
+    n_conv_per_stage_decoder: [2, 2, 2, 2, 2]
+    pool_op_kernel_sizes: [
+      [1, 1, 1],
+      [2, 2, 2],
+      [2, 2, 2],
+      [2, 2, 2],
+      [2, 2, 2],
+      [1, 2, 2] 
+    ]
+    enable_deep_supervision: True
+
+  swinunetr:
+    in_channels: 1
+    out_channels: 1
+    spatial_dims: 3
+    depths: [2, 2, 2, 2]
+    num_heads: [3, 6, 12, 24]   # number of heads in multi-head Attention
+    feature_size: 36
+    use_pretrained: False

dataset-conversion/README.md

@@ -1 +1,84 @@
-This folder contains various scripts used for dataset conversion from BIDS to nnUNet/MONAI formats.
+# Pre-training on multiple datasets
+
+### Download datasets
+
+Download T2w images and spinal cord segmentations for the following datasets.
+
+```commandline
+cd ~/duke/temp/janvalosek/ssl_pretraining_multiple_datasets
+```
+
+`spine-generic multi-subject` (n=267)
+
+```commandline
+git clone https://github.com/spine-generic/data-multi-subject
+cd data-multi-subject
+git checkout sb/156-add-preprocessed-images
+git annex get $(find . -name "*space-other_T2w.nii.gz")
+git annex get $(find . -name "*space-other_T2w_label-SC_seg.nii.gz")
+```
+
+`whole-spine` (n=45)
+
+NOTE: we do not use this dataset right now due to https://github.com/neuropoly/data-management/issues/306
+
+```commandline
+git clone [email protected]:datasets/whole-spine
+cd whole-spine
+git annex dead here
+git annex get $(find . -name "*T2w.nii.gz")
+git annex get $(find . -name "*T2w_seg.nii.gz")
+```
+
+`canproco` (n=321) 
+
+NOTE: we do not use sub-cal* subjects right now due to https://github.com/neuropoly/data-management/issues/305
+
+```commandline
+git clone [email protected]:datasets/canproco
+cd canproco
+git annex dead here
+git annex get $(find . -name "*ses-M0_T2w.nii.gz")
+git annex get $(find . -name "*ses-M0_T2w_seg-manual.nii.gz")
+```
+
+`sci-colorado` (n=80)
+
+```commandline
+git clone [email protected]:datasets/sci-colorado
+cd sci-colorado
+git annex dead here
+git annex get $(find . -name "*T2w.nii.gz")
+git annex get $(find . -name "*T2w_seg-manual.nii.gz")
+```
+
+`dcm-zurich` (n=135)
+
+```commandline
+git clone [email protected]:datasets/dcm-zurich
+cd dcm-zurich
+git annex dead here
+git annex get $(find . -name "*acq-axial_T2w.nii.gz")
+git annex get $(find . -name "*acq-axial_T2w_label-SC_mask-manual.nii.gz")
+```
+
+`sci-paris` (n=14)
+
+```commandline
+git clone [email protected]:datasets/sci-paris
+cd sci-paris
+git annex dead here
+git annex get $(find . -name "*T2w.nii.gz")
+git annex get $(find . -name "*T2w_seg.nii.gz")
+```
+
+### Create MSD-style JSON datalists
+
+```commandline
+python /Users/user/code/model-seg-dcm/vit_unetr_ssl/create_msd_data.py --path-data data-multi-subject --dataset-name spine-generic --path-out . --split 0.8 0.2 --seed 42
+python /Users/user/code/model-seg-dcm/vit_unetr_ssl/create_msd_data.py --path-data whole-spine --dataset-name whole-spine --path-out . --split 0.8 0.2 --seed 42
+python /Users/user/code/model-seg-dcm/vit_unetr_ssl/create_msd_data.py --path-data canproco --dataset-name canproco --path-out . --split 0.8 0.2 --seed 42
+python /Users/user/code/model-seg-dcm/vit_unetr_ssl/create_msd_data.py --path-data sci-colorado --dataset-name sci-colorado --path-out . --split 0.8 0.2 --seed 42
+python /Users/user/code/model-seg-dcm/vit_unetr_ssl/create_msd_data.py --path-data dcm-zurich --dataset-name dcm-zurich --path-out . --split 0.8 0.2 --seed 42
+python /Users/user/code/model-seg-dcm/vit_unetr_ssl/create_msd_data.py --path-data sci-paris --dataset-name sci-paris --path-out . --split 0.8 0.2 --seed 42
+```

dataset-conversion/create_msd_data.py

@@ -0,0 +1,184 @@
+"""
+Create MSD-style JSON datalist file for BIDS datasets.
+The following two keys are included in the JSON file: 'image' and 'label_sc'.
+
+NOTE: the script is meant to be used for SUPERVISED pre-training, meaning that the dataset is split into training and
+validation.
+In other words, NO testing set is created.
+
+The script has to be run for each dataset separately, meaning that one JSON file is created for each dataset.
+
+Example usage:
+    python create_msd_data.py
+        --path-data /Users/user/data/spine-generic
+        --dataset-name spine-generic
+        --path-out /Users/user/data/spine-generic
+
+    python create_msd_data.py
+        --path-data /Users/user/data/dcm-zurich
+        --dataset-name dcm-zurich
+        --path-out /Users/user/data/dcm-zurich
+"""
+
+import os
+import json
+import argparse
+from pathlib import Path
+from loguru import logger
+from sklearn.model_selection import train_test_split
+
+contrast_dict = {
+    'spine-generic': 'space-other_T2w',     # iso T2w (preprocessed data)
+    'whole-spine': 'T2w',                   # iso T2w
+    'canproco': 'ses-M0_T2w',               # iso T2w (session M0)
+    'dcm-zurich': 'acq-axial_T2w',          # axial T2w
+    'sci-paris': 'T2w',                     # iso T2w
+    'sci-colorado': 'T2w'                   # axial T2w
+}
+
+# Spinal cord segmentation file suffixes for different datasets
+sc_fname_suffix_dict = {
+    'spine-generic': 'label-SC_seg',
+    'whole-spine': 'seg',
+    'canproco': 'seg-manual',
+    'dcm-zurich': 'label-SC_mask-manual',
+    'sci-paris': 'seg-manual',
+    'sci-colorado': 'seg-manual'
+}
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description='Create MSD-style JSON datalist file for BIDS datasets.')
+
+    parser.add_argument('--path-data', required=True, type=str,
+                        help='Path to BIDS dataset. Example: /Users/user/data/dcm-zurich')
+    parser.add_argument('--dataset-name', required=True, type=str,
+                        help='Name of the dataset. Example: spine-generic or dcm-zurich.')
+    parser.add_argument('--path-out', type=str, required=True,
+                        help='Path to the output directory where dataset json is saved')
+    parser.add_argument('--split', nargs='+', type=float, default=[0.8, 0.2],
+                        help='Ratios of training and validation 0-1. '
+                             'Example: --split 0.8 0.2')
+    parser.add_argument('--seed', default=42, type=int, help="Seed for reproducibility")
+
+    return parser
+
+
+def main():
+    args = get_parser().parse_args()
+
+    dataset = os.path.abspath(args.path_data)
+    dataset_name = args.dataset_name
+    train_ratio, val_ratio = args.split
+    seed = args.seed
+    path_out = os.path.abspath(args.path_out)
+
+    # Check if the dataset name is valid
+    if dataset_name not in contrast_dict.keys():
+        raise ValueError(f"Dataset name {dataset_name} is not valid. Choose from {contrast_dict.keys()}")
+
+    contrast = contrast_dict[dataset_name]
+    sc_fname_suffix = sc_fname_suffix_dict[dataset_name]
+    datalist_fname = f"{dataset_name}_seed{seed}"
+
+    train_images, val_images = {}, {}
+
+    # For spine-generic, we add 'derivatives/data_preprocessed' to the path to use the preprocessed data with the same
+    # resolution and orientation as the spinal cord segmentations
+    if dataset_name == 'spine-generic':
+        root = Path(dataset) / 'derivatives/data_preprocessed'
+    else:
+        root = Path(dataset)
+    # Path to 'derivatives/labels with spinal cord segmentations
+    labels = Path(dataset) / 'derivatives/labels'
+
+    # Check if the dataset path exists
+    if not os.path.exists(root):
+        raise ValueError(f"Path {root} does not exist.")
+    if not os.path.exists(labels):
+        raise ValueError(f"Path {labels} does not exist.")
+
+    logger.info(f"Root path: {root}")
+    logger.info(f"Labels path: {labels}")
+
+    # get recursively all the subjects from the root folder
+    subjects = [sub for sub in os.listdir(root) if sub.startswith("sub-")]
+
+    # Get the training and validation splits
+    # Note: we are doing SSL pre-training, so we don't need test set
+    tr_subs, val_subs = train_test_split(subjects, test_size=val_ratio, random_state=args.seed)
+
+    # recursively find the spinal cord segmentation files under 'derivatives/labels' for training and validation
+    # subjects
+    tr_seg_files = [str(path) for sub in tr_subs for path in
+                    Path(labels).rglob(f"{sub}_{contrast}_{sc_fname_suffix}.nii.gz")]
+    val_seg_files = [str(path) for sub in val_subs for path in
+                     Path(labels).rglob(f"{sub}_{contrast}_{sc_fname_suffix}.nii.gz")]
+
+    # update the train and validation images dicts with the key as the subject and value as the path to the subject
+    train_images.update({sub: os.path.join(root, sub) for sub in tr_seg_files})
+    val_images.update({sub: os.path.join(root, sub) for sub in val_seg_files})
+
+    logger.info(f"Found subjects in the training set: {len(train_images)}")
+    logger.info(f"Found subjects in the validation set: {len(val_images)}")
+
+    # keys to be defined in the dataset_0.json
+    params = {}
+    params["dataset_name"] = dataset_name
+    params["contrast"] = contrast
+    params["labels"] = {
+        "0": "background",
+        "1": "sc-seg"
+    }
+    params["modality"] = {
+        "0": "MRI"
+    }
+    params["numTraining"] = len(train_images)
+    params["numValidation"] = len(val_images)
+    params["seed"] = args.seed
+    params["tensorImageSize"] = "3D"
+
+    train_images_dict = {"training": train_images}
+    val_images_dict = {"validation": val_images}
+
+    all_images_list = [train_images_dict, val_images_dict]
+
+    for images_dict in all_images_list:
+
+        for name, images_list in images_dict.items():
+
+            temp_list = []
+            for label in images_list:
+
+                temp_data_t2w = {}
+                # create the image path by replacing the label path
+                if dataset_name == 'spine-generic':
+                    temp_data_t2w["image"] = label.replace(f'_{sc_fname_suffix}', '').replace('labels',
+                                                                                              'data_preprocessed')
+                else:
+                    temp_data_t2w["image"] = label.replace(f'_{sc_fname_suffix}', '').replace('/derivatives/labels', '')
+
+                # Spinal cord segmentation file
+                temp_data_t2w["label_sc"] = label
+
+                if os.path.exists(temp_data_t2w["label_sc"]) and os.path.exists(temp_data_t2w["image"]):
+                    temp_list.append(temp_data_t2w)
+                else:
+                    logger.info(f"Either image/label does not exist.")
+
+            params[name] = temp_list
+            logger.info(f"Number of images in {name} set: {len(temp_list)}")
+
+    final_json = json.dumps(params, indent=4, sort_keys=False)
+    if not os.path.exists(path_out):
+        os.makedirs(path_out, exist_ok=True)
+
+    jsonFile = open(path_out + "/" + f"{datalist_fname}.json", "w")
+    jsonFile.write(final_json)
+    jsonFile.close()
+    print(f"JSON file saved to {path_out}/{datalist_fname}.json")
+
+
+if __name__ == "__main__":
+    main()
+

pretraining_and_finetuning/loader.py

@@ -0,0 +1,50 @@
+from monai.data import (DataLoader, DistributedSampler, CacheDataset, load_decathlon_datalist)
+
+from transforms import train_transforms, val_transforms
+
+
+def load_data(datalists_paths, train_batch_size, val_batch_size, num_workers=8, use_distributed=False,
+              crop_size=(64, 192, 320), patch_size=(64, 64, 64), device="cuda", task="pretraining"):
+    """
+    Return train and val dataloaders from datalist json file
+    :param datalists_paths: path(s) to the datalist json file(s)
+    :param train_batch_size: batch size for training dataloader
+    :param val_batch_size: batch size for validation dataloader
+    :param num_workers: number of workers for dataloader
+    :param use_distributed: whether to use distributed training
+    :param crop_size: crop size; e.g., (64, 192, 320)
+    :param patch_size: patch size; e.g., (64, 64, 64)
+    :param device: device to load data and apply transforms
+    :param task: task for train/val transforms; choices: pretraining or finetuning
+    """
+    train_datalist = []
+    val_datalist = []
+    for datalist_path in datalists_paths:
+        train_datalist += load_decathlon_datalist(data_list_file_path=datalist_path, data_list_key="training")
+        val_datalist += load_decathlon_datalist(data_list_file_path=datalist_path, data_list_key="validation")
+
+    train_tfs = train_transforms(crop_size, patch_size, device=device, task=task)
+    val_tfs = val_transforms(crop_size, task=task)
+
+    # training dataset
+    train_ds = CacheDataset(data=train_datalist, transform=train_tfs, cache_rate=0.5, num_workers=4, 
+                            copy_cache=False)
+    # validation dataset
+    val_ds = CacheDataset(data=val_datalist, transform=val_tfs, cache_rate=0.25, num_workers=4,
+                          copy_cache=False)
+
+    if use_distributed:
+        train_sampler = DistributedSampler(dataset=train_ds, even_divisible=True, shuffle=True)
+        val_sampler = DistributedSampler(dataset=val_ds, even_divisible=True, shuffle=False)
+    else:
+        train_sampler = None
+        val_sampler = None
+
+    # training dataloader    
+    train_loader = DataLoader(train_ds, batch_size=train_batch_size, shuffle=True, num_workers=num_workers, 
+                            pin_memory=True, sampler=train_sampler, persistent_workers=True)
+    # validation dataloader
+    val_loader = DataLoader(val_ds, batch_size=val_batch_size, shuffle=False, num_workers=num_workers,
+                            pin_memory=True, sampler=val_sampler, persistent_workers=True)
+
+    return train_loader, val_loader