Optimize data augmentation for 3D kernel

[jcohenadad]

Useful syntax

To look at predictions with GT. Run under the data_seg_mp2rage_20221220_124553/data_processed_lesionseg folder:

for i in 017 025 058 079 082 098 103 111 114 115; do fsleyes -S sub-P${i}/anat/sub-P${i}_UNIT1 derivatives/labels/sub-P${i}/anat/sub-P${i}_UNIT1_lesion-manualHaris -cm yellow ../../model_seg_lesion_mp2rage_xxx/pred_masks/sub-P${i}_UNIT1_pred -cm red -a 50; done

TODO:

• Use less conservative cropping during preprocessing, to avoid the 'black edge', which could limit the performance of training (ie: model would focus on edge detection).
• Use softseg
• Looks like only one lesion is segmented. Maybe check post-processing (keep only one object).

config file

{
    "command": "train",
    "gpu_ids": [4],
    "path_output": "/home/GRAMES.POLYMTL.CA/p101317/data_nvme_p101317/model_seg_lesion_mp2rage_",
    "model_name": "model_seg_lesion_mp2rage",
    "debugging": true,
    "object_detection_params": {
        "object_detection_path": null,
        "safety_factor": [1.0, 1.0, 1.0]
    },
    "wandb": {
        "wandb_api_key": "9095e2bc9e4ab445d478c9c8a81759ae908be8c6",
        "project_name": "basel-mp2rage-lesion",
        "group_name": "3D",
        "run_name": "run-1",
        "log_grads_every": 100
    },
    "loader_parameters": {
        "path_data": ["/home/GRAMES.POLYMTL.CA/p101317/data_nvme_p101317/data_seg_mp2rage_20221217_170634/data_processed_lesionseg"],
        "subject_selection": {"n": [], "metadata": [], "value": []},
        "target_suffix": ["_lesion-manualHaris"],
        "extensions": [".nii.gz"],
        "roi_params": {
            "suffix": null,
            "slice_filter_roi": null
        },
        "contrast_params": {
            "training_validation": ["UNIT1"],
            "testing": ["UNIT1"],
            "balance": {}
        },
        "slice_filter_params": {
            "filter_empty_mask": true,
            "filter_empty_input": true
        },
        "slice_axis": "axial",
        "multichannel": false,
        "soft_gt": false,
        "bids_validate": true
    },
    "split_dataset": {
        "fname_split": null,
        "random_seed": 42,
        "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": 16,
        "loss": {
            "name": "DiceLoss"
        },
        "training_time": {
            "num_epochs": 50,
            "early_stopping_patience": 50,
            "early_stopping_epsilon": 0.001
        },
        "scheduler": {
            "initial_lr": 0.001,
            "lr_scheduler": {
                "name": "CosineAnnealingLR",
                "base_lr": 1e-5,
                "max_lr": 1e-3
            }
        },
        "balance_samples": {
            "applied": false,
            "type": "gt"
        },
        "mixup_alpha": null,
        "transfer_learning": {
            "retrain_model": null,
            "retrain_fraction": 1.0,
            "reset": true
        }
    },
    "default_model": {
        "name": "Unet",
        "dropout_rate": 0.3,
        "bn_momentum": 0.1,
        "final_activation": "sigmoid",
        "depth": 3
    },
    "FiLMedUnet": {
        "applied": false,
        "metadata": "contrasts",
        "film_layers": [0, 1, 0, 0, 0, 0, 0, 0, 0, 0]
    },
    "Modified3DUNet": {
        "applied": true,
        "length_3D": [32, 32, 64],
        "stride_3D": [32, 32, 64],
        "attention": false,
        "n_filters": 8
    },
    "uncertainty": {
        "epistemic": false,
        "aleatoric": false,
        "n_it": 0
    },
    "postprocessing": {
        "remove_noise": {"thr": -1},
        "keep_largest": {},
        "binarize_prediction": {"thr": 0.5},
        "uncertainty": {"thr": -1, "suffix": "_unc-vox.nii.gz"},
        "fill_holes": {},
        "remove_small": {"unit": "vox", "thr": 3}
    },
    "evaluation_parameters": {
        "target_size": {"unit": "vox", "thr": [20, 100]},
        "overlap": {"unit": "vox", "thr": 3}
    },
    "transformation": {
        "Resample":
        {
            "hspace": 0.75,
            "wspace": 0.75,
            "dspace": 0.75
        },
        "CenterCrop": {
            "size": [64, 64, 128]},

        "RandomAffine": {
            "degrees": 5,
            "scale": [0.1, 0.1],
            "translate": [0.1, 0.1],
            "applied_to": ["im", "gt"],
            "dataset_type": ["training"]
        },
        "ElasticTransform": {
            "alpha_range": [28.0, 30.0],
            "sigma_range":  [3.5, 4.5],
            "p": 0.1,
            "applied_to": ["im", "gt"],
            "dataset_type": ["training"]
	},		    
	"NormalizeInstance": {"applied_to": ["im"]}
    }
}

[jcohenadad]

Terminal output

{
    "command": "train",
    "gpu_ids": [4],
    "path_output": "/home/GRAMES.POLYMTL.CA/p101317/data_nvme_p101317/model_seg_lesion_mp2rage_",
    "model_name": "model_seg_lesion_mp2rage",
    "debugging": true,
    "object_detection_params": {
        "object_detection_path": null,
        "safety_factor": [1.0, 1.0, 1.0]
    },
    "wandb": {
        "wandb_api_key": "9095e2bc9e4ab445d478c9c8a81759ae908be8c6",
        "project_name": "basel-mp2rage-lesion",
        "group_name": "3D",
        "run_name": "run-1",
        "log_grads_every": 100
    },
    "loader_parameters": {
        "path_data": ["/home/GRAMES.POLYMTL.CA/p101317/data_nvme_p101317/data_seg_mp2rage_20221217_170634/data_processed_lesionseg"],
        "subject_selection": {"n": [], "metadata": [], "value": []},
        "target_suffix": ["_lesion-manualHaris"],
        "extensions": [".nii.gz"],
        "roi_params": {
            "suffix": null,
            "slice_filter_roi": null
        },
        "contrast_params": {
            "training_validation": ["UNIT1"],
            "testing": ["UNIT1"],
            "balance": {}
        },
        "slice_filter_params": {
            "filter_empty_mask": true,
            "filter_empty_input": true
        },
        "slice_axis": "axial",
        "multichannel": false,
        "soft_gt": false,
        "bids_validate": true
    },
    "split_dataset": {
        "fname_split": null,
        "random_seed": 42,
        "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": 16,
        "loss": {
            "name": "DiceLoss"
        },
        "training_time": {
            "num_epochs": 50,
            "early_stopping_patience": 50,
            "early_stopping_epsilon": 0.001
        },
        "scheduler": {
            "initial_lr": 0.001,
            "lr_scheduler": {
                "name": "CosineAnnealingLR",
                "base_lr": 1e-5,
                "max_lr": 1e-3
            }
        },
        "balance_samples": {
            "applied": false,
            "type": "gt"
        },
        "mixup_alpha": null,
        "transfer_learning": {
            "retrain_model": null,
            "retrain_fraction": 1.0,
            "reset": true
        }
    },
    "default_model": {
        "name": "Unet",
        "dropout_rate": 0.3,
        "bn_momentum": 0.1,
        "final_activation": "sigmoid",
        "depth": 3
    },
    "FiLMedUnet": {
        "applied": false,
        "metadata": "contrasts",
        "film_layers": [0, 1, 0, 0, 0, 0, 0, 0, 0, 0]
    },
    "Modified3DUNet": {
        "applied": true,
        "length_3D": [32, 32, 64],
        "stride_3D": [32, 32, 64],
        "attention": false,
        "n_filters": 8
    },
    "uncertainty": {
        "epistemic": false,
        "aleatoric": false,
        "n_it": 0
    },
    "postprocessing": {
        "remove_noise": {"thr": -1},
        "keep_largest": {},
        "binarize_prediction": {"thr": 0.5},
        "uncertainty": {"thr": -1, "suffix": "_unc-vox.nii.gz"},
        "fill_holes": {},
        "remove_small": {"unit": "vox", "thr": 3}
    },
    "evaluation_parameters": {
        "target_size": {"unit": "vox", "thr": [20, 100]},
        "overlap": {"unit": "vox", "thr": 3}
    },
    "transformation": {
        "Resample":
        {
            "hspace": 0.75,
            "wspace": 0.75,
            "dspace": 0.75
        },
        "CenterCrop": {
            "size": [128, 128, 128]},
	"NormalizeInstance": {"applied_to": ["im"]}
    }
}

[jcohenadad]

Change:

        {
            "hspace": 1,
            "wspace": 1,
            "dspace": 1
        },
        "CenterCrop": {
            "size": [64, 64, 64]},

[jcohenadad]

Investigating the orientation and kernel size. With the following config:

        "slice_axis": "axial",
...
    "Modified3DUNet": {
        "applied": true,
        "length_3D": [
            32,
            32,
            16
        ],

Here is a patch:

imgplot = plt.imshow(img[4, 0, :, :, 8])

So, it seems like the orientation is consistent, ie: the superior-inferior direction is along the 3rd dim (the "8").

[jcohenadad]

Now, let's look whether a modified CenterCrop enables to reach the entire SI axis:

        "CenterCrop": {"size": [32, 32, 128]},

Provides the proper coverage. Below are one image corresponding to the center of each patch (16 vox in SI direction), using stride of 14 along the SI direction (hence, there are 9 batches to cover 128):

One thing to note, though, is that the centercrop along the AP direction might be too aggressive.

[jcohenadad]

This is how a CenterCrop below look like:

        "CenterCrop": {"size": [64, 32, 128]},

So we need to go with 32x64x128.

[jcohenadad]

OK well, 64x128x96 is not appropriate:

[jcohenadad]

that seems more appropriate, but I don't like the cropping in the AP direction with:

        "length_3D": [32, 32, 64],
        "stride_3D": [1, 2, 2],
        "CenterCrop": {"size": [32, 64, 128]},

[jcohenadad]

Changing for:

        "length_3D": [32, 64, 64],
        "stride_3D": [1, 1, 2],
        "CenterCrop": {"size": [32, 64, 128]},

Looks much better: