nnUNet PICAI

[jewelltaylor]

PR Type

[Feature | Fix | Documentation | Other ]

Short Description

Clickup Ticket(s): Initial nnUNet Integration for PICAI research

This PR does some initial work to leverage the nnUNet framework in FL4Health. This includes:

• Moving preprocessing into the repo so that we have the flexibility to produce preprocessed datasets in the format that nnUNet expects.
• Run preprocessing on raw dataset on cluster to produce desired dataset on cluster. Remove old datasets.
• Update the existing FL and central examples to use the new dataset format.
• Add scripts to train nnUNet on the Whole Gland Segmentation task. One script to train from scratch on PICAI dataset. Another script to finetune a model already trained on a private dataset in the Haider Lab.
• Update scripts to requeue if they are still running when they timeout (previously a huge issue since training always takes longer than max job time)
• Update documentation in the picai research folder to include the new functionality

In following tickets I will:

• Do some initial experiments to establish some baselines.
• More sophisticated data augmentation pipeline, possibly borrowed from nnUNet, to apply to current central and FL examples.
• Integrate architecture and training pipeline configuration from nnUNet to the federated setting.

Tests Added

Since a lot of this is research code and subject to change, I have not included tests. However, let me know if you think there are any tests that would be useful.