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.

[emersodb]

Overall, I think all the additions look awesome. Just some fairly straightforward comments. I didn't fully check the correctness of the preprocessing calculations, because that's pretty challenging without examples. So I'll assume you have already verified their "correctness"

[emersodb]

Just some very minor typo type comments. Good to go otherwise.