├── .gitignore ├── LICENSE ├── README.md ├── documentation ├── __init__.py ├── benchmarking.md ├── changelog.md ├── convert_msd_dataset.md ├── dataset_format.md ├── dataset_format_inference.md ├── explanation_normalization.md ├── explanation_plans_files.md ├── extending_nnunet.md ├── how_to_use_nnunet.md ├── installation_instructions.md ├── manual_data_splits.md ├── pretraining_and_finetuning.md ├── region_based_training.md ├── run_inference_with_pretrained_models.md ├── set_environment_variables.md ├── setting_up_paths.md └── tldr_migration_guide_from_v1.md ├── nnunetv2 ├── __init__.py ├── batch_running │ ├── __init__.py │ ├── benchmarking │ │ ├── __init__.py │ │ ├── generate_benchmarking_commands.py │ │ └── summarize_benchmark_results.py │ ├── collect_results_custom_Decathlon.py │ ├── collect_results_custom_Decathlon_2d.py │ ├── generate_lsf_runs_customDecathlon.py │ └── release_trainings │ │ ├── __init__.py │ │ └── nnunetv2_v1 │ │ ├── __init__.py │ │ ├── collect_results.py │ │ └── generate_lsf_commands.py ├── configuration.py ├── dataset_conversion │ ├── Dataset073_Fluo_C3DH_A549_SIM.py │ ├── Dataset120_RoadSegmentation.py │ ├── Dataset137_BraTS21.py │ ├── Dataset218_Amos2022_task1.py │ ├── Dataset219_Amos2022_task2.py │ ├── __init__.py │ ├── convert_MSD_dataset.py │ ├── convert_raw_dataset_from_old_nnunet_format.py │ ├── datasets_for_integration_tests │ │ ├── Dataset996_IntegrationTest_Hippocampus_regions_ignore.py │ │ ├── Dataset997_IntegrationTest_Hippocampus_regions.py │ │ ├── Dataset998_IntegrationTest_Hippocampus_ignore.py │ │ ├── Dataset999_IntegrationTest_Hippocampus.py │ │ └── __init__.py │ └── generate_dataset_json.py ├── ensembling │ ├── __init__.py │ └── ensemble.py ├── evaluation │ ├── __init__.py │ ├── accumulate_cv_results.py │ ├── evaluate_predictions.py │ └── find_best_configuration.py ├── experiment_planning │ ├── __init__.py │ ├── dataset_fingerprint │ │ ├── __init__.py │ │ └── fingerprint_extractor.py │ ├── experiment_planners │ │ ├── __init__.py │ │ ├── default_experiment_planner.py │ │ ├── network_topology.py │ │ └── resencUNet_planner.py │ ├── plan_and_preprocess_api.py │ ├── plan_and_preprocess_entrypoints.py │ ├── plans_for_pretraining │ │ ├── __init__.py │ │ └── move_plans_between_datasets.py │ └── verify_dataset_integrity.py ├── imageio │ ├── __init__.py │ ├── base_reader_writer.py │ ├── natural_image_reager_writer.py │ ├── nibabel_reader_writer.py │ ├── reader_writer_registry.py │ ├── simpleitk_reader_writer.py │ └── tif_reader_writer.py ├── inference │ ├── __init__.py │ ├── export_prediction.py │ ├── predict_from_raw_data.py │ └── sliding_window_prediction.py ├── model_sharing │ ├── __init__.py │ ├── entry_points.py │ ├── model_download.py │ ├── model_export.py │ └── model_import.py ├── paths.py ├── postprocessing │ ├── __init__.py │ └── remove_connected_components.py ├── preprocessing │ ├── __init__.py │ ├── cropping │ │ ├── __init__.py │ │ └── cropping.py │ ├── normalization │ │ ├── __init__.py │ │ ├── default_normalization_schemes.py │ │ └── map_channel_name_to_normalization.py │ ├── preprocessors │ │ ├── __init__.py │ │ └── default_preprocessor.py │ └── resampling │ │ ├── __init__.py │ │ ├── default_resampling.py │ │ └── utils.py ├── run │ ├── __init__.py │ ├── load_pretrained_weights.py │ └── run_training.py ├── tests │ ├── __init__.py │ └── integration_tests │ │ ├── __init__.py │ │ ├── add_lowres_and_cascade.py │ │ ├── cleanup_integration_test.py │ │ └── run_integration_test_bestconfig_inference.py ├── training │ ├── __init__.py │ ├── data_augmentation │ │ ├── __init__.py │ │ ├── compute_initial_patch_size.py │ │ └── custom_transforms │ │ │ ├── __init__.py │ │ │ ├── cascade_transforms.py │ │ │ ├── deep_supervision_donwsampling.py │ │ │ ├── limited_length_multithreaded_augmenter.py │ │ │ ├── manipulating_data_dict.py │ │ │ ├── masking.py │ │ │ ├── region_based_training.py │ │ │ └── transforms_for_dummy_2d.py │ ├── dataloading │ │ ├── __init__.py │ │ ├── base_data_loader.py │ │ ├── data_loader_2d.py │ │ ├── data_loader_3d.py │ │ ├── nnunet_dataset.py │ │ └── utils.py │ ├── logging │ │ ├── __init__.py │ │ └── nnunet_logger.py │ ├── loss │ │ ├── __init__.py │ │ ├── compound_losses.py │ │ ├── deep_supervision.py │ │ ├── dice.py │ │ └── robust_ce_loss.py │ ├── lr_scheduler │ │ ├── __init__.py │ │ └── polylr.py │ └── nnUNetTrainer │ │ ├── __init__.py │ │ ├── nnUNetTrainer.py │ │ └── variants │ │ ├── __init__.py │ │ ├── benchmarking │ │ ├── __init__.py │ │ ├── nnUNetTrainerBenchmark_5epochs.py │ │ └── nnUNetTrainerBenchmark_5epochs_noDataLoading.py │ │ ├── data_augmentation │ │ ├── __init__.py │ │ ├── nnUNetTrainerDA5.py │ │ ├── nnUNetTrainerDAOrd0.py │ │ ├── nnUNetTrainerNoDA.py │ │ └── nnUNetTrainerNoMirroring.py │ │ ├── loss │ │ ├── __init__.py │ │ ├── nnUNetTrainerCELoss.py │ │ ├── nnUNetTrainerDiceLoss.py │ │ └── nnUNetTrainerTopkLoss.py │ │ ├── lr_schedule │ │ ├── __init__.py │ │ └── nnUNetTrainerCosAnneal.py │ │ ├── network_architecture │ │ ├── __init__.py │ │ ├── nnUNetTrainerBN.py │ │ └── nnUNetTrainerNoDeepSupervision.py │ │ ├── optimizer │ │ ├── __init__.py │ │ ├── nnUNetTrainerAdam.py │ │ └── nnUNetTrainerAdan.py │ │ ├── sampling │ │ ├── __init__.py │ │ └── nnUNetTrainer_probabilisticOversampling.py │ │ └── training_length │ │ ├── __init__.py │ │ ├── nnUNetTrainer_Xepochs.py │ │ └── nnUNetTrainer_Xepochs_NoMirroring.py └── utilities │ ├── __init__.py │ ├── collate_outputs.py │ ├── dataset_name_id_conversion.py │ ├── ddp_allgather.py │ ├── default_n_proc_DA.py │ ├── file_path_utilities.py │ ├── find_class_by_name.py │ ├── get_network_from_plans.py │ ├── helpers.py │ ├── json_export.py │ ├── label_handling │ ├── __init__.py │ └── label_handling.py │ ├── network_initialization.py │ ├── overlay_plots.py │ ├── plans_handling │ ├── __init__.py │ └── plans_handler.py │ ├── tensor_utilities.py │ ├── unet.py │ ├── unet_decoder.py │ └── utils.py ├── pyproject.toml └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | env/ 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | 27 | # PyInstaller 28 | # Usually these files are written by a python script from a template 29 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 30 | *.manifest 31 | *.spec 32 | 33 | # Installer logs 34 | pip-log.txt 35 | pip-delete-this-directory.txt 36 | 37 | # Unit test / coverage reports 38 | htmlcov/ 39 | .tox/ 40 | .coverage 41 | .coverage.* 42 | .cache 43 | nosetests.xml 44 | coverage.xml 45 | *,cover 46 | .hypothesis/ 47 | 48 | # Translations 49 | *.mo 50 | *.pot 51 | 52 | # Django stuff: 53 | *.log 54 | local_settings.py 55 | 56 | # Flask stuff: 57 | instance/ 58 | .webassets-cache 59 | 60 | # Scrapy stuff: 61 | .scrapy 62 | 63 | # Sphinx documentation 64 | docs/_build/ 65 | 66 | # PyBuilder 67 | target/ 68 | 69 | # IPython Notebook 70 | .ipynb_checkpoints 71 | 72 | # pyenv 73 | .python-version 74 | 75 | # celery beat schedule file 76 | celerybeat-schedule 77 | 78 | # dotenv 79 | .env 80 | 81 | # virtualenv 82 | venv/ 83 | ENV/ 84 | 85 | # Spyder project settings 86 | .spyderproject 87 | 88 | # Rope project settings 89 | .ropeproject 90 | 91 | *.memmap 92 | *.png 93 | *.zip 94 | *.npz 95 | *.npy 96 | *.jpg 97 | *.jpeg 98 | .idea 99 | *.txt 100 | .idea/* 101 | *.png 102 | *.nii.gz 103 | *.nii 104 | *.tif 105 | *.bmp 106 | *.pkl 107 | *.xml 108 | *.pkl 109 | *.pdf 110 | *.png 111 | *.jpg 112 | *.jpeg 113 | 114 | *.model 115 | 116 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # nnSAM: Plug-and-play Segment Anything Model Improves nnUNet Performance 2 | 3 | ## Our entire code is built based on nnUNet, and you can follow the nnUNet instructions exactly. 4 | 5 | 6 | Install nnSAM depending on your use case: 7 | 8 | ```bash 9 | conda create -n nnsam python=3.9 10 | conda activate nnsam 11 | pip install git+https://github.com/ChaoningZhang/MobileSAM.git 12 | pip install timm 13 | pip install git+https://github.com/Kent0n-Li/nnSAM.git 14 | ``` 15 | 16 | It is important to input "set MODEL_NAME=nnsam" before using it. 17 | ```bash 18 | set MODEL_NAME=nnsam 19 | ``` 20 | 21 | ```bash 22 | nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity 23 | 24 | nnUNetv2_train DATASET_NAME_OR_ID UNET_CONFIGURATION FOLD [additional options, see -h] 25 | 26 | nnUNetv2_train DATASET_NAME_OR_ID UNET_CONFIGURATION FOLD --val --npz 27 | 28 | nnUNetv2_train DATASET_NAME_OR_ID 2d FOLD 29 | 30 | nnUNetv2_train DATASET_NAME_OR_ID 3d_fullres FOLD 31 | 32 | nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d DATASET_NAME_OR_ID -c CONFIGURATION --save_probabilities 33 | ``` 34 | 35 | 36 | ## How to get started? 37 | Read these: 38 | - [Dataset conversion](documentation/dataset_format.md) 39 | - [Usage instructions](documentation/how_to_use_nnunet.md) 40 | 41 | Additional information: 42 | - [Region-based training](documentation/region_based_training.md) 43 | - [Manual data splits](documentation/manual_data_splits.md) 44 | - [Pretraining and finetuning](documentation/pretraining_and_finetuning.md) 45 | - [Intensity Normalization in nnU-Net](documentation/explanation_normalization.md) 46 | - [Manually editing nnU-Net configurations](documentation/explanation_plans_files.md) 47 | - [Extending nnU-Net](documentation/extending_nnunet.md) 48 | - [What is different in V2?](documentation/changelog.md) 49 | 50 | [//]: # (- [Ignore label](documentation/ignore_label.md)) 51 | 52 | 53 | 54 | # Acknowledgements 55 | 56 | nnU-Net is developed and maintained by the Applied Computer Vision Lab (ACVL) of [Helmholtz Imaging](http://helmholtz-imaging.de) 57 | and the [Division of Medical Image Computing](https://www.dkfz.de/en/mic/index.php) at the 58 | [German Cancer Research Center (DKFZ)](https://www.dkfz.de/en/index.html). 59 | -------------------------------------------------------------------------------- /documentation/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/documentation/__init__.py -------------------------------------------------------------------------------- /documentation/changelog.md: -------------------------------------------------------------------------------- 1 | # What is different in v2? 2 | 3 | - We now support **hierarchical labels** (named regions in nnU-Net). For example, instead of training BraTS with the 4 | 'edema', 'necrosis' and 'enhancing tumor' labels you can directly train it on the target areas 'whole tumor', 5 | 'tumor core' and 'enhancing tumor'. See [here](region_based_training.md) for a detailed description + also have a look at the 6 | [BraTS 2021 conversion script](../nnunetv2/dataset_conversion/Dataset137_BraTS21.py). 7 | - Cross-platform support. Cuda, mps (Apple M1/M2) and of course CPU support! Simply select the device with 8 | `-device` in `nnUNetv2_train` and `nnUNetv2_predict`. 9 | - Unified trainer class: nnUNetTrainer. No messing around with cascaded trainer, DDP trainer, region-based trainer, 10 | ignore trainer etc. All default functionality is in there! 11 | - Supports more input/output data formats through ImageIO classes. 12 | - I/O formats can be extended by implementing new Adapters based on `BaseReaderWriter`. 13 | - The nnUNet_raw_cropped folder no longer exists -> saves disk space at no performance penalty. magic! (no jk the 14 | saving of cropped npz files was really slow, so it's actually faster to crop on the fly). 15 | - Preprocessed data and segmentation are stored in different files when unpacked. Seg is stored as int8 and thus 16 | takes 1/4 of the disk space per pixel (and I/O throughput) as in v1. 17 | - Native support for multi-GPU (DDP) TRAINING. 18 | Multi-GPU INFERENCE should still be run with `CUDA_VISIBLE_DEVICES=X nnUNetv2_predict [...] -num_parts Y -part_id X`. 19 | There is no cross-GPU communication in inference, so it doesn't make sense to add additional complexity with DDP. 20 | - All nnU-Net functionality is now also accessible via API. Check the corresponding entry point in `setup.py` to see 21 | what functions you need to call. 22 | - Dataset fingerprint is now explicitly created and saved in a json file (see nnUNet_preprocessed). 23 | 24 | - Complete overhaul of plans files (read also [this](explanation_plans_files.md): 25 | - Plans are now .json and can be opened and read more easily 26 | - Configurations are explicitly named ("3d_fullres" , ...) 27 | - Configurations can inherit from each other to make manual experimentation easier 28 | - A ton of additional functionality is now included in and can be changed through the plans, for example normalization strategy, resampling etc. 29 | - Stages of the cascade are now explicitly listed in the plans. 3d_lowres has 'next_stage' (which can also be a 30 | list of configurations!). 3d_cascade_fullres has a 'previous_stage' entry. By manually editing plans files you can 31 | now connect anything you want, for example 2d with 3d_fullres or whatever. Be wild! (But don't create cycles!) 32 | - Multiple configurations can point to the same preprocessed data folder to save disk space. Careful! Only 33 | configurations that use the same spacing, resampling, normalization etc. should share a data source! By default, 34 | 3d_fullres and 3d_cascade_fullres share the same data 35 | - Any number of configurations can be added to the plans (remember to give them a unique "data_identifier"!) 36 | 37 | Folder structures are different and more user-friendly: 38 | - nnUNet_preprocessed 39 | - By default, preprocessed data is now saved as: `nnUNet_preprocessed/DATASET_NAME/PLANS_IDENTIFIER_CONFIGURATION` to clearly link them to their corresponding plans and configuration 40 | - Name of the folder containing the preprocessed images can be adapted with the `data_identifier` key. 41 | - nnUNet_results 42 | - Results are now sorted as follows: DATASET_NAME/TRAINERCLASS__PLANSIDENTIFIER__CONFIGURATION/FOLD 43 | 44 | ## What other changes are planned and not yet implemented? 45 | - Integration into MONAI (together with our friends at Nvidia) 46 | - New pretrained weights for a large number of datasets (coming very soon)) 47 | 48 | 49 | [//]: # (- nnU-Net now also natively supports an **ignore label**. Pixels with this label will not contribute to the loss. ) 50 | 51 | [//]: # (Use this to learn from sparsely annotated data, or excluding irrelevant areas from training. Read more [here](ignore_label.md).) -------------------------------------------------------------------------------- /documentation/convert_msd_dataset.md: -------------------------------------------------------------------------------- 1 | Use `nnUNetv2_convert_MSD_dataset`. 2 | 3 | Read `nnUNetv2_convert_MSD_dataset -h` for usage instructions. -------------------------------------------------------------------------------- /documentation/dataset_format_inference.md: -------------------------------------------------------------------------------- 1 | # Data format for Inference 2 | Read the documentation on the overall [data format](dataset_format.md) first! 3 | 4 | The data format for inference must match the one used for the raw data (**specifically, the images must be in exactly 5 | the same format as in the imagesTr folder**). As before, the filenames must start with a 6 | unique identifier, followed by a 4-digit modality identifier. Here is an example for two different datasets: 7 | 8 | 1) Task005_Prostate: 9 | 10 | This task has 2 modalities, so the files in the input folder must look like this: 11 | 12 | input_folder 13 | ├── prostate_03_0000.nii.gz 14 | ├── prostate_03_0001.nii.gz 15 | ├── prostate_05_0000.nii.gz 16 | ├── prostate_05_0001.nii.gz 17 | ├── prostate_08_0000.nii.gz 18 | ├── prostate_08_0001.nii.gz 19 | ├── ... 20 | 21 | _0000 has to be the T2 image and _0001 has to be the ADC image (as specified by 'channel_names' in the 22 | dataset.json), exactly the same as was used for training. 23 | 24 | 2) Task002_Heart: 25 | 26 | imagesTs 27 | ├── la_001_0000.nii.gz 28 | ├── la_002_0000.nii.gz 29 | ├── la_006_0000.nii.gz 30 | ├── ... 31 | 32 | Task002 only has one modality, so each case only has one _0000.nii.gz file. 33 | 34 | 35 | The segmentations in the output folder will be named {CASE_IDENTIFIER}.nii.gz (omitting the modality identifier). 36 | 37 | Remember that the file format used for inference (.nii.gz in this example) must be the same as was used for training 38 | (and as was specified in 'file_ending' in the dataset.json)! 39 | -------------------------------------------------------------------------------- /documentation/explanation_normalization.md: -------------------------------------------------------------------------------- 1 | # Intensity normalization in nnU-Net 2 | 3 | The type of intensity normalization applied in nnU-Net can be controlled via the `channel_names` (former `modalities`) 4 | entry in the dataset.json. Just like the old nnU-Net, per-channel z-scoring as well as dataset-wide z-scoring based on 5 | foreground intensities are supported. However, there have been a few additions as well. 6 | 7 | Reminder: The `channel_names` entry typically looks like this: 8 | 9 | "channel_names": { 10 | "0": "T2", 11 | "1": "ADC" 12 | }, 13 | 14 | It has as many entries as there are input channels for the given dataset. 15 | 16 | To tell you a secret, nnU-Net does not really care what your channels are called. We just use this to determine what normalization 17 | scheme will be used for the given dataset. nnU-Net requires you to specify a normalization strategy for each of your input channels! 18 | If you enter a channel name that is not in the following list, the default (`zscore`) will be used. 19 | 20 | Here is a list of currently available normalization schemes: 21 | 22 | - `CT`: Perform CT normalization. Specifically, collect intensity values from the foreground classes (all but the 23 | background and ignore) from all training cases, compute the mean, standard deviation as well as the 0.5 and 24 | 99.5 percentile of the values. Then clip to the percentiles, followed by subtraction of the mean and division with the 25 | standard deviation. The normalization that is applied is the same for each training case (for this input channel). 26 | The values used by nnU-Net for normalization are stored in the `foreground_intensity_properties_per_channel` entry in the 27 | corresponding plans file. This normalization is suitable for modalities presenting physical quantities such as CT 28 | images and ADC maps. 29 | - `noNorm` : do not perform any normalization at all 30 | - `rescale_to_0_1`: rescale the intensities to [0, 1] 31 | - `rgb_to_0_1`: assumes uint8 inputs. Divides by 255 to rescale uint8 to [0, 1] 32 | - `zscore`/anything else: perform z-scoring (subtract mean and standard deviation) separately for each train case 33 | 34 | **Important:** The nnU-Net default is to perform 'CT' normalization for CT images and 'zscore' for everything else! If 35 | you deviate from that path, make sure to benchmark whether that actually improves results! 36 | 37 | # How to implement custom normalization strategies? 38 | - Head over to nnunetv2/preprocessing/normalization 39 | - implement a new image normalization class by deriving from ImageNormalization 40 | - register it in nnunetv2/preprocessing/normalization/map_channel_name_to_normalization.py:channel_name_to_normalization_mapping. 41 | This is where you specify a channel name that should be associated with it 42 | - use it by specifying the correct channel_name 43 | 44 | Normalization can only be applied to one channel at a time. There is currently no way of implementing a normalization scheme 45 | that gets multiple channels as input to be used jointly! -------------------------------------------------------------------------------- /documentation/extending_nnunet.md: -------------------------------------------------------------------------------- 1 | # Extending nnU-Net 2 | We hope that the new structure of nnU-Net v2 makes it much more intuitive on how to modify it! We cannot give an 3 | extensive tutorial on how each and every bit of it can be modified. It is better for you to search for the position 4 | in the repository where the thing you intend to change is implemented and start working your way through the code from 5 | there. Setting breakpoints and debugging into nnU-Net really helps in understanding it and thus will help you make the 6 | necessary modifications! 7 | 8 | Here are some things you might want to read before you start: 9 | - Editing nnU-Net configurations through plans files is really powerful now and allows you to change a lot of things regarding 10 | preprocessing, resampling, network topology etc. Read [this](explanation_plans_files.md)! 11 | - [Image normalization](explanation_normalization.md) and [i/o formats](dataset_format.md#supported-file-formats) are easy to extend! 12 | - Manual data splits can be defined as described [here](manual_data_splits.md) 13 | - You can chain arbitrary configurations together into cascades, see [this again](explanation_plans_files.md) 14 | - Read about our support for [region-based training](region_based_training.md) 15 | - If you intend to modify the training procedure (loss, sampling, data augmentation, lr scheduler, etc) then you need 16 | to implement your own trainer class. Best practice is to create a class that inherits from nnUNetTrainer and 17 | implements the necessary changes. Head over to our [trainer classes folder](../nnunetv2/training/nnUNetTrainer) for 18 | inspiration! There will be similar trainers for what you intend to change and you can take them as a guide. nnUNetTrainer 19 | are structured similarly to PyTorch lightning trainers, this should also make things easier! 20 | - Integrating new network architectures can be done in two ways: 21 | - Quick and dirty: implement a new nnUNetTrainer class and overwrite its `build_network_architecture` function. 22 | Make sure your architecture is compatible with deep supervision (if not, use `nnUNetTrainerNoDeepSupervision` 23 | as basis!) and that it can handle the patch sizes that are thrown at it! Your architecture should NOT apply any 24 | nonlinearities at the end (softmax, sigmoid etc). nnU-Net does that! 25 | - The 'proper' (but difficult) way: Build a dynamically configurable architecture such as the `PlainConvUNet` class 26 | used by default. It needs to have some sort of GPU memory estimation method that can be used to evaluate whether 27 | certain patch sizes and 28 | topologies fit into a specified GPU memory target. Build a new `ExperimentPlanner` that can configure your new 29 | class and communicate with its memory budget estimation. Run `nnUNetv2_plan_and_preprocess` while specifying your 30 | custom `ExperimentPlanner` and a custom `plans_name`. Implement a nnUNetTrainer that can use the plans generated by 31 | your `ExperimentPlanner` to instantiate the network architecture. Specify your plans and trainer when running `nnUNetv2_train`. 32 | It always pays off to first read and understand the corresponding nnU-Net code and use it as a template for your implementation! 33 | - Remember that multi-GPU training, region-based training, ignore label and cascaded training are now simply integrated 34 | into one unified nnUNetTrainer class. No separate classes needed (remember that when implementing your own trainer 35 | classes and ensure support for all of these features! Or raise `NotImplementedError`) 36 | 37 | [//]: # (- Read about our support for [ignore label](ignore_label.md) and [region-based training](region_based_training.md)) 38 | -------------------------------------------------------------------------------- /documentation/manual_data_splits.md: -------------------------------------------------------------------------------- 1 | # How to generate custom splits in nnU-Net 2 | 3 | Sometimes, the default 5-fold cross-validation split by nnU-Net does not fit a project. Maybe you want to run 3-fold 4 | cross-validation instead? Or maybe your training cases cannot be split randomly and require careful stratification. 5 | Fear not, for nnU-Net has got you covered (it really can do anything <3). 6 | 7 | The splits nnU-Net uses are generated in the `do_split` function of nnUNetTrainer. This function will first look for 8 | existing splits, stored as a file, and if no split exists it will create one. So if you wish to influence the split, 9 | manually creating a split file that will then be recognized and used is the way to go! 10 | 11 | The split file is located in the `nnUNet_preprocessed/DATASETXXX_NAME` folder. So it is best practice to first 12 | populate this folder by running `nnUNetv2_plan_and_preproccess`. 13 | 14 | Splits are stored as a .json file. They are a simple python list. The length of that list is the number of splits it 15 | contains (so it's 5 in the default nnU-Net). Each list entry is a dictionary with keys 'train' and 'val'. Values are 16 | again simply lists with the train identifiers in each set. To illustrate this, I am just messing with the Dataset002 17 | file as an example: 18 | 19 | ```commandline 20 | In [1]: from batchgenerators.utilities.file_and_folder_operations import load_json 21 | 22 | In [2]: splits = load_json('splits_final.json') 23 | 24 | In [3]: len(splits) 25 | Out[3]: 5 26 | 27 | In [4]: splits[0].keys() 28 | Out[4]: dict_keys(['train', 'val']) 29 | 30 | In [5]: len(splits[0]['train']) 31 | Out[5]: 16 32 | 33 | In [6]: len(splits[0]['val']) 34 | Out[6]: 4 35 | 36 | In [7]: print(splits[0]) 37 | {'train': ['la_003', 'la_004', 'la_005', 'la_009', 'la_010', 'la_011', 'la_014', 'la_017', 'la_018', 'la_019', 'la_020', 'la_022', 'la_023', 'la_026', 'la_029', 'la_030'], 38 | 'val': ['la_007', 'la_016', 'la_021', 'la_024']} 39 | ``` 40 | 41 | If you are still not sure what splits are supposed to look like, simply download some reference dataset from the 42 | [Medical Decathlon](http://medicaldecathlon.com/), start some training (to generate the splits) and manually inspect 43 | the .json file with your text editor of choice! 44 | 45 | In order to generate your custom splits, all you need to do is reproduce the data structure explained above and save it as 46 | `splits_final.json` in the `nnUNet_preprocessed/DATASETXXX_NAME` folder. Then use `nnUNetv2_train` etc. as usual. -------------------------------------------------------------------------------- /documentation/pretraining_and_finetuning.md: -------------------------------------------------------------------------------- 1 | # Pretraining with nnU-Net 2 | 3 | ## Intro 4 | 5 | So far nnU-Net only supports supervised pre-training, meaning that you train a regular nnU-Net on some source dataset 6 | and then use the final network weights as initialization for your target dataset. 7 | 8 | As a reminder, many training hyperparameters such as patch size and network topology differ between datasets as a 9 | result of the automated dataset analysis and experiment planning nnU-Net is known for. So, out of the box, it is not 10 | possible to simply take the network weights from some dataset and then reuse them for another. 11 | 12 | Consequently, the plans need to be aligned between the two tasks. In this README we show how this can be achieved and 13 | how the resulting weights can then be used for initialization. 14 | 15 | ### Terminology 16 | 17 | Throughout this README we use the following terminology: 18 | 19 | - `source dataset` is the dataset you intend to run the pretraining on 20 | - `target dataset` is the dataset you are interested in; the one you wish to fine tune on 21 | 22 | 23 | ## Pretraining on the source dataset 24 | 25 | In order to obtain matching network topologies we need to transfer the plans from one dataset to another. Since we are 26 | only interested in the target dataset, we first need to run experiment planning (and preprocessing) for it: 27 | 28 | ```bash 29 | nnUNetv2_plan_and_preprocess -d TARGET_DATASET 30 | ``` 31 | 32 | Then we need to extract the dataset fingerprint of the source dataset, if not yet available: 33 | 34 | ```bash 35 | nnUNetv2_extract_fingerprint -d SOURCE_DATASET 36 | ``` 37 | 38 | Now we can take the plans from the target dataset and transfer it to the source: 39 | 40 | ```bash 41 | nnUNetv2_move_plans_between_datasets -s TARGET_DATASET -t SOURCE_DATASET -sp TARGET_PLANS_IDENTIFIER -tp SOURCE_PLANS_IDENTIFIER 42 | ``` 43 | 44 | `SOURCE_PLANS_IDENTIFIER` is hereby probably nnUNetPlans unless you changed the experiment planner in 45 | nnUNetv2_plan_and_preprocess. For `TARGET_PLANS_IDENTIFIER` we recommend you set something custom in order to not 46 | overwrite default plans. 47 | 48 | Note that EVERYTHING is transferred between the datasets. Not just the network topology, batch size and patch size but 49 | also the normalization scheme! Therefore, a transfer between datasets that use different normalization schemes may not 50 | work well (but it could, depending on the schemes!). 51 | 52 | Note on CT normalization: Yes, also the clip values, mean and std are transferred! 53 | 54 | Now you can run the preprocessing on the source task: 55 | 56 | ```bash 57 | nnUNetv2_preprocess -d SOURCE_DATSET -plans_name TARGET_PLANS_IDENTIFIER 58 | ``` 59 | 60 | And run the training as usual: 61 | 62 | ```bash 63 | nnUNetv2_train SOURCE_DATSET CONFIG all -p TARGET_PLANS_IDENTIFIER 64 | ``` 65 | 66 | Note how we use the 'all' fold to train on all available data. For pretraining it does not make sense to split the data. 67 | 68 | ## Using pretrained weights 69 | 70 | Once pretraining is completed (or you obtain compatible weights by other means) you can use them to initialize your model: 71 | 72 | ```bash 73 | nnUNetv2_train TARGET_DATASET CONFIG FOLD -pretrained_weights PATH_TO_CHECKPOINT 74 | ``` 75 | 76 | Specify the checkpoint in PATH_TO_CHECKPOINT. 77 | 78 | When loading pretrained weights, all layers except the segmentation layers will be used! 79 | 80 | So far there are no specific nnUNet trainers for fine tuning, so the current recommendation is to just use 81 | nnUNetTrainer. You can however easily write your own trainers with learning rate ramp up, fine-tuning of segmentation 82 | heads or shorter training time. -------------------------------------------------------------------------------- /documentation/region_based_training.md: -------------------------------------------------------------------------------- 1 | # Region-based training 2 | 3 | ## What is this about? 4 | In some segmentation tasks, most prominently the 5 | [Brain Tumor Segmentation Challenge](http://braintumorsegmentation.org/), the target areas (based on which the metric 6 | will be computed) are different from the labels provided in the training data. This is the case because for some 7 | clinical applications, it is more relevant to detect the whole tumor, tumor core and enhancing tumor instead of the 8 | individual labels (edema, necrosis and non-enhancing tumor, enhancing tumor). 9 | 10 | 11 | 12 | The figure shows an example BraTS case along with label-based representation of the task (top) and region-based 13 | representation (bottom). The challenge evaluation is done on the regions. As we have shown in our 14 | [BraTS 2018 contribution](https://arxiv.org/abs/1809.10483), directly optimizing those 15 | overlapping areas over the individual labels yields better scoring models! 16 | 17 | ## What can nnU-Net do? 18 | nnU-Net's region-based training allows you to learn areas that are constructed by merging individual labels. For 19 | some segmentation tasks this provides a benefit, as this shifts the importance allocated to different labels during training. 20 | Most prominently, this feature can be used to represent **hierarchical classes**, for example when organs + 21 | substructures are to be segmented. Imagine a liver segmentation problem, where vessels and tumors are also to be 22 | segmented. The first target region could thus be the entire liver (including the substructures), while the remaining 23 | targets are the individual substructues. 24 | 25 | Important: nnU-Net still requires integer label maps as input and will produce integer label maps as output! 26 | Region-based training can be used to learn overlapping labels, but there must be a way to model these overlaps 27 | for nnU-Net to work (see below how this is done). 28 | 29 | ## How do you use it? 30 | 31 | When declaring the labels in the `dataset.json` file, BraTS would typically look like this: 32 | 33 | ```python 34 | ... 35 | "labels": { 36 | "background": 0, 37 | "edema": 1, 38 | "non_enhancing_and_necrosis": 2, 39 | "enhancing_tumor": 3 40 | }, 41 | ... 42 | ``` 43 | (we use different int values than the challenge because nnU-Net needs consecutive integers!) 44 | 45 | This representation corresponds to the upper row in the figure above. 46 | 47 | For region-based training, the labels need to be changed to the following: 48 | 49 | ```python 50 | ... 51 | "labels": { 52 | "background": 0, 53 | "whole_tumor": [1, 2, 3], 54 | "tumor_core": [2, 3], 55 | "enhancing_tumor": 3 # or [3] 56 | }, 57 | "regions_class_order": [1, 2, 3], 58 | ... 59 | ``` 60 | This corresponds to the bottom row in the figure above. Note how an additional entry in the dataset.json is 61 | required: `regions_class_order`. This tells nnU-Net how to convert the region representations back to an integer map. 62 | It essentially just tells nnU-Net what labels to place for which region in what order. The length of the 63 | list here needs to be the same as the number of regions (excl background). Each element in the list corresponds 64 | to the label that is placed instead of the region into the final segmentation. Later entries will overwrite earlier ones! 65 | Concretely, for the example given here, nnU-Net 66 | will firstly place the label 1 (edema) where the 'whole_tumor' region was predicted, then place the label 2 67 | (non-enhancing tumor and necrosis) where the "tumor_core" was predicted and finally place the label 3 in the 68 | predicted 'enhancing_tumor' area. With each step, part of the previously set pixels 69 | will be overwritten with the new label! So when setting your `regions_class_order`, place encompassing regions 70 | (like whole tumor etc) first, followed by substructures. 71 | 72 | **IMPORTANT** Because the conversion back to a segmentation map is sensitive to the order in which the regions are 73 | declared ("place label X in the first region") you need to make sure that this order is not perturbed! When 74 | automatically generating the dataset.json, make sure the dictionary keys do not get sorted alphabetically! Set 75 | `sort_keys=False` in `json.dump()`!!! 76 | 77 | nnU-Net will perform the evaluation + model selection also on the regions, not the individual labels! 78 | 79 | That's all. Easy, huh? -------------------------------------------------------------------------------- /documentation/run_inference_with_pretrained_models.md: -------------------------------------------------------------------------------- 1 | # How to run inference with pretrained models 2 | **Important:** Pretrained weights from nnU-Net v1 are NOT compatible with V2. You will need to retrain with the new 3 | version. But honestly, you already have a fully trained model with which you can run inference (in v1), so 4 | just continue using that! 5 | 6 | Not yet available for V2 :-( 7 | If you wish to run inference with pretrained models, check out the old nnU-Net for now. We are working on this full steam! 8 | -------------------------------------------------------------------------------- /documentation/set_environment_variables.md: -------------------------------------------------------------------------------- 1 | # How to set environment variables 2 | 3 | nnU-Net requires some environment variables so that it always knows where the raw data, preprocessed data and trained 4 | models are. Depending on the operating system, these environment variables need to be set in different ways. 5 | 6 | Variables can either be set permanently (recommended!) or you can decide to set them everytime you call nnU-Net. 7 | 8 | # Linux & MacOS 9 | 10 | ## Permanent 11 | Locate the `.bashrc` file in your home folder and add the following lines to the bottom: 12 | 13 | ```bash 14 | export nnUNet_raw="/media/fabian/nnUNet_raw" 15 | export nnUNet_preprocessed="/media/fabian/nnUNet_preprocessed" 16 | export nnUNet_results="/media/fabian/nnUNet_results" 17 | ``` 18 | 19 | (Of course you need to adapt the paths to the actual folders you intend to use). 20 | If you are using a different shell, such as zsh, you will need to find the correct script for it. For zsh this is `.zshrc`. 21 | 22 | ## Temporary 23 | Just execute the following lines whenever you run nnU-Net: 24 | ```bash 25 | export nnUNet_raw="/media/fabian/nnUNet_raw" 26 | export nnUNet_preprocessed="/media/fabian/nnUNet_preprocessed" 27 | export nnUNet_results="/media/fabian/nnUNet_results" 28 | ``` 29 | (Of course you need to adapt the paths to the actual folders you intend to use). 30 | 31 | Important: These variables will be deleted if you close your terminal! They will also only apply to the current 32 | terminal window and DO NOT transfer to other terminals! 33 | 34 | Alternatively you can also just prefix them to your nnU-Net commands: 35 | 36 | `nnUNet_results="/media/fabian/nnUNet_results" nnUNet_preprocessed="/media/fabian/nnUNet_preprocessed" nnUNetv2_train[...]` 37 | 38 | ## Verify that environment parameters are set 39 | You can always execute `echo ${nnUNet_raw}` etc to print the environment variables. This will return an empty string if 40 | they were not set. 41 | 42 | # Windows 43 | Useful links: 44 | - [https://www3.ntu.edu.sg](https://www3.ntu.edu.sg/home/ehchua/programming/howto/Environment_Variables.html#:~:text=To%20set%20(or%20change)%20a,it%20to%20an%20empty%20string.) 45 | - [https://phoenixnap.com](https://phoenixnap.com/kb/windows-set-environment-variable) 46 | 47 | ## Permanent 48 | See `Set Environment Variable in Windows via GUI` [here](https://phoenixnap.com/kb/windows-set-environment-variable). 49 | Or read about setx (command prompt). 50 | 51 | ## Temporary 52 | Just execute the following before you run nnU-Net: 53 | 54 | (PowerShell) 55 | ```PowerShell 56 | $Env:nnUNet_raw = "C:/Users/fabian/nnUNet_raw" 57 | $Env:nnUNet_preprocessed = "C:/Users/fabian/nnUNet_preprocessed" 58 | $Env:nnUNet_results = "C:/Users/fabian/nnUNet_results" 59 | ``` 60 | 61 | (Command Prompt) 62 | ```Command Prompt 63 | set nnUNet_raw=C:/Users/fabian/nnUNet_raw 64 | set nnUNet_preprocessed=C:/Users/fabian/nnUNet_preprocessed 65 | set nnUNet_results=C:/Users/fabian/fabian/nnUNet_results 66 | ``` 67 | 68 | (Of course you need to adapt the paths to the actual folders you intend to use). 69 | 70 | Important: These variables will be deleted if you close your session! They will also only apply to the current 71 | window and DO NOT transfer to other sessions! 72 | 73 | ## Verify that environment parameters are set 74 | Printing in Windows works differently depending on the environment you are in: 75 | 76 | PowerShell: `echo $Env:[variable_name]` 77 | 78 | Command Prompt: `echo %[variable_name]%` 79 | -------------------------------------------------------------------------------- /documentation/setting_up_paths.md: -------------------------------------------------------------------------------- 1 | # Setting up Paths 2 | 3 | nnU-Net relies on environment variables to know where raw data, preprocessed data and trained model weights are stored. 4 | To use the full functionality of nnU-Net, the following three environment variables must be set: 5 | 6 | 1) `nnUNet_raw`: This is where you place the raw datasets. This folder will have one subfolder for each dataset names 7 | DatasetXXX_YYY where XXX is a 3-digit identifier (such as 001, 002, 043, 999, ...) and YYY is the (unique) 8 | dataset name. The datasets must be in nnU-Net format, see [here](dataset_format.md). 9 | 10 | Example tree structure: 11 | ``` 12 | nnUNet_raw/Dataset001_NAME1 13 | ├── dataset.json 14 | ├── imagesTr 15 | │   ├── ... 16 | ├── imagesTs 17 | │   ├── ... 18 | └── labelsTr 19 | ├── ... 20 | nnUNet_raw/Dataset002_NAME2 21 | ├── dataset.json 22 | ├── imagesTr 23 | │   ├── ... 24 | ├── imagesTs 25 | │   ├── ... 26 | └── labelsTr 27 | ├── ... 28 | ``` 29 | 30 | 2) `nnUNet_preprocessed`: This is the folder where the preprocessed data will be saved. The data will also be read from 31 | this folder during training. It is important that this folder is located on a drive with low access latency and high 32 | throughput (such as a nvme SSD (PCIe gen 3 is sufficient)). 33 | 34 | 3) `nnUNet_results`: This specifies where nnU-Net will save the model weights. If pretrained models are downloaded, this 35 | is where it will save them. 36 | 37 | ### How to set environment variables 38 | See [here](set_environment_variables.md). -------------------------------------------------------------------------------- /documentation/tldr_migration_guide_from_v1.md: -------------------------------------------------------------------------------- 1 | # TLDR Migration Guide from nnU-Net V1 2 | 3 | - nnU-Net V2 can be installed simultaneously with V1. They won't get in each other's way 4 | - The environment variables needed for V2 have slightly different names. Read [this](setting_up_paths.md). 5 | - nnU-Net V2 datasets are called DatasetXXX_NAME. Not Task. 6 | - Datasets have the same structure (imagesTr, labelsTr, dataset.json) but we now support more 7 | [file types](dataset_format.md#supported-file-formats). The dataset.json is simplified. Use `generate_dataset_json` 8 | from nnunetv2.dataset_conversion.generate_dataset_json.py. 9 | - Careful: labels are now no longer declared as value:name but name:value. This has to do with [hierarchical labels](region_based_training.md). 10 | - nnU-Net v2 commands start with `nnUNetv2...`. They work mostly (but not entirely) the same. Just use the `-h` option. 11 | - You can transfer your V1 raw datasets to V2 with `nnUNetv2_convert_old_nnUNet_dataset`. You cannot transfer trained 12 | models. Continue to use the old nnU-Net Version for making inference with those. 13 | - These are the commands you are most likely to be using (in that order) 14 | - `nnUNetv2_plan_and_preprocess`. Example: `nnUNetv2_plan_and_preprocess -d 2` 15 | - `nnUNetv2_train`. Example: `nnUNetv2_train 2 3d_fullres 0` 16 | - `nnUNetv2_find_best_configuration`. Example: `nnUNetv2_find_best_configuration 2 -c 2d 3d_fullres`. This command 17 | will now create a `inference_instructions.txt` file in your `nnUNet_preprocessed/DatasetXXX_NAME/` folder which 18 | tells you exactly how to do inference. 19 | - `nnUNetv2_predict`. Example: `nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -c 3d_fullres -d 2` 20 | - `nnUNetv2_apply_postprocessing` (see inference_instructions.txt) 21 | -------------------------------------------------------------------------------- /nnunetv2/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/__init__.py -------------------------------------------------------------------------------- /nnunetv2/batch_running/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/batch_running/__init__.py -------------------------------------------------------------------------------- /nnunetv2/batch_running/benchmarking/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/batch_running/benchmarking/__init__.py -------------------------------------------------------------------------------- /nnunetv2/batch_running/benchmarking/generate_benchmarking_commands.py: -------------------------------------------------------------------------------- 1 | if __name__ == '__main__': 2 | """ 3 | This code probably only works within the DKFZ infrastructure (using LSF). You will need to adapt it to your scheduler! 4 | """ 5 | gpu_models = ['NVIDIAA100_PCIE_40GB', 'NVIDIAGeForceRTX2080Ti', 'NVIDIATITANRTX', 'TeslaV100_SXM2_32GB', 6 | 'NVIDIAA100_SXM4_40GB', 'TeslaV100_PCIE_32GB'] 7 | datasets = [2, 3, 4, 5] 8 | trainers = ['nnUNetTrainerBenchmark_5epochs', 'nnUNetTrainerBenchmark_5epochs_noDataLoading'] 9 | plans = ['nnUNetPlans'] 10 | configs = ['2d', '3d_fullres'] 11 | num_gpus = 1 12 | 13 | benchmark_configurations = {d: configs for d in datasets} 14 | 15 | exclude_hosts = "" 16 | resources = "-R \"tensorcore\"" 17 | queue = "-q gpu" 18 | preamble = "-L /bin/bash \"source ~/load_env_cluster4.sh && " 19 | train_command = 'nnUNetv2_train' 20 | 21 | folds = (0, ) 22 | 23 | use_these_modules = { 24 | tr: plans for tr in trainers 25 | } 26 | 27 | additional_arguments = f' -num_gpus {num_gpus}' # '' 28 | 29 | output_file = "/home/isensee/deleteme.txt" 30 | with open(output_file, 'w') as f: 31 | for g in gpu_models: 32 | gpu_requirements = f"-gpu num={num_gpus}:j_exclusive=yes:gmodel={g}" 33 | for tr in use_these_modules.keys(): 34 | for p in use_these_modules[tr]: 35 | for dataset in benchmark_configurations.keys(): 36 | for config in benchmark_configurations[dataset]: 37 | for fl in folds: 38 | command = f'bsub {exclude_hosts} {resources} {queue} {gpu_requirements} {preamble} {train_command} {dataset} {config} {fl} -tr {tr} -p {p}' 39 | if additional_arguments is not None and len(additional_arguments) > 0: 40 | command += f' {additional_arguments}' 41 | f.write(f'{command}\"\n') -------------------------------------------------------------------------------- /nnunetv2/batch_running/benchmarking/summarize_benchmark_results.py: -------------------------------------------------------------------------------- 1 | from batchgenerators.utilities.file_and_folder_operations import join, load_json, isfile 2 | from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name 3 | from nnunetv2.paths import nnUNet_results 4 | from nnunetv2.utilities.file_path_utilities import get_output_folder 5 | 6 | if __name__ == '__main__': 7 | trainers = ['nnUNetTrainerBenchmark_5epochs', 'nnUNetTrainerBenchmark_5epochs_noDataLoading'] 8 | datasets = [2, 3, 4, 5] 9 | plans = ['nnUNetPlans'] 10 | configs = ['2d', '3d_fullres'] 11 | output_file = join(nnUNet_results, 'benchmark_results.csv') 12 | 13 | torch_version = "1.12.0a0+git664058f" #"1.11.0a0+gitbc2c6ed" # 14 | cudnn_version = 8500 # 8302 # 15 | num_gpus = 1 16 | 17 | unique_gpus = set() 18 | 19 | # collect results in the most janky way possible. Amazing coding skills! 20 | all_results = {} 21 | for tr in trainers: 22 | all_results[tr] = {} 23 | for p in plans: 24 | all_results[tr][p] = {} 25 | for c in configs: 26 | all_results[tr][p][c] = {} 27 | for d in datasets: 28 | dataset_name = maybe_convert_to_dataset_name(d) 29 | output_folder = get_output_folder(dataset_name, tr, p, c, fold=0) 30 | expected_benchmark_file = join(output_folder, 'benchmark_result.json') 31 | all_results[tr][p][c][d] = {} 32 | if isfile(expected_benchmark_file): 33 | # filter results for what we want 34 | results = [i for i in load_json(expected_benchmark_file).values() 35 | if i['num_gpus'] == num_gpus and i['cudnn_version'] == cudnn_version and 36 | i['torch_version'] == torch_version] 37 | for r in results: 38 | all_results[tr][p][c][d][r['gpu_name']] = r 39 | unique_gpus.add(r['gpu_name']) 40 | 41 | # haha. Fuck this. Collect GPUs in the code above. 42 | # unique_gpus = np.unique([i["gpu_name"] for tr in trainers for p in plans for c in configs for d in datasets for i in all_results[tr][p][c][d]]) 43 | 44 | unique_gpus = list(unique_gpus) 45 | unique_gpus.sort() 46 | 47 | with open(output_file, 'w') as f: 48 | f.write('Dataset,Trainer,Plans,Config') 49 | for g in unique_gpus: 50 | f.write(f",{g}") 51 | f.write("\n") 52 | for d in datasets: 53 | for tr in trainers: 54 | for p in plans: 55 | for c in configs: 56 | gpu_results = [] 57 | for g in unique_gpus: 58 | if g in all_results[tr][p][c][d].keys(): 59 | gpu_results.append(round(all_results[tr][p][c][d][g]["fastest_epoch"], ndigits=2)) 60 | else: 61 | gpu_results.append("MISSING") 62 | # skip if all are missing 63 | if all([i == 'MISSING' for i in gpu_results]): 64 | continue 65 | f.write(f"{d},{tr},{p},{c}") 66 | for g in gpu_results: 67 | f.write(f",{g}") 68 | f.write("\n") 69 | f.write("\n") 70 | 71 | -------------------------------------------------------------------------------- /nnunetv2/batch_running/collect_results_custom_Decathlon_2d.py: -------------------------------------------------------------------------------- 1 | from batchgenerators.utilities.file_and_folder_operations import * 2 | 3 | from nnunetv2.batch_running.collect_results_custom_Decathlon import collect_results, summarize 4 | from nnunetv2.paths import nnUNet_results 5 | 6 | if __name__ == '__main__': 7 | use_these_trainers = { 8 | 'nnUNetTrainer': ('nnUNetPlans', ), 9 | } 10 | all_results_file = join(nnUNet_results, 'hrnet_results.csv') 11 | datasets = [2, 3, 4, 17, 20, 24, 27, 38, 55, 64, 82] 12 | collect_results(use_these_trainers, datasets, all_results_file) 13 | 14 | folds = (0, ) 15 | configs = ('2d', ) 16 | output_file = join(nnUNet_results, 'hrnet_results_summary_fold0.csv') 17 | summarize(all_results_file, output_file, folds, configs, datasets, use_these_trainers) 18 | 19 | -------------------------------------------------------------------------------- /nnunetv2/batch_running/generate_lsf_runs_customDecathlon.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | import numpy as np 3 | 4 | 5 | def merge(dict1, dict2): 6 | keys = np.unique(list(dict1.keys()) + list(dict2.keys())) 7 | keys = np.unique(keys) 8 | res = {} 9 | for k in keys: 10 | all_configs = [] 11 | if dict1.get(k) is not None: 12 | all_configs += list(dict1[k]) 13 | if dict2.get(k) is not None: 14 | all_configs += list(dict2[k]) 15 | if len(all_configs) > 0: 16 | res[k] = tuple(np.unique(all_configs)) 17 | return res 18 | 19 | 20 | if __name__ == "__main__": 21 | # after the Nature Methods paper we switch our evaluation to a different (more stable/high quality) set of 22 | # datasets for evaluation and future development 23 | configurations_all = { 24 | 2: ("3d_fullres", "2d"), 25 | 3: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 26 | 4: ("2d", "3d_fullres"), 27 | 17: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 28 | 20: ("2d", "3d_fullres"), 29 | 24: ("2d", "3d_fullres"), 30 | 27: ("2d", "3d_fullres"), 31 | 38: ("2d", "3d_fullres"), 32 | 55: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 33 | 64: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 34 | 82: ("2d", "3d_fullres"), 35 | # 83: ("2d", "3d_fullres"), 36 | } 37 | 38 | configurations_3d_fr_only = { 39 | i: ("3d_fullres", ) for i in configurations_all if "3d_fullres" in configurations_all[i] 40 | } 41 | 42 | configurations_3d_c_only = { 43 | i: ("3d_cascade_fullres", ) for i in configurations_all if "3d_cascade_fullres" in configurations_all[i] 44 | } 45 | 46 | configurations_3d_lr_only = { 47 | i: ("3d_lowres", ) for i in configurations_all if "3d_lowres" in configurations_all[i] 48 | } 49 | 50 | configurations_2d_only = { 51 | i: ("2d", ) for i in configurations_all if "2d" in configurations_all[i] 52 | } 53 | 54 | num_gpus = 1 55 | exclude_hosts = "-R \"select[hname!='e230-dgx2-2']\" -R \"select[hname!='e230-dgx2-1']\" -R \"select[hname!='e230-dgx1-1']\" -R \"select[hname!='e230-dgxa100-1']\" -R \"select[hname!='e230-dgxa100-2']\" -R \"select[hname!='e230-dgxa100-3']\" -R \"select[hname!='e230-dgxa100-4']\"" 56 | resources = "-R \"tensorcore\"" 57 | gpu_requirements = f"-gpu num={num_gpus}:j_exclusive=yes:gmem=33G" 58 | queue = "-q gpu-lowprio" 59 | preamble = "-L /bin/bash \"source ~/load_env_cluster4.sh && " 60 | train_command = 'nnUNet_results=/dkfz/cluster/gpu/checkpoints/OE0441/isensee/nnUNet_results_remake_release nnUNetv2_train' 61 | 62 | folds = (0, ) 63 | # use_this = configurations_2d_only 64 | use_this = merge(configurations_3d_fr_only, configurations_3d_lr_only) 65 | # use_this = merge(use_this, configurations_3d_c_only) 66 | 67 | use_these_modules = { 68 | 'nnUNetTrainer': ('nnUNetPlans',), 69 | 'nnUNetTrainerDiceCELoss_noSmooth': ('nnUNetPlans',), 70 | # 'nnUNetTrainer_DASegOrd0': ('nnUNetPlans',), 71 | } 72 | 73 | additional_arguments = f'--disable_checkpointing -num_gpus {num_gpus}' # '' 74 | 75 | output_file = "/home/isensee/deleteme.txt" 76 | with open(output_file, 'w') as f: 77 | for tr in use_these_modules.keys(): 78 | for p in use_these_modules[tr]: 79 | for dataset in use_this.keys(): 80 | for config in use_this[dataset]: 81 | for fl in folds: 82 | command = f'bsub {exclude_hosts} {resources} {queue} {gpu_requirements} {preamble} {train_command} {dataset} {config} {fl} -tr {tr} -p {p}' 83 | if additional_arguments is not None and len(additional_arguments) > 0: 84 | command += f' {additional_arguments}' 85 | f.write(f'{command}\"\n') 86 | 87 | -------------------------------------------------------------------------------- /nnunetv2/batch_running/release_trainings/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/batch_running/release_trainings/__init__.py -------------------------------------------------------------------------------- /nnunetv2/batch_running/release_trainings/nnunetv2_v1/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/batch_running/release_trainings/nnunetv2_v1/__init__.py -------------------------------------------------------------------------------- /nnunetv2/batch_running/release_trainings/nnunetv2_v1/generate_lsf_commands.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | import numpy as np 3 | 4 | 5 | def merge(dict1, dict2): 6 | keys = np.unique(list(dict1.keys()) + list(dict2.keys())) 7 | keys = np.unique(keys) 8 | res = {} 9 | for k in keys: 10 | all_configs = [] 11 | if dict1.get(k) is not None: 12 | all_configs += list(dict1[k]) 13 | if dict2.get(k) is not None: 14 | all_configs += list(dict2[k]) 15 | if len(all_configs) > 0: 16 | res[k] = tuple(np.unique(all_configs)) 17 | return res 18 | 19 | 20 | if __name__ == "__main__": 21 | # after the Nature Methods paper we switch our evaluation to a different (more stable/high quality) set of 22 | # datasets for evaluation and future development 23 | configurations_all = { 24 | # 1: ("3d_fullres", "2d"), 25 | 2: ("3d_fullres", "2d"), 26 | # 3: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 27 | # 4: ("2d", "3d_fullres"), 28 | 5: ("2d", "3d_fullres"), 29 | # 6: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 30 | # 7: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 31 | # 8: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 32 | # 9: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 33 | # 10: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 34 | # 17: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 35 | 20: ("2d", "3d_fullres"), 36 | 24: ("2d", "3d_fullres"), 37 | 27: ("2d", "3d_fullres"), 38 | 35: ("2d", "3d_fullres"), 39 | 38: ("2d", "3d_fullres"), 40 | # 55: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 41 | # 64: ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), 42 | # 82: ("2d", "3d_fullres"), 43 | # 83: ("2d", "3d_fullres"), 44 | } 45 | 46 | configurations_3d_fr_only = { 47 | i: ("3d_fullres", ) for i in configurations_all if "3d_fullres" in configurations_all[i] 48 | } 49 | 50 | configurations_3d_c_only = { 51 | i: ("3d_cascade_fullres", ) for i in configurations_all if "3d_cascade_fullres" in configurations_all[i] 52 | } 53 | 54 | configurations_3d_lr_only = { 55 | i: ("3d_lowres", ) for i in configurations_all if "3d_lowres" in configurations_all[i] 56 | } 57 | 58 | configurations_2d_only = { 59 | i: ("2d", ) for i in configurations_all if "2d" in configurations_all[i] 60 | } 61 | 62 | num_gpus = 1 63 | exclude_hosts = "-R \"select[hname!='e230-dgx2-2']\" -R \"select[hname!='e230-dgx2-1']\"" 64 | resources = "-R \"tensorcore\"" 65 | gpu_requirements = f"-gpu num={num_gpus}:j_exclusive=yes:gmem=1G" 66 | queue = "-q gpu-lowprio" 67 | preamble = "-L /bin/bash \"source ~/load_env_cluster4.sh && " 68 | train_command = 'nnUNet_keep_files_open=True nnUNet_results=/dkfz/cluster/gpu/data/OE0441/isensee/nnUNet_results_remake_release_normfix nnUNetv2_train' 69 | 70 | folds = (0, 1, 2, 3, 4) 71 | # use_this = configurations_2d_only 72 | # use_this = merge(configurations_3d_fr_only, configurations_3d_lr_only) 73 | # use_this = merge(use_this, configurations_3d_c_only) 74 | use_this = configurations_all 75 | 76 | use_these_modules = { 77 | 'nnUNetTrainer': ('nnUNetPlans',), 78 | } 79 | 80 | additional_arguments = f'--disable_checkpointing -num_gpus {num_gpus}' # '' 81 | 82 | output_file = "/home/isensee/deleteme.txt" 83 | with open(output_file, 'w') as f: 84 | for tr in use_these_modules.keys(): 85 | for p in use_these_modules[tr]: 86 | for dataset in use_this.keys(): 87 | for config in use_this[dataset]: 88 | for fl in folds: 89 | command = f'bsub {exclude_hosts} {resources} {queue} {gpu_requirements} {preamble} {train_command} {dataset} {config} {fl} -tr {tr} -p {p}' 90 | if additional_arguments is not None and len(additional_arguments) > 0: 91 | command += f' {additional_arguments}' 92 | f.write(f'{command}\"\n') 93 | 94 | -------------------------------------------------------------------------------- /nnunetv2/configuration.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from nnunetv2.utilities.default_n_proc_DA import get_allowed_n_proc_DA 4 | 5 | default_num_processes = 8 if 'nnUNet_def_n_proc' not in os.environ else int(os.environ['nnUNet_def_n_proc']) 6 | 7 | ANISO_THRESHOLD = 3 # determines when a sample is considered anisotropic (3 means that the spacing in the low 8 | # resolution axis must be 3x as large as the next largest spacing) 9 | 10 | default_n_proc_DA = get_allowed_n_proc_DA() 11 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/Dataset073_Fluo_C3DH_A549_SIM.py: -------------------------------------------------------------------------------- 1 | from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json 2 | from nnunetv2.paths import nnUNet_raw, nnUNet_preprocessed 3 | import tifffile 4 | from batchgenerators.utilities.file_and_folder_operations import * 5 | import shutil 6 | 7 | 8 | if __name__ == '__main__': 9 | """ 10 | This is going to be my test dataset for working with tif as input and output images 11 | 12 | All we do here is copy the files and rename them. Not file conversions take place 13 | """ 14 | dataset_name = 'Dataset073_Fluo_C3DH_A549_SIM' 15 | 16 | imagestr = join(nnUNet_raw, dataset_name, 'imagesTr') 17 | imagests = join(nnUNet_raw, dataset_name, 'imagesTs') 18 | labelstr = join(nnUNet_raw, dataset_name, 'labelsTr') 19 | maybe_mkdir_p(imagestr) 20 | maybe_mkdir_p(imagests) 21 | maybe_mkdir_p(labelstr) 22 | 23 | # we extract the downloaded train and test datasets to two separate folders and name them Fluo-C3DH-A549-SIM_train 24 | # and Fluo-C3DH-A549-SIM_test 25 | train_source = '/home/fabian/Downloads/Fluo-C3DH-A549-SIM_train' 26 | test_source = '/home/fabian/Downloads/Fluo-C3DH-A549-SIM_test' 27 | 28 | # with the old nnU-Net we had to convert all the files to nifti. This is no longer required. We can just copy the 29 | # tif files 30 | 31 | # tif is broken when it comes to spacing. No standards. Grr. So when we use tif nnU-Net expects a separate file 32 | # that specifies the spacing. This file needs to exist for EVERY training/test case to allow for different spacings 33 | # between files. Important! The spacing must align with the axes. 34 | # Here when we do print(tifffile.imread('IMAGE').shape) we get (29, 300, 350). The low resolution axis is the first. 35 | # The spacing on the website is griven in the wrong axis order. Great. 36 | spacing = (1, 0.126, 0.126) 37 | 38 | # train set 39 | for seq in ['01', '02']: 40 | images_dir = join(train_source, seq) 41 | seg_dir = join(train_source, seq + '_GT', 'SEG') 42 | # if we were to be super clean we would go by IDs but here we just trust the files are sorted the correct way. 43 | # Simpler filenames in the cell tracking challenge would be soooo nice. 44 | images = subfiles(images_dir, suffix='.tif', sort=True, join=False) 45 | segs = subfiles(seg_dir, suffix='.tif', sort=True, join=False) 46 | for i, (im, se) in enumerate(zip(images, segs)): 47 | target_name = f'{seq}_image_{i:03d}' 48 | # we still need the '_0000' suffix for images! Otherwise we would not be able to support multiple input 49 | # channels distributed over separate files 50 | shutil.copy(join(images_dir, im), join(imagestr, target_name + '_0000.tif')) 51 | # spacing file! 52 | save_json({'spacing': spacing}, join(imagestr, target_name + '.json')) 53 | shutil.copy(join(seg_dir, se), join(labelstr, target_name + '.tif')) 54 | # spacing file! 55 | save_json({'spacing': spacing}, join(labelstr, target_name + '.json')) 56 | 57 | # test set, same a strain just without the segmentations 58 | for seq in ['01', '02']: 59 | images_dir = join(test_source, seq) 60 | images = subfiles(images_dir, suffix='.tif', sort=True, join=False) 61 | for i, im in enumerate(images): 62 | target_name = f'{seq}_image_{i:03d}' 63 | shutil.copy(join(images_dir, im), join(imagests, target_name + '_0000.tif')) 64 | # spacing file! 65 | save_json({'spacing': spacing}, join(imagests, target_name + '.json')) 66 | 67 | # now we generate the dataset json 68 | generate_dataset_json( 69 | join(nnUNet_raw, dataset_name), 70 | {0: 'fluorescence_microscopy'}, 71 | {'background': 0, 'cell': 1}, 72 | 60, 73 | '.tif' 74 | ) 75 | 76 | # custom split to ensure we are stratifying properly. This dataset only has 2 folds 77 | caseids = [i[:-4] for i in subfiles(labelstr, suffix='.tif', join=False)] 78 | splits = [] 79 | splits.append( 80 | {'train': [i for i in caseids if i.startswith('01_')], 'val': [i for i in caseids if i.startswith('02_')]} 81 | ) 82 | splits.append( 83 | {'train': [i for i in caseids if i.startswith('02_')], 'val': [i for i in caseids if i.startswith('01_')]} 84 | ) 85 | save_json(splits, join(nnUNet_preprocessed, dataset_name, 'splits_final.json')) -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/Dataset120_RoadSegmentation.py: -------------------------------------------------------------------------------- 1 | import multiprocessing 2 | import shutil 3 | from multiprocessing import Pool 4 | 5 | from batchgenerators.utilities.file_and_folder_operations import * 6 | 7 | from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json 8 | from nnunetv2.paths import nnUNet_raw 9 | from skimage import io 10 | from acvl_utils.morphology.morphology_helper import generic_filter_components 11 | from scipy.ndimage import binary_fill_holes 12 | 13 | 14 | def load_and_covnert_case(input_image: str, input_seg: str, output_image: str, output_seg: str, 15 | min_component_size: int = 50): 16 | seg = io.imread(input_seg) 17 | seg[seg == 255] = 1 18 | image = io.imread(input_image) 19 | image = image.sum(2) 20 | mask = image == (3 * 255) 21 | # the dataset has large white areas in which road segmentations can exist but no image information is available. 22 | # Remove the road label in these areas 23 | mask = generic_filter_components(mask, filter_fn=lambda ids, sizes: [i for j, i in enumerate(ids) if 24 | sizes[j] > min_component_size]) 25 | mask = binary_fill_holes(mask) 26 | seg[mask] = 0 27 | io.imsave(output_seg, seg, check_contrast=False) 28 | shutil.copy(input_image, output_image) 29 | 30 | 31 | if __name__ == "__main__": 32 | # extracted archive from https://www.kaggle.com/datasets/insaff/massachusetts-roads-dataset?resource=download 33 | source = '/media/fabian/data/raw_datasets/Massachussetts_road_seg/road_segmentation_ideal' 34 | 35 | dataset_name = 'Dataset120_RoadSegmentation' 36 | 37 | imagestr = join(nnUNet_raw, dataset_name, 'imagesTr') 38 | imagests = join(nnUNet_raw, dataset_name, 'imagesTs') 39 | labelstr = join(nnUNet_raw, dataset_name, 'labelsTr') 40 | labelsts = join(nnUNet_raw, dataset_name, 'labelsTs') 41 | maybe_mkdir_p(imagestr) 42 | maybe_mkdir_p(imagests) 43 | maybe_mkdir_p(labelstr) 44 | maybe_mkdir_p(labelsts) 45 | 46 | train_source = join(source, 'training') 47 | test_source = join(source, 'testing') 48 | 49 | with multiprocessing.get_context("spawn").Pool(8) as p: 50 | 51 | # not all training images have a segmentation 52 | valid_ids = subfiles(join(train_source, 'output'), join=False, suffix='png') 53 | num_train = len(valid_ids) 54 | r = [] 55 | for v in valid_ids: 56 | r.append( 57 | p.starmap_async( 58 | load_and_covnert_case, 59 | (( 60 | join(train_source, 'input', v), 61 | join(train_source, 'output', v), 62 | join(imagestr, v[:-4] + '_0000.png'), 63 | join(labelstr, v), 64 | 50 65 | ),) 66 | ) 67 | ) 68 | 69 | # test set 70 | valid_ids = subfiles(join(test_source, 'output'), join=False, suffix='png') 71 | for v in valid_ids: 72 | r.append( 73 | p.starmap_async( 74 | load_and_covnert_case, 75 | (( 76 | join(test_source, 'input', v), 77 | join(test_source, 'output', v), 78 | join(imagests, v[:-4] + '_0000.png'), 79 | join(labelsts, v), 80 | 50 81 | ),) 82 | ) 83 | ) 84 | _ = [i.get() for i in r] 85 | 86 | generate_dataset_json(join(nnUNet_raw, dataset_name), {0: 'R', 1: 'G', 2: 'B'}, {'background': 0, 'road': 1}, 87 | num_train, '.png', dataset_name) 88 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/Dataset137_BraTS21.py: -------------------------------------------------------------------------------- 1 | import multiprocessing 2 | import shutil 3 | from multiprocessing import Pool 4 | 5 | import SimpleITK as sitk 6 | import numpy as np 7 | from batchgenerators.utilities.file_and_folder_operations import * 8 | from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json 9 | from nnunetv2.paths import nnUNet_raw 10 | 11 | 12 | def copy_BraTS_segmentation_and_convert_labels_to_nnUNet(in_file: str, out_file: str) -> None: 13 | # use this for segmentation only!!! 14 | # nnUNet wants the labels to be continuous. BraTS is 0, 1, 2, 4 -> we make that into 0, 1, 2, 3 15 | img = sitk.ReadImage(in_file) 16 | img_npy = sitk.GetArrayFromImage(img) 17 | 18 | uniques = np.unique(img_npy) 19 | for u in uniques: 20 | if u not in [0, 1, 2, 4]: 21 | raise RuntimeError('unexpected label') 22 | 23 | seg_new = np.zeros_like(img_npy) 24 | seg_new[img_npy == 4] = 3 25 | seg_new[img_npy == 2] = 1 26 | seg_new[img_npy == 1] = 2 27 | img_corr = sitk.GetImageFromArray(seg_new) 28 | img_corr.CopyInformation(img) 29 | sitk.WriteImage(img_corr, out_file) 30 | 31 | 32 | def convert_labels_back_to_BraTS(seg: np.ndarray): 33 | new_seg = np.zeros_like(seg) 34 | new_seg[seg == 1] = 2 35 | new_seg[seg == 3] = 4 36 | new_seg[seg == 2] = 1 37 | return new_seg 38 | 39 | 40 | def load_convert_labels_back_to_BraTS(filename, input_folder, output_folder): 41 | a = sitk.ReadImage(join(input_folder, filename)) 42 | b = sitk.GetArrayFromImage(a) 43 | c = convert_labels_back_to_BraTS(b) 44 | d = sitk.GetImageFromArray(c) 45 | d.CopyInformation(a) 46 | sitk.WriteImage(d, join(output_folder, filename)) 47 | 48 | 49 | def convert_folder_with_preds_back_to_BraTS_labeling_convention(input_folder: str, output_folder: str, num_processes: int = 12): 50 | """ 51 | reads all prediction files (nifti) in the input folder, converts the labels back to BraTS convention and saves the 52 | """ 53 | maybe_mkdir_p(output_folder) 54 | nii = subfiles(input_folder, suffix='.nii.gz', join=False) 55 | with multiprocessing.get_context("spawn").Pool(num_processes) as p: 56 | p.starmap(load_convert_labels_back_to_BraTS, zip(nii, [input_folder] * len(nii), [output_folder] * len(nii))) 57 | 58 | 59 | if __name__ == '__main__': 60 | brats_data_dir = '/home/isensee/drives/E132-Rohdaten/BraTS_2021/training' 61 | 62 | task_id = 137 63 | task_name = "BraTS2021" 64 | 65 | foldername = "Dataset%03.0d_%s" % (task_id, task_name) 66 | 67 | # setting up nnU-Net folders 68 | out_base = join(nnUNet_raw, foldername) 69 | imagestr = join(out_base, "imagesTr") 70 | labelstr = join(out_base, "labelsTr") 71 | maybe_mkdir_p(imagestr) 72 | maybe_mkdir_p(labelstr) 73 | 74 | case_ids = subdirs(brats_data_dir, prefix='BraTS', join=False) 75 | 76 | for c in case_ids: 77 | shutil.copy(join(brats_data_dir, c, c + "_t1.nii.gz"), join(imagestr, c + '_0000.nii.gz')) 78 | shutil.copy(join(brats_data_dir, c, c + "_t1ce.nii.gz"), join(imagestr, c + '_0001.nii.gz')) 79 | shutil.copy(join(brats_data_dir, c, c + "_t2.nii.gz"), join(imagestr, c + '_0002.nii.gz')) 80 | shutil.copy(join(brats_data_dir, c, c + "_flair.nii.gz"), join(imagestr, c + '_0003.nii.gz')) 81 | 82 | copy_BraTS_segmentation_and_convert_labels_to_nnUNet(join(brats_data_dir, c, c + "_seg.nii.gz"), 83 | join(labelstr, c + '.nii.gz')) 84 | 85 | generate_dataset_json(out_base, 86 | channel_names={0: 'T1', 1: 'T1ce', 2: 'T2', 3: 'Flair'}, 87 | labels={ 88 | 'background': 0, 89 | 'whole tumor': (1, 2, 3), 90 | 'tumor core': (2, 3), 91 | 'enhancing tumor': (3, ) 92 | }, 93 | num_training_cases=len(case_ids), 94 | file_ending='.nii.gz', 95 | regions_class_order=(1, 2, 3), 96 | license='see https://www.synapse.org/#!Synapse:syn25829067/wiki/610863', 97 | reference='see https://www.synapse.org/#!Synapse:syn25829067/wiki/610863', 98 | dataset_release='1.0') 99 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/Dataset218_Amos2022_task1.py: -------------------------------------------------------------------------------- 1 | from batchgenerators.utilities.file_and_folder_operations import * 2 | import shutil 3 | from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json 4 | from nnunetv2.paths import nnUNet_raw 5 | 6 | 7 | def convert_amos_task1(amos_base_dir: str, nnunet_dataset_id: int = 218): 8 | """ 9 | AMOS doesn't say anything about how the validation set is supposed to be used. So we just incorporate that into 10 | the train set. Having a 5-fold cross-validation is superior to a single train:val split 11 | """ 12 | task_name = "AMOS2022_postChallenge_task1" 13 | 14 | foldername = "Dataset%03.0d_%s" % (nnunet_dataset_id, task_name) 15 | 16 | # setting up nnU-Net folders 17 | out_base = join(nnUNet_raw, foldername) 18 | imagestr = join(out_base, "imagesTr") 19 | imagests = join(out_base, "imagesTs") 20 | labelstr = join(out_base, "labelsTr") 21 | maybe_mkdir_p(imagestr) 22 | maybe_mkdir_p(imagests) 23 | maybe_mkdir_p(labelstr) 24 | 25 | dataset_json_source = load_json(join(amos_base_dir, 'dataset.json')) 26 | 27 | training_identifiers = [i['image'].split('/')[-1][:-7] for i in dataset_json_source['training']] 28 | tr_ctr = 0 29 | for tr in training_identifiers: 30 | if int(tr.split("_")[-1]) <= 410: # these are the CT images 31 | tr_ctr += 1 32 | shutil.copy(join(amos_base_dir, 'imagesTr', tr + '.nii.gz'), join(imagestr, f'{tr}_0000.nii.gz')) 33 | shutil.copy(join(amos_base_dir, 'labelsTr', tr + '.nii.gz'), join(labelstr, f'{tr}.nii.gz')) 34 | 35 | test_identifiers = [i['image'].split('/')[-1][:-7] for i in dataset_json_source['test']] 36 | for ts in test_identifiers: 37 | if int(ts.split("_")[-1]) <= 500: # these are the CT images 38 | shutil.copy(join(amos_base_dir, 'imagesTs', ts + '.nii.gz'), join(imagests, f'{ts}_0000.nii.gz')) 39 | 40 | val_identifiers = [i['image'].split('/')[-1][:-7] for i in dataset_json_source['validation']] 41 | for vl in val_identifiers: 42 | if int(vl.split("_")[-1]) <= 409: # these are the CT images 43 | tr_ctr += 1 44 | shutil.copy(join(amos_base_dir, 'imagesVa', vl + '.nii.gz'), join(imagestr, f'{vl}_0000.nii.gz')) 45 | shutil.copy(join(amos_base_dir, 'labelsVa', vl + '.nii.gz'), join(labelstr, f'{vl}.nii.gz')) 46 | 47 | generate_dataset_json(out_base, {0: "CT"}, labels={v: int(k) for k,v in dataset_json_source['labels'].items()}, 48 | num_training_cases=tr_ctr, file_ending='.nii.gz', 49 | dataset_name=task_name, reference='https://amos22.grand-challenge.org/', 50 | release='https://zenodo.org/record/7262581', 51 | overwrite_image_reader_writer='NibabelIOWithReorient', 52 | description="This is the dataset as released AFTER the challenge event. It has the " 53 | "validation set gt in it! We just use the validation images as additional " 54 | "training cases because AMOS doesn't specify how they should be used. nnU-Net's" 55 | " 5-fold CV is better than some random train:val split.") 56 | 57 | 58 | if __name__ == '__main__': 59 | import argparse 60 | parser = argparse.ArgumentParser() 61 | parser.add_argument('input_folder', type=str, 62 | help="The downloaded and extracted AMOS2022 (https://amos22.grand-challenge.org/) data. " 63 | "Use this link: https://zenodo.org/record/7262581." 64 | "You need to specify the folder with the imagesTr, imagesVal, labelsTr etc subfolders here!") 65 | parser.add_argument('-d', required=False, type=int, default=218, help='nnU-Net Dataset ID, default: 218') 66 | args = parser.parse_args() 67 | amos_base = args.input_folder 68 | convert_amos_task1(amos_base, args.d) 69 | 70 | 71 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/Dataset219_Amos2022_task2.py: -------------------------------------------------------------------------------- 1 | from batchgenerators.utilities.file_and_folder_operations import * 2 | import shutil 3 | from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json 4 | from nnunetv2.paths import nnUNet_raw 5 | 6 | 7 | def convert_amos_task2(amos_base_dir: str, nnunet_dataset_id: int = 219): 8 | """ 9 | AMOS doesn't say anything about how the validation set is supposed to be used. So we just incorporate that into 10 | the train set. Having a 5-fold cross-validation is superior to a single train:val split 11 | """ 12 | task_name = "AMOS2022_postChallenge_task2" 13 | 14 | foldername = "Dataset%03.0d_%s" % (nnunet_dataset_id, task_name) 15 | 16 | # setting up nnU-Net folders 17 | out_base = join(nnUNet_raw, foldername) 18 | imagestr = join(out_base, "imagesTr") 19 | imagests = join(out_base, "imagesTs") 20 | labelstr = join(out_base, "labelsTr") 21 | maybe_mkdir_p(imagestr) 22 | maybe_mkdir_p(imagests) 23 | maybe_mkdir_p(labelstr) 24 | 25 | dataset_json_source = load_json(join(amos_base_dir, 'dataset.json')) 26 | 27 | training_identifiers = [i['image'].split('/')[-1][:-7] for i in dataset_json_source['training']] 28 | for tr in training_identifiers: 29 | shutil.copy(join(amos_base_dir, 'imagesTr', tr + '.nii.gz'), join(imagestr, f'{tr}_0000.nii.gz')) 30 | shutil.copy(join(amos_base_dir, 'labelsTr', tr + '.nii.gz'), join(labelstr, f'{tr}.nii.gz')) 31 | 32 | test_identifiers = [i['image'].split('/')[-1][:-7] for i in dataset_json_source['test']] 33 | for ts in test_identifiers: 34 | shutil.copy(join(amos_base_dir, 'imagesTs', ts + '.nii.gz'), join(imagests, f'{ts}_0000.nii.gz')) 35 | 36 | val_identifiers = [i['image'].split('/')[-1][:-7] for i in dataset_json_source['validation']] 37 | for vl in val_identifiers: 38 | shutil.copy(join(amos_base_dir, 'imagesVa', vl + '.nii.gz'), join(imagestr, f'{vl}_0000.nii.gz')) 39 | shutil.copy(join(amos_base_dir, 'labelsVa', vl + '.nii.gz'), join(labelstr, f'{vl}.nii.gz')) 40 | 41 | generate_dataset_json(out_base, {0: "either_CT_or_MR"}, labels={v: int(k) for k,v in dataset_json_source['labels'].items()}, 42 | num_training_cases=len(training_identifiers) + len(val_identifiers), file_ending='.nii.gz', 43 | dataset_name=task_name, reference='https://amos22.grand-challenge.org/', 44 | release='https://zenodo.org/record/7262581', 45 | overwrite_image_reader_writer='NibabelIOWithReorient', 46 | description="This is the dataset as released AFTER the challenge event. It has the " 47 | "validation set gt in it! We just use the validation images as additional " 48 | "training cases because AMOS doesn't specify how they should be used. nnU-Net's" 49 | " 5-fold CV is better than some random train:val split.") 50 | 51 | 52 | if __name__ == '__main__': 53 | import argparse 54 | parser = argparse.ArgumentParser() 55 | parser.add_argument('input_folder', type=str, 56 | help="The downloaded and extracted AMOS2022 (https://amos22.grand-challenge.org/) data. " 57 | "Use this link: https://zenodo.org/record/7262581." 58 | "You need to specify the folder with the imagesTr, imagesVal, labelsTr etc subfolders here!") 59 | parser.add_argument('-d', required=False, type=int, default=219, help='nnU-Net Dataset ID, default: 219') 60 | args = parser.parse_args() 61 | amos_base = args.input_folder 62 | convert_amos_task2(amos_base, args.d) 63 | 64 | # /home/isensee/Downloads/amos22/amos22/ 65 | 66 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/dataset_conversion/__init__.py -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/convert_raw_dataset_from_old_nnunet_format.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | from copy import deepcopy 3 | 4 | from batchgenerators.utilities.file_and_folder_operations import join, maybe_mkdir_p, isdir, load_json, save_json 5 | from nnunetv2.paths import nnUNet_raw 6 | 7 | 8 | def convert(source_folder, target_dataset_name): 9 | """ 10 | remember that old tasks were called TaskXXX_YYY and new ones are called DatasetXXX_YYY 11 | source_folder 12 | """ 13 | if isdir(join(nnUNet_raw, target_dataset_name)): 14 | raise RuntimeError(f'Target dataset name {target_dataset_name} already exists. Aborting... ' 15 | f'(we might break something). If you are sure you want to proceed, please manually ' 16 | f'delete {join(nnUNet_raw, target_dataset_name)}') 17 | maybe_mkdir_p(join(nnUNet_raw, target_dataset_name)) 18 | shutil.copytree(join(source_folder, 'imagesTr'), join(nnUNet_raw, target_dataset_name, 'imagesTr')) 19 | shutil.copytree(join(source_folder, 'labelsTr'), join(nnUNet_raw, target_dataset_name, 'labelsTr')) 20 | if isdir(join(source_folder, 'imagesTs')): 21 | shutil.copytree(join(source_folder, 'imagesTs'), join(nnUNet_raw, target_dataset_name, 'imagesTs')) 22 | if isdir(join(source_folder, 'labelsTs')): 23 | shutil.copytree(join(source_folder, 'labelsTs'), join(nnUNet_raw, target_dataset_name, 'labelsTs')) 24 | if isdir(join(source_folder, 'imagesVal')): 25 | shutil.copytree(join(source_folder, 'imagesVal'), join(nnUNet_raw, target_dataset_name, 'imagesVal')) 26 | if isdir(join(source_folder, 'labelsVal')): 27 | shutil.copytree(join(source_folder, 'labelsVal'), join(nnUNet_raw, target_dataset_name, 'labelsVal')) 28 | shutil.copy(join(source_folder, 'dataset.json'), join(nnUNet_raw, target_dataset_name)) 29 | 30 | dataset_json = load_json(join(nnUNet_raw, target_dataset_name, 'dataset.json')) 31 | del dataset_json['tensorImageSize'] 32 | del dataset_json['numTest'] 33 | del dataset_json['training'] 34 | del dataset_json['test'] 35 | dataset_json['channel_names'] = deepcopy(dataset_json['modality']) 36 | del dataset_json['modality'] 37 | 38 | dataset_json['labels'] = {j: int(i) for i, j in dataset_json['labels'].items()} 39 | dataset_json['file_ending'] = ".nii.gz" 40 | save_json(dataset_json, join(nnUNet_raw, target_dataset_name, 'dataset.json'), sort_keys=False) 41 | 42 | 43 | def convert_entry_point(): 44 | import argparse 45 | parser = argparse.ArgumentParser() 46 | parser.add_argument("input_folder", type=str, 47 | help='Raw old nnUNet dataset. This must be the folder with imagesTr,labelsTr etc subfolders! ' 48 | 'Please provide the PATH to the old Task, not just the task name. nnU-Net V2 does not ' 49 | 'know where v1 tasks are.') 50 | parser.add_argument("output_dataset_name", type=str, 51 | help='New dataset NAME (not path!). Must follow the DatasetXXX_NAME convention!') 52 | args = parser.parse_args() 53 | convert(args.input_folder, args.output_dataset_name) 54 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/datasets_for_integration_tests/Dataset996_IntegrationTest_Hippocampus_regions_ignore.py: -------------------------------------------------------------------------------- 1 | import SimpleITK as sitk 2 | import shutil 3 | 4 | import numpy as np 5 | from batchgenerators.utilities.file_and_folder_operations import isdir, join, load_json, save_json, nifti_files 6 | 7 | from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name 8 | from nnunetv2.paths import nnUNet_raw 9 | from nnunetv2.utilities.label_handling.label_handling import LabelManager 10 | from nnunetv2.utilities.plans_handling.plans_handler import PlansManager, ConfigurationManager 11 | 12 | 13 | def sparsify_segmentation(seg: np.ndarray, label_manager: LabelManager, percent_of_slices: float) -> np.ndarray: 14 | assert label_manager.has_ignore_label, "This preprocessor only works with datasets that have an ignore label!" 15 | seg_new = np.ones_like(seg) * label_manager.ignore_label 16 | x, y, z = seg.shape 17 | # x 18 | num_slices = max(1, round(x * percent_of_slices)) 19 | selected_slices = np.random.choice(x, num_slices, replace=False) 20 | seg_new[selected_slices] = seg[selected_slices] 21 | # y 22 | num_slices = max(1, round(y * percent_of_slices)) 23 | selected_slices = np.random.choice(y, num_slices, replace=False) 24 | seg_new[:, selected_slices] = seg[:, selected_slices] 25 | # z 26 | num_slices = max(1, round(z * percent_of_slices)) 27 | selected_slices = np.random.choice(z, num_slices, replace=False) 28 | seg_new[:, :, selected_slices] = seg[:, :, selected_slices] 29 | return seg_new 30 | 31 | 32 | if __name__ == '__main__': 33 | dataset_name = 'IntegrationTest_Hippocampus_regions_ignore' 34 | dataset_id = 996 35 | dataset_name = f"Dataset{dataset_id:03d}_{dataset_name}" 36 | 37 | try: 38 | existing_dataset_name = maybe_convert_to_dataset_name(dataset_id) 39 | if existing_dataset_name != dataset_name: 40 | raise FileExistsError(f"A different dataset with id {dataset_id} already exists :-(: {existing_dataset_name}. If " 41 | f"you intent to delete it, remember to also remove it in nnUNet_preprocessed and " 42 | f"nnUNet_results!") 43 | except RuntimeError: 44 | pass 45 | 46 | if isdir(join(nnUNet_raw, dataset_name)): 47 | shutil.rmtree(join(nnUNet_raw, dataset_name)) 48 | 49 | source_dataset = maybe_convert_to_dataset_name(4) 50 | shutil.copytree(join(nnUNet_raw, source_dataset), join(nnUNet_raw, dataset_name)) 51 | 52 | # additionally optimize entire hippocampus region, remove Posterior 53 | dj = load_json(join(nnUNet_raw, dataset_name, 'dataset.json')) 54 | dj['labels'] = { 55 | 'background': 0, 56 | 'hippocampus': (1, 2), 57 | 'anterior': 1, 58 | 'ignore': 3 59 | } 60 | dj['regions_class_order'] = (2, 1) 61 | save_json(dj, join(nnUNet_raw, dataset_name, 'dataset.json'), sort_keys=False) 62 | 63 | # now add ignore label to segmentation images 64 | np.random.seed(1234) 65 | lm = LabelManager(label_dict=dj['labels'], regions_class_order=dj.get('regions_class_order')) 66 | 67 | segs = nifti_files(join(nnUNet_raw, dataset_name, 'labelsTr')) 68 | for s in segs: 69 | seg_itk = sitk.ReadImage(s) 70 | seg_npy = sitk.GetArrayFromImage(seg_itk) 71 | seg_npy = sparsify_segmentation(seg_npy, lm, 0.1 / 3) 72 | seg_itk_new = sitk.GetImageFromArray(seg_npy) 73 | seg_itk_new.CopyInformation(seg_itk) 74 | sitk.WriteImage(seg_itk_new, s) 75 | 76 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/datasets_for_integration_tests/Dataset997_IntegrationTest_Hippocampus_regions.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | 3 | from batchgenerators.utilities.file_and_folder_operations import isdir, join, load_json, save_json 4 | 5 | from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name 6 | from nnunetv2.paths import nnUNet_raw 7 | 8 | if __name__ == '__main__': 9 | dataset_name = 'IntegrationTest_Hippocampus_regions' 10 | dataset_id = 997 11 | dataset_name = f"Dataset{dataset_id:03d}_{dataset_name}" 12 | 13 | try: 14 | existing_dataset_name = maybe_convert_to_dataset_name(dataset_id) 15 | if existing_dataset_name != dataset_name: 16 | raise FileExistsError( 17 | f"A different dataset with id {dataset_id} already exists :-(: {existing_dataset_name}. If " 18 | f"you intent to delete it, remember to also remove it in nnUNet_preprocessed and " 19 | f"nnUNet_results!") 20 | except RuntimeError: 21 | pass 22 | 23 | if isdir(join(nnUNet_raw, dataset_name)): 24 | shutil.rmtree(join(nnUNet_raw, dataset_name)) 25 | 26 | source_dataset = maybe_convert_to_dataset_name(4) 27 | shutil.copytree(join(nnUNet_raw, source_dataset), join(nnUNet_raw, dataset_name)) 28 | 29 | # additionally optimize entire hippocampus region, remove Posterior 30 | dj = load_json(join(nnUNet_raw, dataset_name, 'dataset.json')) 31 | dj['labels'] = { 32 | 'background': 0, 33 | 'hippocampus': (1, 2), 34 | 'anterior': 1 35 | } 36 | dj['regions_class_order'] = (2, 1) 37 | save_json(dj, join(nnUNet_raw, dataset_name, 'dataset.json'), sort_keys=False) 38 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/datasets_for_integration_tests/Dataset998_IntegrationTest_Hippocampus_ignore.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | 3 | from batchgenerators.utilities.file_and_folder_operations import isdir, join, load_json, save_json 4 | 5 | from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name 6 | from nnunetv2.paths import nnUNet_raw 7 | 8 | 9 | if __name__ == '__main__': 10 | dataset_name = 'IntegrationTest_Hippocampus_ignore' 11 | dataset_id = 998 12 | dataset_name = f"Dataset{dataset_id:03d}_{dataset_name}" 13 | 14 | try: 15 | existing_dataset_name = maybe_convert_to_dataset_name(dataset_id) 16 | if existing_dataset_name != dataset_name: 17 | raise FileExistsError(f"A different dataset with id {dataset_id} already exists :-(: {existing_dataset_name}. If " 18 | f"you intent to delete it, remember to also remove it in nnUNet_preprocessed and " 19 | f"nnUNet_results!") 20 | except RuntimeError: 21 | pass 22 | 23 | if isdir(join(nnUNet_raw, dataset_name)): 24 | shutil.rmtree(join(nnUNet_raw, dataset_name)) 25 | 26 | source_dataset = maybe_convert_to_dataset_name(4) 27 | shutil.copytree(join(nnUNet_raw, source_dataset), join(nnUNet_raw, dataset_name)) 28 | 29 | # set class 2 to ignore label 30 | dj = load_json(join(nnUNet_raw, dataset_name, 'dataset.json')) 31 | dj['labels']['ignore'] = 2 32 | del dj['labels']['Posterior'] 33 | save_json(dj, join(nnUNet_raw, dataset_name, 'dataset.json'), sort_keys=False) 34 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/datasets_for_integration_tests/Dataset999_IntegrationTest_Hippocampus.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | 3 | from batchgenerators.utilities.file_and_folder_operations import isdir, join 4 | 5 | from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name 6 | from nnunetv2.paths import nnUNet_raw 7 | 8 | 9 | if __name__ == '__main__': 10 | dataset_name = 'IntegrationTest_Hippocampus' 11 | dataset_id = 999 12 | dataset_name = f"Dataset{dataset_id:03d}_{dataset_name}" 13 | 14 | try: 15 | existing_dataset_name = maybe_convert_to_dataset_name(dataset_id) 16 | if existing_dataset_name != dataset_name: 17 | raise FileExistsError(f"A different dataset with id {dataset_id} already exists :-(: {existing_dataset_name}. If " 18 | f"you intent to delete it, remember to also remove it in nnUNet_preprocessed and " 19 | f"nnUNet_results!") 20 | except RuntimeError: 21 | pass 22 | 23 | if isdir(join(nnUNet_raw, dataset_name)): 24 | shutil.rmtree(join(nnUNet_raw, dataset_name)) 25 | 26 | source_dataset = maybe_convert_to_dataset_name(4) 27 | shutil.copytree(join(nnUNet_raw, source_dataset), join(nnUNet_raw, dataset_name)) 28 | -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/datasets_for_integration_tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/dataset_conversion/datasets_for_integration_tests/__init__.py -------------------------------------------------------------------------------- /nnunetv2/dataset_conversion/generate_dataset_json.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | 3 | from batchgenerators.utilities.file_and_folder_operations import save_json, join 4 | 5 | 6 | def generate_dataset_json(output_folder: str, 7 | channel_names: dict, 8 | labels: dict, 9 | num_training_cases: int, 10 | file_ending: str, 11 | regions_class_order: Tuple[int, ...] = None, 12 | dataset_name: str = None, reference: str = None, release: str = None, license: str = None, 13 | description: str = None, 14 | overwrite_image_reader_writer: str = None, **kwargs): 15 | """ 16 | Generates a dataset.json file in the output folder 17 | 18 | channel_names: 19 | Channel names must map the index to the name of the channel, example: 20 | { 21 | 0: 'T1', 22 | 1: 'CT' 23 | } 24 | Note that the channel names may influence the normalization scheme!! Learn more in the documentation. 25 | 26 | labels: 27 | This will tell nnU-Net what labels to expect. Important: This will also determine whether you use region-based training or not. 28 | Example regular labels: 29 | { 30 | 'background': 0, 31 | 'left atrium': 1, 32 | 'some other label': 2 33 | } 34 | Example region-based training: 35 | { 36 | 'background': 0, 37 | 'whole tumor': (1, 2, 3), 38 | 'tumor core': (2, 3), 39 | 'enhancing tumor': 3 40 | } 41 | 42 | Remember that nnU-Net expects consecutive values for labels! nnU-Net also expects 0 to be background! 43 | 44 | num_training_cases: is used to double check all cases are there! 45 | 46 | file_ending: needed for finding the files correctly. IMPORTANT! File endings must match between images and 47 | segmentations! 48 | 49 | dataset_name, reference, release, license, description: self-explanatory and not used by nnU-Net. Just for 50 | completeness and as a reminder that these would be great! 51 | 52 | overwrite_image_reader_writer: If you need a special IO class for your dataset you can derive it from 53 | BaseReaderWriter, place it into nnunet.imageio and reference it here by name 54 | 55 | kwargs: whatever you put here will be placed in the dataset.json as well 56 | 57 | """ 58 | has_regions: bool = any([isinstance(i, (tuple, list)) and len(i) > 1 for i in labels.values()]) 59 | if has_regions: 60 | assert regions_class_order is not None, f"You have defined regions but regions_class_order is not set. " \ 61 | f"You need that." 62 | # channel names need strings as keys 63 | keys = list(channel_names.keys()) 64 | for k in keys: 65 | if not isinstance(k, str): 66 | channel_names[str(k)] = channel_names[k] 67 | del channel_names[k] 68 | 69 | # labels need ints as values 70 | for l in labels.keys(): 71 | value = labels[l] 72 | if isinstance(value, (tuple, list)): 73 | value = tuple([int(i) for i in value]) 74 | labels[l] = value 75 | else: 76 | labels[l] = int(labels[l]) 77 | 78 | dataset_json = { 79 | 'channel_names': channel_names, # previously this was called 'modality'. I didnt like this so this is 80 | # channel_names now. Live with it. 81 | 'labels': labels, 82 | 'numTraining': num_training_cases, 83 | 'file_ending': file_ending, 84 | } 85 | 86 | if dataset_name is not None: 87 | dataset_json['name'] = dataset_name 88 | if reference is not None: 89 | dataset_json['reference'] = reference 90 | if release is not None: 91 | dataset_json['release'] = release 92 | if license is not None: 93 | dataset_json['licence'] = license 94 | if description is not None: 95 | dataset_json['description'] = description 96 | if overwrite_image_reader_writer is not None: 97 | dataset_json['overwrite_image_reader_writer'] = overwrite_image_reader_writer 98 | if regions_class_order is not None: 99 | dataset_json['regions_class_order'] = regions_class_order 100 | 101 | dataset_json.update(kwargs) 102 | 103 | save_json(dataset_json, join(output_folder, 'dataset.json'), sort_keys=False) 104 | -------------------------------------------------------------------------------- /nnunetv2/ensembling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/ensembling/__init__.py -------------------------------------------------------------------------------- /nnunetv2/evaluation/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/evaluation/__init__.py -------------------------------------------------------------------------------- /nnunetv2/evaluation/accumulate_cv_results.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | from typing import Union, List, Tuple 3 | 4 | from batchgenerators.utilities.file_and_folder_operations import load_json, join, isdir, maybe_mkdir_p, subfiles, isfile 5 | 6 | from nnunetv2.configuration import default_num_processes 7 | from nnunetv2.evaluation.evaluate_predictions import compute_metrics_on_folder 8 | from nnunetv2.paths import nnUNet_raw 9 | from nnunetv2.utilities.plans_handling.plans_handler import PlansManager 10 | 11 | 12 | def accumulate_cv_results(trained_model_folder, 13 | merged_output_folder: str, 14 | folds: Union[List[int], Tuple[int, ...]], 15 | num_processes: int = default_num_processes, 16 | overwrite: bool = True): 17 | """ 18 | There are a lot of things that can get fucked up, so the simplest way to deal with potential problems is to 19 | collect the cv results into a separate folder and then evaluate them again. No messing with summary_json files! 20 | """ 21 | 22 | if overwrite and isdir(merged_output_folder): 23 | shutil.rmtree(merged_output_folder) 24 | maybe_mkdir_p(merged_output_folder) 25 | 26 | dataset_json = load_json(join(trained_model_folder, 'dataset.json')) 27 | plans_manager = PlansManager(join(trained_model_folder, 'plans.json')) 28 | rw = plans_manager.image_reader_writer_class() 29 | shutil.copy(join(trained_model_folder, 'dataset.json'), join(merged_output_folder, 'dataset.json')) 30 | shutil.copy(join(trained_model_folder, 'plans.json'), join(merged_output_folder, 'plans.json')) 31 | 32 | did_we_copy_something = False 33 | for f in folds: 34 | expected_validation_folder = join(trained_model_folder, f'fold_{f}', 'validation') 35 | if not isdir(expected_validation_folder): 36 | raise RuntimeError(f"fold {f} of model {trained_model_folder} is missing. Please train it!") 37 | predicted_files = subfiles(expected_validation_folder, suffix=dataset_json['file_ending'], join=False) 38 | for pf in predicted_files: 39 | if overwrite and isfile(join(merged_output_folder, pf)): 40 | raise RuntimeError(f'More than one of your folds has a prediction for case {pf}') 41 | if overwrite or not isfile(join(merged_output_folder, pf)): 42 | shutil.copy(join(expected_validation_folder, pf), join(merged_output_folder, pf)) 43 | did_we_copy_something = True 44 | 45 | if did_we_copy_something or not isfile(join(merged_output_folder, 'summary.json')): 46 | label_manager = plans_manager.get_label_manager(dataset_json) 47 | compute_metrics_on_folder(join(nnUNet_raw, plans_manager.dataset_name, 'labelsTr'), 48 | merged_output_folder, 49 | join(merged_output_folder, 'summary.json'), 50 | rw, 51 | dataset_json['file_ending'], 52 | label_manager.foreground_regions if label_manager.has_regions else 53 | label_manager.foreground_labels, 54 | label_manager.ignore_label, 55 | num_processes) 56 | -------------------------------------------------------------------------------- /nnunetv2/experiment_planning/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/experiment_planning/__init__.py -------------------------------------------------------------------------------- /nnunetv2/experiment_planning/dataset_fingerprint/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/experiment_planning/dataset_fingerprint/__init__.py -------------------------------------------------------------------------------- /nnunetv2/experiment_planning/experiment_planners/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/experiment_planning/experiment_planners/__init__.py -------------------------------------------------------------------------------- /nnunetv2/experiment_planning/experiment_planners/network_topology.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | import numpy as np 3 | 4 | 5 | def get_shape_must_be_divisible_by(net_numpool_per_axis): 6 | return 2 ** np.array(net_numpool_per_axis) 7 | 8 | 9 | def pad_shape(shape, must_be_divisible_by): 10 | """ 11 | pads shape so that it is divisible by must_be_divisible_by 12 | :param shape: 13 | :param must_be_divisible_by: 14 | :return: 15 | """ 16 | if not isinstance(must_be_divisible_by, (tuple, list, np.ndarray)): 17 | must_be_divisible_by = [must_be_divisible_by] * len(shape) 18 | else: 19 | assert len(must_be_divisible_by) == len(shape) 20 | 21 | new_shp = [shape[i] + must_be_divisible_by[i] - shape[i] % must_be_divisible_by[i] for i in range(len(shape))] 22 | 23 | for i in range(len(shape)): 24 | if shape[i] % must_be_divisible_by[i] == 0: 25 | new_shp[i] -= must_be_divisible_by[i] 26 | new_shp = np.array(new_shp).astype(int) 27 | return new_shp 28 | 29 | 30 | def get_pool_and_conv_props(spacing, patch_size, min_feature_map_size, max_numpool): 31 | """ 32 | this is the same as get_pool_and_conv_props_v2 from old nnunet 33 | 34 | :param spacing: 35 | :param patch_size: 36 | :param min_feature_map_size: min edge length of feature maps in bottleneck 37 | :param max_numpool: 38 | :return: 39 | """ 40 | # todo review this code 41 | dim = len(spacing) 42 | 43 | current_spacing = deepcopy(list(spacing)) 44 | current_size = deepcopy(list(patch_size)) 45 | 46 | pool_op_kernel_sizes = [[1] * len(spacing)] 47 | conv_kernel_sizes = [] 48 | 49 | num_pool_per_axis = [0] * dim 50 | kernel_size = [1] * dim 51 | 52 | while True: 53 | # exclude axes that we cannot pool further because of min_feature_map_size constraint 54 | valid_axes_for_pool = [i for i in range(dim) if current_size[i] >= 2*min_feature_map_size] 55 | if len(valid_axes_for_pool) < 1: 56 | break 57 | 58 | spacings_of_axes = [current_spacing[i] for i in valid_axes_for_pool] 59 | 60 | # find axis that are within factor of 2 within smallest spacing 61 | min_spacing_of_valid = min(spacings_of_axes) 62 | valid_axes_for_pool = [i for i in valid_axes_for_pool if current_spacing[i] / min_spacing_of_valid < 2] 63 | 64 | # max_numpool constraint 65 | valid_axes_for_pool = [i for i in valid_axes_for_pool if num_pool_per_axis[i] < max_numpool] 66 | 67 | if len(valid_axes_for_pool) == 1: 68 | if current_size[valid_axes_for_pool[0]] >= 3 * min_feature_map_size: 69 | pass 70 | else: 71 | break 72 | if len(valid_axes_for_pool) < 1: 73 | break 74 | 75 | # now we need to find kernel sizes 76 | # kernel sizes are initialized to 1. They are successively set to 3 when their associated axis becomes within 77 | # factor 2 of min_spacing. Once they are 3 they remain 3 78 | for d in range(dim): 79 | if kernel_size[d] == 3: 80 | continue 81 | else: 82 | if spacings_of_axes[d] / min(current_spacing) < 2: 83 | kernel_size[d] = 3 84 | 85 | other_axes = [i for i in range(dim) if i not in valid_axes_for_pool] 86 | 87 | pool_kernel_sizes = [0] * dim 88 | for v in valid_axes_for_pool: 89 | pool_kernel_sizes[v] = 2 90 | num_pool_per_axis[v] += 1 91 | current_spacing[v] *= 2 92 | current_size[v] = np.ceil(current_size[v] / 2) 93 | for nv in other_axes: 94 | pool_kernel_sizes[nv] = 1 95 | 96 | pool_op_kernel_sizes.append(pool_kernel_sizes) 97 | conv_kernel_sizes.append(deepcopy(kernel_size)) 98 | #print(conv_kernel_sizes) 99 | 100 | must_be_divisible_by = get_shape_must_be_divisible_by(num_pool_per_axis) 101 | patch_size = pad_shape(patch_size, must_be_divisible_by) 102 | 103 | # we need to add one more conv_kernel_size for the bottleneck. We always use 3x3(x3) conv here 104 | conv_kernel_sizes.append([3]*dim) 105 | return num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, patch_size, must_be_divisible_by 106 | -------------------------------------------------------------------------------- /nnunetv2/experiment_planning/experiment_planners/resencUNet_planner.py: -------------------------------------------------------------------------------- 1 | from typing import Union, List, Tuple 2 | 3 | from torch import nn 4 | 5 | from nnunetv2.experiment_planning.experiment_planners.default_experiment_planner import ExperimentPlanner 6 | from dynamic_network_architectures.architectures.unet import ResidualEncoderUNet 7 | 8 | 9 | class ResEncUNetPlanner(ExperimentPlanner): 10 | def __init__(self, dataset_name_or_id: Union[str, int], 11 | gpu_memory_target_in_gb: float = 8, 12 | preprocessor_name: str = 'DefaultPreprocessor', plans_name: str = 'nnUNetResEncUNetPlans', 13 | overwrite_target_spacing: Union[List[float], Tuple[float, ...]] = None, 14 | suppress_transpose: bool = False): 15 | super().__init__(dataset_name_or_id, gpu_memory_target_in_gb, preprocessor_name, plans_name, 16 | overwrite_target_spacing, suppress_transpose) 17 | 18 | self.UNet_base_num_features = 32 19 | self.UNet_class = ResidualEncoderUNet 20 | # the following two numbers are really arbitrary and were set to reproduce default nnU-Net's configurations as 21 | # much as possible 22 | self.UNet_reference_val_3d = 680000000 23 | self.UNet_reference_val_2d = 135000000 24 | self.UNet_reference_com_nfeatures = 32 25 | self.UNet_reference_val_corresp_GB = 8 26 | self.UNet_reference_val_corresp_bs_2d = 12 27 | self.UNet_reference_val_corresp_bs_3d = 2 28 | self.UNet_featuremap_min_edge_length = 4 29 | self.UNet_blocks_per_stage_encoder = (1, 3, 4, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6) 30 | self.UNet_blocks_per_stage_decoder = (1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1) 31 | self.UNet_min_batch_size = 2 32 | self.UNet_max_features_2d = 512 33 | self.UNet_max_features_3d = 320 34 | 35 | 36 | if __name__ == '__main__': 37 | # we know both of these networks run with batch size 2 and 12 on ~8-10GB, respectively 38 | net = ResidualEncoderUNet(input_channels=1, n_stages=6, features_per_stage=(32, 64, 128, 256, 320, 320), 39 | conv_op=nn.Conv3d, kernel_sizes=3, strides=(1, 2, 2, 2, 2, 2), 40 | n_blocks_per_stage=(1, 3, 4, 6, 6, 6), num_classes=3, 41 | n_conv_per_stage_decoder=(1, 1, 1, 1, 1), 42 | conv_bias=True, norm_op=nn.InstanceNorm3d, norm_op_kwargs={}, dropout_op=None, 43 | nonlin=nn.LeakyReLU, nonlin_kwargs={'inplace': True}, deep_supervision=True) 44 | print(net.compute_conv_feature_map_size((128, 128, 128))) # -> 558319104. The value you see above was finetuned 45 | # from this one to match the regular nnunetplans more closely 46 | 47 | net = ResidualEncoderUNet(input_channels=1, n_stages=7, features_per_stage=(32, 64, 128, 256, 512, 512, 512), 48 | conv_op=nn.Conv2d, kernel_sizes=3, strides=(1, 2, 2, 2, 2, 2, 2), 49 | n_blocks_per_stage=(1, 3, 4, 6, 6, 6, 6), num_classes=3, 50 | n_conv_per_stage_decoder=(1, 1, 1, 1, 1, 1), 51 | conv_bias=True, norm_op=nn.InstanceNorm2d, norm_op_kwargs={}, dropout_op=None, 52 | nonlin=nn.LeakyReLU, nonlin_kwargs={'inplace': True}, deep_supervision=True) 53 | print(net.compute_conv_feature_map_size((512, 512))) # -> 129793792 54 | 55 | -------------------------------------------------------------------------------- /nnunetv2/experiment_planning/plans_for_pretraining/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/experiment_planning/plans_for_pretraining/__init__.py -------------------------------------------------------------------------------- /nnunetv2/experiment_planning/plans_for_pretraining/move_plans_between_datasets.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from typing import Union 3 | 4 | from batchgenerators.utilities.file_and_folder_operations import join, isdir, isfile, load_json, subfiles, save_json 5 | 6 | from nnunetv2.imageio.reader_writer_registry import determine_reader_writer_from_dataset_json 7 | from nnunetv2.paths import nnUNet_preprocessed, nnUNet_raw 8 | from nnunetv2.utilities.file_path_utilities import maybe_convert_to_dataset_name 9 | 10 | 11 | def move_plans_between_datasets( 12 | source_dataset_name_or_id: Union[int, str], 13 | target_dataset_name_or_id: Union[int, str], 14 | source_plans_identifier: str, 15 | target_plans_identifier: str = None): 16 | source_dataset_name = maybe_convert_to_dataset_name(source_dataset_name_or_id) 17 | target_dataset_name = maybe_convert_to_dataset_name(target_dataset_name_or_id) 18 | 19 | if target_plans_identifier is None: 20 | target_plans_identifier = source_plans_identifier 21 | 22 | source_folder = join(nnUNet_preprocessed, source_dataset_name) 23 | assert isdir(source_folder), f"Cannot move plans because preprocessed directory of source dataset is missing. " \ 24 | f"Run nnUNetv2_plan_and_preprocess for source dataset first!" 25 | 26 | source_plans_file = join(source_folder, source_plans_identifier + '.json') 27 | assert isfile(source_plans_file), f"Source plans are missing. Run the corresponding experiment planning first! " \ 28 | f"Expected file: {source_plans_file}" 29 | 30 | source_plans = load_json(source_plans_file) 31 | source_plans['dataset_name'] = target_dataset_name 32 | 33 | # we need to change data_identifier to use target_plans_identifier 34 | if target_plans_identifier != source_plans_identifier: 35 | for c in source_plans['configurations'].keys(): 36 | old_identifier = source_plans['configurations'][c]["data_identifier"] 37 | if old_identifier.startswith(source_plans_identifier): 38 | new_identifier = target_plans_identifier + old_identifier[len(source_plans_identifier):] 39 | else: 40 | new_identifier = target_plans_identifier + '_' + old_identifier 41 | source_plans['configurations'][c]["data_identifier"] = new_identifier 42 | 43 | # we need to change the reader writer class! 44 | target_raw_data_dir = join(nnUNet_raw, target_dataset_name) 45 | target_dataset_json = load_json(join(target_raw_data_dir, 'dataset.json')) 46 | file_ending = target_dataset_json['file_ending'] 47 | # pick any file from the imagesTr folder 48 | some_file = subfiles(join(target_raw_data_dir, 'imagesTr'), suffix=file_ending)[0] 49 | rw = determine_reader_writer_from_dataset_json(target_dataset_json, some_file, allow_nonmatching_filename=True, 50 | verbose=False) 51 | 52 | source_plans["image_reader_writer"] = rw.__name__ 53 | 54 | save_json(source_plans, join(nnUNet_preprocessed, target_dataset_name, target_plans_identifier + '.json'), 55 | sort_keys=False) 56 | 57 | 58 | def entry_point_move_plans_between_datasets(): 59 | parser = argparse.ArgumentParser() 60 | parser.add_argument('-s', type=str, required=True, 61 | help='Source dataset name or id') 62 | parser.add_argument('-t', type=str, required=True, 63 | help='Target dataset name or id') 64 | parser.add_argument('-sp', type=str, required=True, 65 | help='Source plans identifier. If your plans are named "nnUNetPlans.json" then the ' 66 | 'identifier would be nnUNetPlans') 67 | parser.add_argument('-tp', type=str, required=False, default=None, 68 | help='Target plans identifier. Default is None meaning the source plans identifier will ' 69 | 'be kept. Not recommended if the source plans identifier is a default nnU-Net identifier ' 70 | 'such as nnUNetPlans!!!') 71 | args = parser.parse_args() 72 | move_plans_between_datasets(args.s, args.t, args.sp, args.tp) 73 | 74 | 75 | if __name__ == '__main__': 76 | move_plans_between_datasets(2, 4, 'nnUNetPlans', 'nnUNetPlansFrom2') 77 | -------------------------------------------------------------------------------- /nnunetv2/imageio/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/imageio/__init__.py -------------------------------------------------------------------------------- /nnunetv2/imageio/natural_image_reager_writer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 HIP Applied Computer Vision Lab, Division of Medical Image Computing, German Cancer Research Center 2 | # (DKFZ), Heidelberg, Germany 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | from typing import Tuple, Union, List 17 | import numpy as np 18 | from nnunetv2.imageio.base_reader_writer import BaseReaderWriter 19 | from skimage import io 20 | 21 | 22 | class NaturalImage2DIO(BaseReaderWriter): 23 | """ 24 | ONLY SUPPORTS 2D IMAGES!!! 25 | """ 26 | 27 | # there are surely more we could add here. Everything that can be read by skimage.io should be supported 28 | supported_file_endings = [ 29 | '.png', 30 | # '.jpg', 31 | # '.jpeg', # jpg not supported because we cannot allow lossy compression! segmentation maps! 32 | '.bmp', 33 | '.tif' 34 | ] 35 | 36 | def read_images(self, image_fnames: Union[List[str], Tuple[str, ...]]) -> Tuple[np.ndarray, dict]: 37 | images = [] 38 | for f in image_fnames: 39 | npy_img = io.imread(f) 40 | if len(npy_img.shape) == 3: 41 | # rgb image, last dimension should be the color channel and the size of that channel should be 3 42 | # (or 4 if we have alpha) 43 | assert npy_img.shape[-1] == 3 or npy_img.shape[-1] == 4, "If image has three dimensions then the last " \ 44 | "dimension must have shape 3 or 4 " \ 45 | f"(RGB or RGBA). Image shape here is {npy_img.shape}" 46 | # move RGB(A) to front, add additional dim so that we have shape (1, c, X, Y), where c is either 3 or 4 47 | images.append(npy_img.transpose((2, 0, 1))[:, None]) 48 | elif len(npy_img.shape) == 2: 49 | # grayscale image 50 | images.append(npy_img[None, None]) 51 | 52 | if not self._check_all_same([i.shape for i in images]): 53 | print('ERROR! Not all input images have the same shape!') 54 | print('Shapes:') 55 | print([i.shape for i in images]) 56 | print('Image files:') 57 | print(image_fnames) 58 | raise RuntimeError() 59 | return np.vstack(images).astype(np.float32), {'spacing': (999, 1, 1)} 60 | 61 | def read_seg(self, seg_fname: str) -> Tuple[np.ndarray, dict]: 62 | return self.read_images((seg_fname, )) 63 | 64 | def write_seg(self, seg: np.ndarray, output_fname: str, properties: dict) -> None: 65 | io.imsave(output_fname, seg[0].astype(np.uint8), check_contrast=False) 66 | 67 | 68 | if __name__ == '__main__': 69 | images = ('/media/fabian/data/nnUNet_raw/Dataset120_RoadSegmentation/imagesTr/img-11_0000.png',) 70 | segmentation = '/media/fabian/data/nnUNet_raw/Dataset120_RoadSegmentation/labelsTr/img-11.png' 71 | imgio = NaturalImage2DIO() 72 | img, props = imgio.read_images(images) 73 | seg, segprops = imgio.read_seg(segmentation) -------------------------------------------------------------------------------- /nnunetv2/imageio/reader_writer_registry.py: -------------------------------------------------------------------------------- 1 | import traceback 2 | from typing import Type 3 | 4 | from batchgenerators.utilities.file_and_folder_operations import join 5 | 6 | import nnunetv2 7 | from nnunetv2.imageio.natural_image_reager_writer import NaturalImage2DIO 8 | from nnunetv2.imageio.nibabel_reader_writer import NibabelIO, NibabelIOWithReorient 9 | from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO 10 | from nnunetv2.imageio.tif_reader_writer import Tiff3DIO 11 | from nnunetv2.imageio.base_reader_writer import BaseReaderWriter 12 | from nnunetv2.utilities.find_class_by_name import recursive_find_python_class 13 | 14 | LIST_OF_IO_CLASSES = [ 15 | NaturalImage2DIO, 16 | SimpleITKIO, 17 | Tiff3DIO, 18 | NibabelIO, 19 | NibabelIOWithReorient 20 | ] 21 | 22 | 23 | def determine_reader_writer_from_dataset_json(dataset_json_content: dict, example_file: str = None, 24 | allow_nonmatching_filename: bool = False, verbose: bool = True 25 | ) -> Type[BaseReaderWriter]: 26 | if 'overwrite_image_reader_writer' in dataset_json_content.keys() and \ 27 | dataset_json_content['overwrite_image_reader_writer'] != 'None': 28 | ioclass_name = dataset_json_content['overwrite_image_reader_writer'] 29 | # trying to find that class in the nnunetv2.imageio module 30 | try: 31 | ret = recursive_find_reader_writer_by_name(ioclass_name) 32 | if verbose: print('Using %s reader/writer' % ret) 33 | return ret 34 | except RuntimeError: 35 | if verbose: print('Warning: Unable to find ioclass specified in dataset.json: %s' % ioclass_name) 36 | if verbose: print('Trying to automatically determine desired class') 37 | return determine_reader_writer_from_file_ending(dataset_json_content['file_ending'], example_file, 38 | allow_nonmatching_filename, verbose) 39 | 40 | 41 | def determine_reader_writer_from_file_ending(file_ending: str, example_file: str = None, allow_nonmatching_filename: bool = False, 42 | verbose: bool = True): 43 | for rw in LIST_OF_IO_CLASSES: 44 | if file_ending.lower() in rw.supported_file_endings: 45 | if example_file is not None: 46 | # if an example file is provided, try if we can actually read it. If not move on to the next reader 47 | try: 48 | tmp = rw() 49 | _ = tmp.read_images((example_file,)) 50 | if verbose: print('Using %s as reader/writer' % rw) 51 | return rw 52 | except: 53 | if verbose: print(f'Failed to open file {example_file} with reader {rw}:') 54 | traceback.print_exc() 55 | pass 56 | else: 57 | if verbose: print('Using %s as reader/writer' % rw) 58 | return rw 59 | else: 60 | if allow_nonmatching_filename and example_file is not None: 61 | try: 62 | tmp = rw() 63 | _ = tmp.read_images((example_file,)) 64 | if verbose: print('Using %s as reader/writer' % rw) 65 | return rw 66 | except: 67 | if verbose: print(f'Failed to open file {example_file} with reader {rw}:') 68 | if verbose: traceback.print_exc() 69 | pass 70 | raise RuntimeError("Unable to determine a reader for file ending %s and file %s (file None means no file provided)." % (file_ending, example_file)) 71 | 72 | 73 | def recursive_find_reader_writer_by_name(rw_class_name: str) -> Type[BaseReaderWriter]: 74 | ret = recursive_find_python_class(join(nnunetv2.__path__[0], "imageio"), rw_class_name, 'nnunetv2.imageio') 75 | if ret is None: 76 | raise RuntimeError("Unable to find reader writer class '%s'. Please make sure this class is located in the " 77 | "nnunetv2.imageio module." % rw_class_name) 78 | else: 79 | return ret 80 | -------------------------------------------------------------------------------- /nnunetv2/inference/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/inference/__init__.py -------------------------------------------------------------------------------- /nnunetv2/model_sharing/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/model_sharing/__init__.py -------------------------------------------------------------------------------- /nnunetv2/model_sharing/entry_points.py: -------------------------------------------------------------------------------- 1 | from nnunetv2.model_sharing.model_download import download_and_install_from_url 2 | from nnunetv2.model_sharing.model_export import export_pretrained_model 3 | from nnunetv2.model_sharing.model_import import install_model_from_zip_file 4 | 5 | 6 | def print_license_warning(): 7 | print('') 8 | print('######################################################') 9 | print('!!!!!!!!!!!!!!!!!!!!!!!!WARNING!!!!!!!!!!!!!!!!!!!!!!!') 10 | print('######################################################') 11 | print("Using the pretrained model weights is subject to the license of the dataset they were trained on. Some " 12 | "allow commercial use, others don't. It is your responsibility to make sure you use them appropriately! Use " 13 | "nnUNet_print_pretrained_model_info(task_name) to see a summary of the dataset and where to find its license!") 14 | print('######################################################') 15 | print('') 16 | 17 | 18 | def download_by_url(): 19 | import argparse 20 | parser = argparse.ArgumentParser( 21 | description="Use this to download pretrained models. This script is intended to download models via url only. " 22 | "CAREFUL: This script will overwrite " 23 | "existing models (if they share the same trainer class and plans as " 24 | "the pretrained model.") 25 | parser.add_argument("url", type=str, help='URL of the pretrained model') 26 | args = parser.parse_args() 27 | url = args.url 28 | download_and_install_from_url(url) 29 | 30 | 31 | def install_from_zip_entry_point(): 32 | import argparse 33 | parser = argparse.ArgumentParser( 34 | description="Use this to install a zip file containing a pretrained model.") 35 | parser.add_argument("zip", type=str, help='zip file') 36 | args = parser.parse_args() 37 | zip = args.zip 38 | install_model_from_zip_file(zip) 39 | 40 | 41 | def export_pretrained_model_entry(): 42 | import argparse 43 | parser = argparse.ArgumentParser( 44 | description="Use this to export a trained model as a zip file.") 45 | parser.add_argument('-d', type=str, required=True, help='Dataset name or id') 46 | parser.add_argument('-o', type=str, required=True, help='Output file name') 47 | parser.add_argument('-c', nargs='+', type=str, required=False, 48 | default=('3d_lowres', '3d_fullres', '2d', '3d_cascade_fullres'), 49 | help="List of configuration names") 50 | parser.add_argument('-tr', required=False, type=str, default='nnUNetTrainer', help='Trainer class') 51 | parser.add_argument('-p', required=False, type=str, default='nnUNetPlans', help='plans identifier') 52 | parser.add_argument('-f', required=False, nargs='+', type=str, default=(0, 1, 2, 3, 4), help='list of fold ids') 53 | parser.add_argument('-chk', required=False, nargs='+', type=str, default=('checkpoint_final.pth', ), 54 | help='Lis tof checkpoint names to export. Default: checkpoint_final.pth') 55 | parser.add_argument('--not_strict', action='store_false', default=False, required=False, help='Set this to allow missing folds and/or configurations') 56 | parser.add_argument('--exp_cv_preds', action='store_true', required=False, help='Set this to export the cross-validation predictions as well') 57 | args = parser.parse_args() 58 | 59 | export_pretrained_model(dataset_name_or_id=args.d, output_file=args.o, configurations=args.c, trainer=args.tr, 60 | plans_identifier=args.p, folds=args.f, strict=not args.not_strict, save_checkpoints=args.chk, 61 | export_crossval_predictions=args.exp_cv_preds) 62 | -------------------------------------------------------------------------------- /nnunetv2/model_sharing/model_download.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import requests 4 | from batchgenerators.utilities.file_and_folder_operations import * 5 | from time import time 6 | from nnunetv2.model_sharing.model_import import install_model_from_zip_file 7 | from nnunetv2.paths import nnUNet_results 8 | from tqdm import tqdm 9 | 10 | 11 | def download_and_install_from_url(url): 12 | assert nnUNet_results is not None, "Cannot install model because network_training_output_dir is not " \ 13 | "set (RESULTS_FOLDER missing as environment variable, see " \ 14 | "Installation instructions)" 15 | print('Downloading pretrained model from url:', url) 16 | import http.client 17 | http.client.HTTPConnection._http_vsn = 10 18 | http.client.HTTPConnection._http_vsn_str = 'HTTP/1.0' 19 | 20 | import os 21 | home = os.path.expanduser('~') 22 | random_number = int(time() * 1e7) 23 | tempfile = join(home, '.nnunetdownload_%s' % str(random_number)) 24 | 25 | try: 26 | download_file(url=url, local_filename=tempfile, chunk_size=8192 * 16) 27 | print("Download finished. Extracting...") 28 | install_model_from_zip_file(tempfile) 29 | print("Done") 30 | except Exception as e: 31 | raise e 32 | finally: 33 | if isfile(tempfile): 34 | os.remove(tempfile) 35 | 36 | 37 | def download_file(url: str, local_filename: str, chunk_size: Optional[int] = 8192 * 16) -> str: 38 | # borrowed from https://stackoverflow.com/questions/16694907/download-large-file-in-python-with-requests 39 | # NOTE the stream=True parameter below 40 | with requests.get(url, stream=True, timeout=100) as r: 41 | r.raise_for_status() 42 | with tqdm.wrapattr(open(local_filename, 'wb'), "write", total=int(r.headers.get("Content-Length"))) as f: 43 | for chunk in r.iter_content(chunk_size=chunk_size): 44 | f.write(chunk) 45 | return local_filename 46 | 47 | 48 | -------------------------------------------------------------------------------- /nnunetv2/model_sharing/model_import.py: -------------------------------------------------------------------------------- 1 | import zipfile 2 | 3 | from nnunetv2.paths import nnUNet_results 4 | 5 | 6 | def install_model_from_zip_file(zip_file: str): 7 | with zipfile.ZipFile(zip_file, 'r') as zip_ref: 8 | zip_ref.extractall(nnUNet_results) -------------------------------------------------------------------------------- /nnunetv2/paths.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | 17 | """ 18 | PLEASE READ paths.md FOR INFORMATION TO HOW TO SET THIS UP 19 | """ 20 | nnUNet_raw = os.environ.get('nnUNet_raw') 21 | nnUNet_preprocessed = os.environ.get('nnUNet_preprocessed') 22 | nnUNet_results = os.path.join(os.environ.get('nnUNet_results'), os.environ.get('MODEL_NAME')) 23 | os.makedirs(nnUNet_results, exist_ok=True) 24 | 25 | if nnUNet_raw is None: 26 | print("nnUNet_raw is not defined and nnU-Net can only be used on data for which preprocessed files " 27 | "are already present on your system. nnU-Net cannot be used for experiment planning and preprocessing like " 28 | "this. If this is not intended, please read documentation/setting_up_paths.md for information on how to set " 29 | "this up properly.") 30 | 31 | if nnUNet_preprocessed is None: 32 | print("nnUNet_preprocessed is not defined and nnU-Net can not be used for preprocessing " 33 | "or training. If this is not intended, please read documentation/setting_up_paths.md for information on how " 34 | "to set this up.") 35 | 36 | if nnUNet_results is None: 37 | print("nnUNet_results is not defined and nnU-Net cannot be used for training or " 38 | "inference. If this is not intended behavior, please read documentation/setting_up_paths.md for information " 39 | "on how to set this up.") 40 | -------------------------------------------------------------------------------- /nnunetv2/postprocessing/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/postprocessing/__init__.py -------------------------------------------------------------------------------- /nnunetv2/preprocessing/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/preprocessing/__init__.py -------------------------------------------------------------------------------- /nnunetv2/preprocessing/cropping/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/preprocessing/cropping/__init__.py -------------------------------------------------------------------------------- /nnunetv2/preprocessing/cropping/cropping.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | # Hello! crop_to_nonzero is the function you are looking for. Ignore the rest. 5 | from acvl_utils.cropping_and_padding.bounding_boxes import get_bbox_from_mask, crop_to_bbox, bounding_box_to_slice 6 | 7 | 8 | def create_nonzero_mask(data): 9 | """ 10 | 11 | :param data: 12 | :return: the mask is True where the data is nonzero 13 | """ 14 | from scipy.ndimage import binary_fill_holes 15 | assert len(data.shape) == 4 or len(data.shape) == 3, "data must have shape (C, X, Y, Z) or shape (C, X, Y)" 16 | nonzero_mask = np.zeros(data.shape[1:], dtype=bool) 17 | for c in range(data.shape[0]): 18 | this_mask = data[c] != 0 19 | nonzero_mask = nonzero_mask | this_mask 20 | nonzero_mask = binary_fill_holes(nonzero_mask) 21 | return nonzero_mask 22 | 23 | 24 | def crop_to_nonzero(data, seg=None, nonzero_label=-1): 25 | """ 26 | 27 | :param data: 28 | :param seg: 29 | :param nonzero_label: this will be written into the segmentation map 30 | :return: 31 | """ 32 | nonzero_mask = create_nonzero_mask(data) 33 | bbox = get_bbox_from_mask(nonzero_mask) 34 | 35 | slicer = bounding_box_to_slice(bbox) 36 | data = data[tuple([slice(None), *slicer])] 37 | 38 | if seg is not None: 39 | seg = seg[tuple([slice(None), *slicer])] 40 | 41 | nonzero_mask = nonzero_mask[slicer][None] 42 | if seg is not None: 43 | seg[(seg == 0) & (~nonzero_mask)] = nonzero_label 44 | else: 45 | nonzero_mask = nonzero_mask.astype(np.int8) 46 | nonzero_mask[nonzero_mask == 0] = nonzero_label 47 | nonzero_mask[nonzero_mask > 0] = 0 48 | seg = nonzero_mask 49 | return data, seg, bbox 50 | 51 | 52 | -------------------------------------------------------------------------------- /nnunetv2/preprocessing/normalization/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/preprocessing/normalization/__init__.py -------------------------------------------------------------------------------- /nnunetv2/preprocessing/normalization/default_normalization_schemes.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from typing import Type 3 | 4 | import numpy as np 5 | from numpy import number 6 | 7 | 8 | class ImageNormalization(ABC): 9 | leaves_pixels_outside_mask_at_zero_if_use_mask_for_norm_is_true = None 10 | 11 | def __init__(self, use_mask_for_norm: bool = None, intensityproperties: dict = None, 12 | target_dtype: Type[number] = np.float32): 13 | assert use_mask_for_norm is None or isinstance(use_mask_for_norm, bool) 14 | self.use_mask_for_norm = use_mask_for_norm 15 | assert isinstance(intensityproperties, dict) 16 | self.intensityproperties = intensityproperties 17 | self.target_dtype = target_dtype 18 | 19 | @abstractmethod 20 | def run(self, image: np.ndarray, seg: np.ndarray = None) -> np.ndarray: 21 | """ 22 | Image and seg must have the same shape. Seg is not always used 23 | """ 24 | pass 25 | 26 | 27 | class ZScoreNormalization(ImageNormalization): 28 | leaves_pixels_outside_mask_at_zero_if_use_mask_for_norm_is_true = True 29 | 30 | def run(self, image: np.ndarray, seg: np.ndarray = None) -> np.ndarray: 31 | """ 32 | here seg is used to store the zero valued region. The value for that region in the segmentation is -1 by 33 | default. 34 | """ 35 | image = image.astype(self.target_dtype) 36 | if self.use_mask_for_norm is not None and self.use_mask_for_norm: 37 | # negative values in the segmentation encode the 'outside' region (think zero values around the brain as 38 | # in BraTS). We want to run the normalization only in the brain region, so we need to mask the image. 39 | # The default nnU-net sets use_mask_for_norm to True if cropping to the nonzero region substantially 40 | # reduced the image size. 41 | mask = seg >= 0 42 | mean = image[mask].mean() 43 | std = image[mask].std() 44 | image[mask] = (image[mask] - mean) / (max(std, 1e-8)) 45 | else: 46 | mean = image.mean() 47 | std = image.std() 48 | image = (image - mean) / (max(std, 1e-8)) 49 | return image 50 | 51 | 52 | class CTNormalization(ImageNormalization): 53 | leaves_pixels_outside_mask_at_zero_if_use_mask_for_norm_is_true = False 54 | 55 | def run(self, image: np.ndarray, seg: np.ndarray = None) -> np.ndarray: 56 | assert self.intensityproperties is not None, "CTNormalization requires intensity properties" 57 | image = image.astype(self.target_dtype) 58 | mean_intensity = self.intensityproperties['mean'] 59 | std_intensity = self.intensityproperties['std'] 60 | lower_bound = self.intensityproperties['percentile_00_5'] 61 | upper_bound = self.intensityproperties['percentile_99_5'] 62 | image = np.clip(image, lower_bound, upper_bound) 63 | image = (image - mean_intensity) / max(std_intensity, 1e-8) 64 | return image 65 | 66 | 67 | class NoNormalization(ImageNormalization): 68 | leaves_pixels_outside_mask_at_zero_if_use_mask_for_norm_is_true = False 69 | 70 | def run(self, image: np.ndarray, seg: np.ndarray = None) -> np.ndarray: 71 | return image.astype(self.target_dtype) 72 | 73 | 74 | class RescaleTo01Normalization(ImageNormalization): 75 | leaves_pixels_outside_mask_at_zero_if_use_mask_for_norm_is_true = False 76 | 77 | def run(self, image: np.ndarray, seg: np.ndarray = None) -> np.ndarray: 78 | image = image.astype(self.target_dtype) 79 | image = image - image.min() 80 | image = image / np.clip(image.max(), a_min=1e-8, a_max=None) 81 | return image 82 | 83 | 84 | class RGBTo01Normalization(ImageNormalization): 85 | leaves_pixels_outside_mask_at_zero_if_use_mask_for_norm_is_true = False 86 | 87 | def run(self, image: np.ndarray, seg: np.ndarray = None) -> np.ndarray: 88 | assert image.min() >= 0, "RGB images are uint 8, for whatever reason I found pixel values smaller than 0. " \ 89 | "Your images do not seem to be RGB images" 90 | assert image.max() <= 255, "RGB images are uint 8, for whatever reason I found pixel values greater than 255" \ 91 | ". Your images do not seem to be RGB images" 92 | image = image.astype(self.target_dtype) 93 | image = image / 255. 94 | return image 95 | 96 | -------------------------------------------------------------------------------- /nnunetv2/preprocessing/normalization/map_channel_name_to_normalization.py: -------------------------------------------------------------------------------- 1 | from typing import Type 2 | 3 | from nnunetv2.preprocessing.normalization.default_normalization_schemes import CTNormalization, NoNormalization, \ 4 | ZScoreNormalization, RescaleTo01Normalization, RGBTo01Normalization, ImageNormalization 5 | 6 | channel_name_to_normalization_mapping = { 7 | 'CT': CTNormalization, 8 | 'noNorm': NoNormalization, 9 | 'zscore': ZScoreNormalization, 10 | 'rescale_0_1': RescaleTo01Normalization, 11 | 'rgb_to_0_1': RGBTo01Normalization 12 | } 13 | 14 | 15 | def get_normalization_scheme(channel_name: str) -> Type[ImageNormalization]: 16 | """ 17 | If we find the channel_name in channel_name_to_normalization_mapping return the corresponding normalization. If it is 18 | not found, use the default (ZScoreNormalization) 19 | """ 20 | norm_scheme = channel_name_to_normalization_mapping.get(channel_name) 21 | if norm_scheme is None: 22 | norm_scheme = ZScoreNormalization 23 | # print('Using %s for image normalization' % norm_scheme.__name__) 24 | return norm_scheme 25 | -------------------------------------------------------------------------------- /nnunetv2/preprocessing/preprocessors/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/preprocessing/preprocessors/__init__.py -------------------------------------------------------------------------------- /nnunetv2/preprocessing/resampling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/preprocessing/resampling/__init__.py -------------------------------------------------------------------------------- /nnunetv2/preprocessing/resampling/utils.py: -------------------------------------------------------------------------------- 1 | from typing import Callable 2 | 3 | import nnunetv2 4 | from batchgenerators.utilities.file_and_folder_operations import join 5 | from nnunetv2.utilities.find_class_by_name import recursive_find_python_class 6 | 7 | 8 | def recursive_find_resampling_fn_by_name(resampling_fn: str) -> Callable: 9 | ret = recursive_find_python_class(join(nnunetv2.__path__[0], "preprocessing", "resampling"), resampling_fn, 10 | 'nnunetv2.preprocessing.resampling') 11 | if ret is None: 12 | raise RuntimeError("Unable to find resampling function named '%s'. Please make sure this fn is located in the " 13 | "nnunetv2.preprocessing.resampling module." % resampling_fn) 14 | else: 15 | return ret 16 | -------------------------------------------------------------------------------- /nnunetv2/run/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/run/__init__.py -------------------------------------------------------------------------------- /nnunetv2/run/load_pretrained_weights.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn.parallel import DistributedDataParallel as DDP 3 | 4 | 5 | def load_pretrained_weights(network, fname, verbose=False): 6 | """ 7 | THIS DOES NOT TRANSFER SEGMENTATION HEADS! 8 | 9 | network can be either a plain model or DDP. We need to account for that in the parameter names 10 | """ 11 | saved_model = torch.load(fname) 12 | pretrained_dict = saved_model['network_weights'] 13 | is_ddp = isinstance(network, DDP) 14 | 15 | skip_strings_in_pretrained = [ 16 | '.seg_layers.', 17 | ] 18 | 19 | model_dict = network.state_dict() 20 | # verify that all but the segmentation layers have the same shape 21 | for key, _ in model_dict.items(): 22 | if is_ddp: 23 | key_pretrained = key[7:] 24 | else: 25 | key_pretrained = key 26 | if all([i not in key for i in skip_strings_in_pretrained]): 27 | assert key_pretrained in pretrained_dict, \ 28 | f"Key {key_pretrained} is missing in the pretrained model weights. The pretrained weights do not seem to be " \ 29 | f"compatible with your network." 30 | assert model_dict[key].shape == pretrained_dict[key_pretrained].shape, \ 31 | f"The shape of the parameters of key {key_pretrained} is not the same. Pretrained model: " \ 32 | f"{pretrained_dict[key_pretrained].shape}; your network: {model_dict[key]}. The pretrained model " \ 33 | f"does not seem to be compatible with your network." 34 | 35 | # fun fact this does allow loading from parameters that do not cover the entire network. For example pretrained 36 | # encoders 37 | # I didnt even know that you could put if statements into dict comprehensions. Damn am I good. 38 | pretrained_dict = {'module.' + k if is_ddp else k: v 39 | for k, v in pretrained_dict.items() 40 | if (('module.' + k if is_ddp else k) in model_dict) and all([i not in k for i in skip_strings_in_pretrained])} 41 | # yet another line of death right here... 42 | 43 | model_dict.update(pretrained_dict) 44 | 45 | print("################### Loading pretrained weights from file ", fname, '###################') 46 | if verbose: 47 | print("Below is the list of overlapping blocks in pretrained model and nnUNet architecture:") 48 | for key, _ in pretrained_dict.items(): 49 | print(key[7:] if is_ddp else key) 50 | print("################### Done ###################") 51 | network.load_state_dict(model_dict) 52 | 53 | 54 | -------------------------------------------------------------------------------- /nnunetv2/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/tests/__init__.py -------------------------------------------------------------------------------- /nnunetv2/tests/integration_tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/tests/integration_tests/__init__.py -------------------------------------------------------------------------------- /nnunetv2/tests/integration_tests/add_lowres_and_cascade.py: -------------------------------------------------------------------------------- 1 | from batchgenerators.utilities.file_and_folder_operations import * 2 | 3 | from nnunetv2.paths import nnUNet_preprocessed 4 | from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name 5 | 6 | if __name__ == '__main__': 7 | import argparse 8 | 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('-d', nargs='+', type=int, help='List of dataset ids') 11 | args = parser.parse_args() 12 | 13 | for d in args.d: 14 | dataset_name = maybe_convert_to_dataset_name(d) 15 | plans = load_json(join(nnUNet_preprocessed, dataset_name, 'nnUNetPlans.json')) 16 | plans['configurations']['3d_lowres'] = { 17 | "data_identifier": "nnUNetPlans_3d_lowres", # do not be a dumbo and forget this. I was a dumbo. And I paid dearly with ~10 min debugging time 18 | 'inherits_from': '3d_fullres', 19 | "patch_size": [20, 28, 20], 20 | "median_image_size_in_voxels": [18.0, 25.0, 18.0], 21 | "spacing": [2.0, 2.0, 2.0], 22 | "n_conv_per_stage_encoder": [2, 2, 2], 23 | "n_conv_per_stage_decoder": [2, 2], 24 | "num_pool_per_axis": [2, 2, 2], 25 | "pool_op_kernel_sizes": [[1, 1, 1], [2, 2, 2], [2, 2, 2]], 26 | "conv_kernel_sizes": [[3, 3, 3], [3, 3, 3], [3, 3, 3]], 27 | "next_stage": "3d_cascade_fullres" 28 | } 29 | plans['configurations']['3d_cascade_fullres'] = { 30 | 'inherits_from': '3d_fullres', 31 | "previous_stage": "3d_lowres" 32 | } 33 | save_json(plans, join(nnUNet_preprocessed, dataset_name, 'nnUNetPlans.json'), sort_keys=False) -------------------------------------------------------------------------------- /nnunetv2/tests/integration_tests/cleanup_integration_test.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | 3 | from batchgenerators.utilities.file_and_folder_operations import isdir, join 4 | 5 | from nnunetv2.paths import nnUNet_raw, nnUNet_results, nnUNet_preprocessed 6 | 7 | if __name__ == '__main__': 8 | # deletes everything! 9 | dataset_names = [ 10 | 'Dataset996_IntegrationTest_Hippocampus_regions_ignore', 11 | 'Dataset997_IntegrationTest_Hippocampus_regions', 12 | 'Dataset998_IntegrationTest_Hippocampus_ignore', 13 | 'Dataset999_IntegrationTest_Hippocampus', 14 | ] 15 | for fld in [nnUNet_raw, nnUNet_preprocessed, nnUNet_results]: 16 | for d in dataset_names: 17 | if isdir(join(fld, d)): 18 | shutil.rmtree(join(fld, d)) 19 | 20 | -------------------------------------------------------------------------------- /nnunetv2/tests/integration_tests/run_integration_test_bestconfig_inference.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from batchgenerators.utilities.file_and_folder_operations import join, load_pickle 4 | 5 | from nnunetv2.ensembling.ensemble import ensemble_folders 6 | from nnunetv2.evaluation.find_best_configuration import find_best_configuration, \ 7 | dumb_trainer_config_plans_to_trained_models_dict 8 | from nnunetv2.inference.predict_from_raw_data import predict_from_raw_data 9 | from nnunetv2.paths import nnUNet_raw, nnUNet_results 10 | from nnunetv2.postprocessing.remove_connected_components import apply_postprocessing_to_folder 11 | from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name 12 | from nnunetv2.utilities.file_path_utilities import get_output_folder 13 | 14 | 15 | if __name__ == '__main__': 16 | """ 17 | Predicts the imagesTs folder with the best configuration and applies postprocessing 18 | """ 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument('-d', type=int, help='dataset id') 21 | args = parser.parse_args() 22 | d = args.d 23 | 24 | dataset_name = maybe_convert_to_dataset_name(d) 25 | source_dir = join(nnUNet_raw, dataset_name, 'imagesTs') 26 | target_dir_base = join(nnUNet_results, dataset_name) 27 | 28 | models = dumb_trainer_config_plans_to_trained_models_dict(['nnUNetTrainer_5epochs'], 29 | ['2d', 30 | '3d_lowres', 31 | '3d_cascade_fullres', 32 | '3d_fullres'], 33 | ['nnUNetPlans']) 34 | ret = find_best_configuration(d, models, allow_ensembling=True, num_processes=8, overwrite=True, 35 | folds=(0, 1, 2, 3, 4), strict=True) 36 | 37 | has_ensemble = len(ret['best_model_or_ensemble']['selected_model_or_models']) > 1 38 | 39 | # we don't use all folds to speed stuff up 40 | used_folds = (0, 3) 41 | output_folders = [] 42 | for im in ret['best_model_or_ensemble']['selected_model_or_models']: 43 | output_dir = join(target_dir_base, f"pred_{im['configuration']}") 44 | model_folder = get_output_folder(d, im['trainer'], im['plans_identifier'], im['configuration']) 45 | # note that if the best model is the enseble of 3d_lowres and 3d cascade then 3d_lowres will be predicted 46 | # twice (once standalone and once to generate the predictions for the cascade) because we don't reuse the 47 | # prediction here. Proper way would be to check for that and 48 | # then give the output of 3d_lowres inference to the folder_with_segs_from_prev_stage kwarg in 49 | # predict_from_raw_data. Since we allow for 50 | # dynamically setting 'previous_stage' in the plans I am too lazy to implement this here. This is just an 51 | # integration test after all. Take a closer look at how this in handled in predict_from_raw_data 52 | predict_from_raw_data(list_of_lists_or_source_folder=source_dir, output_folder=output_dir, 53 | model_training_output_dir=model_folder, use_folds=used_folds, 54 | save_probabilities=has_ensemble, verbose=False, overwrite=True) 55 | output_folders.append(output_dir) 56 | 57 | # if we have an ensemble, we need to ensemble the results 58 | if has_ensemble: 59 | ensemble_folders(output_folders, join(target_dir_base, 'ensemble_predictions'), save_merged_probabilities=False) 60 | folder_for_pp = join(target_dir_base, 'ensemble_predictions') 61 | else: 62 | folder_for_pp = output_folders[0] 63 | 64 | # apply postprocessing 65 | pp_fns, pp_fn_kwargs = load_pickle(ret['best_model_or_ensemble']['postprocessing_file']) 66 | apply_postprocessing_to_folder(folder_for_pp, join(target_dir_base, 'ensemble_predictions_postprocessed'), 67 | pp_fns, 68 | pp_fn_kwargs, plans_file_or_dict=ret['best_model_or_ensemble']['some_plans_file']) 69 | -------------------------------------------------------------------------------- /nnunetv2/training/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/data_augmentation/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/data_augmentation/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/data_augmentation/compute_initial_patch_size.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def get_patch_size(final_patch_size, rot_x, rot_y, rot_z, scale_range): 5 | if isinstance(rot_x, (tuple, list)): 6 | rot_x = max(np.abs(rot_x)) 7 | if isinstance(rot_y, (tuple, list)): 8 | rot_y = max(np.abs(rot_y)) 9 | if isinstance(rot_z, (tuple, list)): 10 | rot_z = max(np.abs(rot_z)) 11 | rot_x = min(90 / 360 * 2. * np.pi, rot_x) 12 | rot_y = min(90 / 360 * 2. * np.pi, rot_y) 13 | rot_z = min(90 / 360 * 2. * np.pi, rot_z) 14 | from batchgenerators.augmentations.utils import rotate_coords_3d, rotate_coords_2d 15 | coords = np.array(final_patch_size) 16 | final_shape = np.copy(coords) 17 | if len(coords) == 3: 18 | final_shape = np.max(np.vstack((np.abs(rotate_coords_3d(coords, rot_x, 0, 0)), final_shape)), 0) 19 | final_shape = np.max(np.vstack((np.abs(rotate_coords_3d(coords, 0, rot_y, 0)), final_shape)), 0) 20 | final_shape = np.max(np.vstack((np.abs(rotate_coords_3d(coords, 0, 0, rot_z)), final_shape)), 0) 21 | elif len(coords) == 2: 22 | final_shape = np.max(np.vstack((np.abs(rotate_coords_2d(coords, rot_x)), final_shape)), 0) 23 | final_shape /= min(scale_range) 24 | return final_shape.astype(int) 25 | -------------------------------------------------------------------------------- /nnunetv2/training/data_augmentation/custom_transforms/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/data_augmentation/custom_transforms/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/data_augmentation/custom_transforms/deep_supervision_donwsampling.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Union, List 2 | 3 | from batchgenerators.augmentations.utils import resize_segmentation 4 | from batchgenerators.transforms.abstract_transforms import AbstractTransform 5 | import numpy as np 6 | 7 | 8 | class DownsampleSegForDSTransform2(AbstractTransform): 9 | ''' 10 | data_dict['output_key'] will be a list of segmentations scaled according to ds_scales 11 | ''' 12 | def __init__(self, ds_scales: Union[List, Tuple], 13 | order: int = 0, input_key: str = "seg", 14 | output_key: str = "seg", axes: Tuple[int] = None): 15 | """ 16 | Downscales data_dict[input_key] according to ds_scales. Each entry in ds_scales specified one deep supervision 17 | output and its resolution relative to the original data, for example 0.25 specifies 1/4 of the original shape. 18 | ds_scales can also be a tuple of tuples, for example ((1, 1, 1), (0.5, 0.5, 0.5)) to specify the downsampling 19 | for each axis independently 20 | """ 21 | self.axes = axes 22 | self.output_key = output_key 23 | self.input_key = input_key 24 | self.order = order 25 | self.ds_scales = ds_scales 26 | 27 | def __call__(self, **data_dict): 28 | if self.axes is None: 29 | axes = list(range(2, len(data_dict[self.input_key].shape))) 30 | else: 31 | axes = self.axes 32 | 33 | output = [] 34 | for s in self.ds_scales: 35 | if not isinstance(s, (tuple, list)): 36 | s = [s] * len(axes) 37 | else: 38 | assert len(s) == len(axes), f'If ds_scales is a tuple for each resolution (one downsampling factor ' \ 39 | f'for each axis) then the number of entried in that tuple (here ' \ 40 | f'{len(s)}) must be the same as the number of axes (here {len(axes)}).' 41 | 42 | if all([i == 1 for i in s]): 43 | output.append(data_dict[self.input_key]) 44 | else: 45 | new_shape = np.array(data_dict[self.input_key].shape).astype(float) 46 | for i, a in enumerate(axes): 47 | new_shape[a] *= s[i] 48 | new_shape = np.round(new_shape).astype(int) 49 | out_seg = np.zeros(new_shape, dtype=data_dict[self.input_key].dtype) 50 | for b in range(data_dict[self.input_key].shape[0]): 51 | for c in range(data_dict[self.input_key].shape[1]): 52 | out_seg[b, c] = resize_segmentation(data_dict[self.input_key][b, c], new_shape[2:], self.order) 53 | output.append(out_seg) 54 | data_dict[self.output_key] = output 55 | return data_dict 56 | -------------------------------------------------------------------------------- /nnunetv2/training/data_augmentation/custom_transforms/limited_length_multithreaded_augmenter.py: -------------------------------------------------------------------------------- 1 | from batchgenerators.dataloading.nondet_multi_threaded_augmenter import NonDetMultiThreadedAugmenter 2 | 3 | 4 | class LimitedLenWrapper(NonDetMultiThreadedAugmenter): 5 | def __init__(self, my_imaginary_length, *args, **kwargs): 6 | super().__init__(*args, **kwargs) 7 | self.len = my_imaginary_length 8 | 9 | def __len__(self): 10 | return self.len 11 | -------------------------------------------------------------------------------- /nnunetv2/training/data_augmentation/custom_transforms/manipulating_data_dict.py: -------------------------------------------------------------------------------- 1 | from batchgenerators.transforms.abstract_transforms import AbstractTransform 2 | 3 | 4 | class RemoveKeyTransform(AbstractTransform): 5 | def __init__(self, key_to_remove: str): 6 | self.key_to_remove = key_to_remove 7 | 8 | def __call__(self, **data_dict): 9 | _ = data_dict.pop(self.key_to_remove, None) 10 | return data_dict 11 | -------------------------------------------------------------------------------- /nnunetv2/training/data_augmentation/custom_transforms/masking.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | from batchgenerators.transforms.abstract_transforms import AbstractTransform 4 | 5 | 6 | class MaskTransform(AbstractTransform): 7 | def __init__(self, apply_to_channels: List[int], mask_idx_in_seg: int = 0, set_outside_to: int = 0, 8 | data_key: str = "data", seg_key: str = "seg"): 9 | """ 10 | Sets everything outside the mask to 0. CAREFUL! outside is defined as < 0, not =0 (in the Mask)!!! 11 | """ 12 | self.apply_to_channels = apply_to_channels 13 | self.seg_key = seg_key 14 | self.data_key = data_key 15 | self.set_outside_to = set_outside_to 16 | self.mask_idx_in_seg = mask_idx_in_seg 17 | 18 | def __call__(self, **data_dict): 19 | mask = data_dict[self.seg_key][:, self.mask_idx_in_seg] < 0 20 | for c in self.apply_to_channels: 21 | data_dict[self.data_key][:, c][mask] = self.set_outside_to 22 | return data_dict 23 | -------------------------------------------------------------------------------- /nnunetv2/training/data_augmentation/custom_transforms/region_based_training.py: -------------------------------------------------------------------------------- 1 | from typing import List, Tuple, Union 2 | 3 | from batchgenerators.transforms.abstract_transforms import AbstractTransform 4 | import numpy as np 5 | 6 | 7 | class ConvertSegmentationToRegionsTransform(AbstractTransform): 8 | def __init__(self, regions: Union[List, Tuple], 9 | seg_key: str = "seg", output_key: str = "seg", seg_channel: int = 0): 10 | """ 11 | regions are tuple of tuples where each inner tuple holds the class indices that are merged into one region, 12 | example: 13 | regions= ((1, 2), (2, )) will result in 2 regions: one covering the region of labels 1&2 and the other just 2 14 | :param regions: 15 | :param seg_key: 16 | :param output_key: 17 | """ 18 | self.seg_channel = seg_channel 19 | self.output_key = output_key 20 | self.seg_key = seg_key 21 | self.regions = regions 22 | 23 | def __call__(self, **data_dict): 24 | seg = data_dict.get(self.seg_key) 25 | num_regions = len(self.regions) 26 | if seg is not None: 27 | seg_shp = seg.shape 28 | output_shape = list(seg_shp) 29 | output_shape[1] = num_regions 30 | region_output = np.zeros(output_shape, dtype=seg.dtype) 31 | for b in range(seg_shp[0]): 32 | for region_id, region_source_labels in enumerate(self.regions): 33 | if not isinstance(region_source_labels, (list, tuple)): 34 | region_source_labels = (region_source_labels, ) 35 | for label_value in region_source_labels: 36 | region_output[b, region_id][seg[b, self.seg_channel] == label_value] = 1 37 | data_dict[self.output_key] = region_output 38 | return data_dict 39 | -------------------------------------------------------------------------------- /nnunetv2/training/data_augmentation/custom_transforms/transforms_for_dummy_2d.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Union, List 2 | 3 | from batchgenerators.transforms.abstract_transforms import AbstractTransform 4 | 5 | 6 | class Convert3DTo2DTransform(AbstractTransform): 7 | def __init__(self, apply_to_keys: Union[List[str], Tuple[str]] = ('data', 'seg')): 8 | """ 9 | Transforms a 5D array (b, c, x, y, z) to a 4D array (b, c * x, y, z) by overloading the color channel 10 | """ 11 | self.apply_to_keys = apply_to_keys 12 | 13 | def __call__(self, **data_dict): 14 | for k in self.apply_to_keys: 15 | shp = data_dict[k].shape 16 | assert len(shp) == 5, 'This transform only works on 3D data, so expects 5D tensor (b, c, x, y, z) as input.' 17 | data_dict[k] = data_dict[k].reshape((shp[0], shp[1] * shp[2], shp[3], shp[4])) 18 | shape_key = f'orig_shape_{k}' 19 | assert shape_key not in data_dict.keys(), f'Convert3DTo2DTransform needs to store the original shape. ' \ 20 | f'It does that using the {shape_key} key. That key is ' \ 21 | f'already taken. Bummer.' 22 | data_dict[shape_key] = shp 23 | return data_dict 24 | 25 | 26 | class Convert2DTo3DTransform(AbstractTransform): 27 | def __init__(self, apply_to_keys: Union[List[str], Tuple[str]] = ('data', 'seg')): 28 | """ 29 | Reverts Convert3DTo2DTransform by transforming a 4D array (b, c * x, y, z) back to 5D (b, c, x, y, z) 30 | """ 31 | self.apply_to_keys = apply_to_keys 32 | 33 | def __call__(self, **data_dict): 34 | for k in self.apply_to_keys: 35 | shape_key = f'orig_shape_{k}' 36 | assert shape_key in data_dict.keys(), f'Did not find key {shape_key} in data_dict. Shitty. ' \ 37 | f'Convert2DTo3DTransform only works in tandem with ' \ 38 | f'Convert3DTo2DTransform and you probably forgot to add ' \ 39 | f'Convert3DTo2DTransform to your pipeline. (Convert3DTo2DTransform ' \ 40 | f'is where the missing key is generated)' 41 | original_shape = data_dict[shape_key] 42 | current_shape = data_dict[k].shape 43 | data_dict[k] = data_dict[k].reshape((original_shape[0], original_shape[1], original_shape[2], 44 | current_shape[-2], current_shape[-1])) 45 | return data_dict 46 | -------------------------------------------------------------------------------- /nnunetv2/training/dataloading/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/dataloading/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/dataloading/data_loader_3d.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from nnunetv2.training.dataloading.base_data_loader import nnUNetDataLoaderBase 3 | from nnunetv2.training.dataloading.nnunet_dataset import nnUNetDataset 4 | 5 | 6 | class nnUNetDataLoader3D(nnUNetDataLoaderBase): 7 | def generate_train_batch(self): 8 | selected_keys = self.get_indices() 9 | # preallocate memory for data and seg 10 | data_all = np.zeros(self.data_shape, dtype=np.float32) 11 | seg_all = np.zeros(self.seg_shape, dtype=np.int16) 12 | case_properties = [] 13 | 14 | for j, i in enumerate(selected_keys): 15 | # oversampling foreground will improve stability of model training, especially if many patches are empty 16 | # (Lung for example) 17 | force_fg = self.get_do_oversample(j) 18 | 19 | data, seg, properties = self._data.load_case(i) 20 | 21 | # If we are doing the cascade then the segmentation from the previous stage will already have been loaded by 22 | # self._data.load_case(i) (see nnUNetDataset.load_case) 23 | shape = data.shape[1:] 24 | dim = len(shape) 25 | bbox_lbs, bbox_ubs = self.get_bbox(shape, force_fg, properties['class_locations']) 26 | 27 | # whoever wrote this knew what he was doing (hint: it was me). We first crop the data to the region of the 28 | # bbox that actually lies within the data. This will result in a smaller array which is then faster to pad. 29 | # valid_bbox is just the coord that lied within the data cube. It will be padded to match the patch size 30 | # later 31 | valid_bbox_lbs = [max(0, bbox_lbs[i]) for i in range(dim)] 32 | valid_bbox_ubs = [min(shape[i], bbox_ubs[i]) for i in range(dim)] 33 | 34 | # At this point you might ask yourself why we would treat seg differently from seg_from_previous_stage. 35 | # Why not just concatenate them here and forget about the if statements? Well that's because segneeds to 36 | # be padded with -1 constant whereas seg_from_previous_stage needs to be padded with 0s (we could also 37 | # remove label -1 in the data augmentation but this way it is less error prone) 38 | this_slice = tuple([slice(0, data.shape[0])] + [slice(i, j) for i, j in zip(valid_bbox_lbs, valid_bbox_ubs)]) 39 | data = data[this_slice] 40 | 41 | this_slice = tuple([slice(0, seg.shape[0])] + [slice(i, j) for i, j in zip(valid_bbox_lbs, valid_bbox_ubs)]) 42 | seg = seg[this_slice] 43 | 44 | padding = [(-min(0, bbox_lbs[i]), max(bbox_ubs[i] - shape[i], 0)) for i in range(dim)] 45 | data_all[j] = np.pad(data, ((0, 0), *padding), 'constant', constant_values=0) 46 | seg_all[j] = np.pad(seg, ((0, 0), *padding), 'constant', constant_values=-1) 47 | 48 | return {'data': data_all, 'seg': seg_all, 'properties': case_properties, 'keys': selected_keys} 49 | 50 | 51 | if __name__ == '__main__': 52 | folder = '/media/fabian/data/nnUNet_preprocessed/Dataset002_Heart/3d_fullres' 53 | ds = nnUNetDataset(folder, 0) # this should not load the properties! 54 | dl = nnUNetDataLoader3D(ds, 5, (16, 16, 16), (16, 16, 16), 0.33, None, None) 55 | a = next(dl) 56 | -------------------------------------------------------------------------------- /nnunetv2/training/dataloading/utils.py: -------------------------------------------------------------------------------- 1 | import multiprocessing 2 | import os 3 | from multiprocessing import Pool 4 | from typing import List 5 | 6 | import numpy as np 7 | from batchgenerators.utilities.file_and_folder_operations import isfile, subfiles 8 | from nnunetv2.configuration import default_num_processes 9 | 10 | 11 | def _convert_to_npy(npz_file: str, unpack_segmentation: bool = True, overwrite_existing: bool = False) -> None: 12 | try: 13 | a = np.load(npz_file) # inexpensive, no compression is done here. This just reads metadata 14 | if overwrite_existing or not isfile(npz_file[:-3] + "npy"): 15 | np.save(npz_file[:-3] + "npy", a['data']) 16 | if unpack_segmentation and (overwrite_existing or not isfile(npz_file[:-4] + "_seg.npy")): 17 | np.save(npz_file[:-4] + "_seg.npy", a['seg']) 18 | except KeyboardInterrupt: 19 | if isfile(npz_file[:-3] + "npy"): 20 | os.remove(npz_file[:-3] + "npy") 21 | if isfile(npz_file[:-4] + "_seg.npy"): 22 | os.remove(npz_file[:-4] + "_seg.npy") 23 | raise KeyboardInterrupt 24 | 25 | 26 | def unpack_dataset(folder: str, unpack_segmentation: bool = True, overwrite_existing: bool = False, 27 | num_processes: int = default_num_processes): 28 | """ 29 | all npz files in this folder belong to the dataset, unpack them all 30 | """ 31 | with multiprocessing.get_context("spawn").Pool(num_processes) as p: 32 | npz_files = subfiles(folder, True, None, ".npz", True) 33 | p.starmap(_convert_to_npy, zip(npz_files, 34 | [unpack_segmentation] * len(npz_files), 35 | [overwrite_existing] * len(npz_files)) 36 | ) 37 | 38 | 39 | def get_case_identifiers(folder: str) -> List[str]: 40 | """ 41 | finds all npz files in the given folder and reconstructs the training case names from them 42 | """ 43 | case_identifiers = [i[:-4] for i in os.listdir(folder) if i.endswith("npz") and (i.find("segFromPrevStage") == -1)] 44 | return case_identifiers 45 | 46 | 47 | if __name__ == '__main__': 48 | unpack_dataset('/media/fabian/data/nnUNet_preprocessed/Dataset002_Heart/2d') -------------------------------------------------------------------------------- /nnunetv2/training/logging/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/logging/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/loss/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/loss/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/loss/deep_supervision.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | 4 | class DeepSupervisionWrapper(nn.Module): 5 | def __init__(self, loss, weight_factors=None): 6 | """ 7 | Wraps a loss function so that it can be applied to multiple outputs. Forward accepts an arbitrary number of 8 | inputs. Each input is expected to be a tuple/list. Each tuple/list must have the same length. The loss is then 9 | applied to each entry like this: 10 | l = w0 * loss(input0[0], input1[0], ...) + w1 * loss(input0[1], input1[1], ...) + ... 11 | If weights are None, all w will be 1. 12 | """ 13 | super(DeepSupervisionWrapper, self).__init__() 14 | self.weight_factors = weight_factors 15 | self.loss = loss 16 | 17 | def forward(self, *args): 18 | for i in args: 19 | assert isinstance(i, (tuple, list)), "all args must be either tuple or list, got %s" % type(i) 20 | # we could check for equal lengths here as well but we really shouldn't overdo it with checks because 21 | # this code is executed a lot of times! 22 | 23 | if self.weight_factors is None: 24 | weights = [1] * len(args[0]) 25 | else: 26 | weights = self.weight_factors 27 | 28 | # we initialize the loss like this instead of 0 to ensure it sits on the correct device, not sure if that's 29 | # really necessary 30 | l = weights[0] * self.loss(*[j[0] for j in args]) 31 | for i, inputs in enumerate(zip(*args)): 32 | if i == 0: 33 | continue 34 | l += weights[i] * self.loss(*inputs) 35 | return l -------------------------------------------------------------------------------- /nnunetv2/training/loss/robust_ce_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, Tensor 3 | import numpy as np 4 | 5 | 6 | class RobustCrossEntropyLoss(nn.CrossEntropyLoss): 7 | """ 8 | this is just a compatibility layer because my target tensor is float and has an extra dimension 9 | 10 | input must be logits, not probabilities! 11 | """ 12 | def forward(self, input: Tensor, target: Tensor) -> Tensor: 13 | if len(target.shape) == len(input.shape): 14 | assert target.shape[1] == 1 15 | target = target[:, 0] 16 | return super().forward(input, target.long()) 17 | 18 | 19 | class TopKLoss(RobustCrossEntropyLoss): 20 | """ 21 | input must be logits, not probabilities! 22 | """ 23 | def __init__(self, weight=None, ignore_index: int = -100, k: float = 10, label_smoothing: float = 0): 24 | self.k = k 25 | super(TopKLoss, self).__init__(weight, False, ignore_index, reduce=False, label_smoothing=label_smoothing) 26 | 27 | def forward(self, inp, target): 28 | target = target[:, 0].long() 29 | res = super(TopKLoss, self).forward(inp, target) 30 | num_voxels = np.prod(res.shape, dtype=np.int64) 31 | res, _ = torch.topk(res.view((-1, )), int(num_voxels * self.k / 100), sorted=False) 32 | return res.mean() 33 | 34 | -------------------------------------------------------------------------------- /nnunetv2/training/lr_scheduler/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/lr_scheduler/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/lr_scheduler/polylr.py: -------------------------------------------------------------------------------- 1 | from torch.optim.lr_scheduler import _LRScheduler 2 | 3 | 4 | class PolyLRScheduler(_LRScheduler): 5 | def __init__(self, optimizer, initial_lr: float, max_steps: int, exponent: float = 0.9, current_step: int = None): 6 | self.optimizer = optimizer 7 | self.initial_lr = initial_lr 8 | self.max_steps = max_steps 9 | self.exponent = exponent 10 | self.ctr = 0 11 | super().__init__(optimizer, current_step if current_step is not None else -1, False) 12 | 13 | def step(self, current_step=None): 14 | if current_step is None or current_step == -1: 15 | current_step = self.ctr 16 | self.ctr += 1 17 | 18 | new_lr = self.initial_lr * (1 - current_step / self.max_steps) ** self.exponent 19 | for param_group in self.optimizer.param_groups: 20 | param_group['lr'] = new_lr 21 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/variants/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/benchmarking/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/variants/benchmarking/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/benchmarking/nnUNetTrainerBenchmark_5epochs.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from batchgenerators.utilities.file_and_folder_operations import save_json, join, isfile, load_json 3 | 4 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 5 | from torch import distributed as dist 6 | 7 | 8 | class nnUNetTrainerBenchmark_5epochs(nnUNetTrainer): 9 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 10 | device: torch.device = torch.device('cuda')): 11 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 12 | assert self.fold == 0, "It makes absolutely no sense to specify a certain fold. Stick with 0 so that we can parse the results." 13 | self.disable_checkpointing = True 14 | self.num_epochs = 5 15 | assert torch.cuda.is_available(), "This only works on GPU" 16 | self.crashed_with_runtime_error = False 17 | 18 | def perform_actual_validation(self, save_probabilities: bool = False): 19 | pass 20 | 21 | def save_checkpoint(self, filename: str) -> None: 22 | # do not trust people to remember that self.disable_checkpointing must be True for this trainer 23 | pass 24 | 25 | def run_training(self): 26 | try: 27 | super().run_training() 28 | except RuntimeError: 29 | self.crashed_with_runtime_error = True 30 | 31 | def on_train_end(self): 32 | super().on_train_end() 33 | 34 | if not self.is_ddp or self.local_rank == 0: 35 | torch_version = torch.__version__ 36 | cudnn_version = torch.backends.cudnn.version() 37 | gpu_name = torch.cuda.get_device_name() 38 | if self.crashed_with_runtime_error: 39 | fastest_epoch = 'Not enough VRAM!' 40 | else: 41 | epoch_times = [i - j for i, j in zip(self.logger.my_fantastic_logging['epoch_end_timestamps'], 42 | self.logger.my_fantastic_logging['epoch_start_timestamps'])] 43 | fastest_epoch = min(epoch_times) 44 | 45 | if self.is_ddp: 46 | num_gpus = dist.get_world_size() 47 | else: 48 | num_gpus = 1 49 | 50 | benchmark_result_file = join(self.output_folder, 'benchmark_result.json') 51 | if isfile(benchmark_result_file): 52 | old_results = load_json(benchmark_result_file) 53 | else: 54 | old_results = {} 55 | # generate some unique key 56 | my_key = f"{cudnn_version}__{torch_version.replace(' ', '')}__{gpu_name.replace(' ', '')}__gpus_{num_gpus}" 57 | old_results[my_key] = { 58 | 'torch_version': torch_version, 59 | 'cudnn_version': cudnn_version, 60 | 'gpu_name': gpu_name, 61 | 'fastest_epoch': fastest_epoch, 62 | 'num_gpus': num_gpus, 63 | } 64 | save_json(old_results, 65 | join(self.output_folder, 'benchmark_result.json')) 66 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/benchmarking/nnUNetTrainerBenchmark_5epochs_noDataLoading.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from nnunetv2.training.nnUNetTrainer.variants.benchmarking.nnUNetTrainerBenchmark_5epochs import \ 4 | nnUNetTrainerBenchmark_5epochs 5 | from nnunetv2.utilities.label_handling.label_handling import determine_num_input_channels 6 | 7 | 8 | class nnUNetTrainerBenchmark_5epochs_noDataLoading(nnUNetTrainerBenchmark_5epochs): 9 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 10 | device: torch.device = torch.device('cuda')): 11 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 12 | self._set_batch_size_and_oversample() 13 | num_input_channels = determine_num_input_channels(self.plans_manager, self.configuration_manager, 14 | self.dataset_json) 15 | patch_size = self.configuration_manager.patch_size 16 | dummy_data = torch.rand((self.batch_size, num_input_channels, *patch_size), device=self.device) 17 | dummy_target = [ 18 | torch.round( 19 | torch.rand((self.batch_size, 1, *[int(i * j) for i, j in zip(patch_size, k)]), device=self.device) * 20 | max(self.label_manager.all_labels) 21 | ) for k in self._get_deep_supervision_scales()] 22 | self.dummy_batch = {'data': dummy_data, 'target': dummy_target} 23 | 24 | def get_dataloaders(self): 25 | return None, None 26 | 27 | def run_training(self): 28 | try: 29 | self.on_train_start() 30 | 31 | for epoch in range(self.current_epoch, self.num_epochs): 32 | self.on_epoch_start() 33 | 34 | self.on_train_epoch_start() 35 | train_outputs = [] 36 | for batch_id in range(self.num_iterations_per_epoch): 37 | train_outputs.append(self.train_step(self.dummy_batch)) 38 | self.on_train_epoch_end(train_outputs) 39 | 40 | with torch.no_grad(): 41 | self.on_validation_epoch_start() 42 | val_outputs = [] 43 | for batch_id in range(self.num_val_iterations_per_epoch): 44 | val_outputs.append(self.validation_step(self.dummy_batch)) 45 | self.on_validation_epoch_end(val_outputs) 46 | 47 | self.on_epoch_end() 48 | 49 | self.on_train_end() 50 | except RuntimeError: 51 | self.crashed_with_runtime_error = True 52 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/data_augmentation/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/variants/data_augmentation/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerNoDA.py: -------------------------------------------------------------------------------- 1 | from typing import Union, Tuple, List 2 | 3 | from batchgenerators.transforms.abstract_transforms import AbstractTransform 4 | 5 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 6 | import numpy as np 7 | 8 | 9 | class nnUNetTrainerNoDA(nnUNetTrainer): 10 | @staticmethod 11 | def get_training_transforms(patch_size: Union[np.ndarray, Tuple[int]], 12 | rotation_for_DA: dict, 13 | deep_supervision_scales: Union[List, Tuple], 14 | mirror_axes: Tuple[int, ...], 15 | do_dummy_2d_data_aug: bool, 16 | order_resampling_data: int = 1, 17 | order_resampling_seg: int = 0, 18 | border_val_seg: int = -1, 19 | use_mask_for_norm: List[bool] = None, 20 | is_cascaded: bool = False, 21 | foreground_labels: Union[Tuple[int, ...], List[int]] = None, 22 | regions: List[Union[List[int], Tuple[int, ...], int]] = None, 23 | ignore_label: int = None) -> AbstractTransform: 24 | return nnUNetTrainer.get_validation_transforms(deep_supervision_scales, is_cascaded, foreground_labels, 25 | regions, ignore_label) 26 | 27 | def get_plain_dataloaders(self, initial_patch_size: Tuple[int, ...], dim: int): 28 | return super().get_plain_dataloaders( 29 | initial_patch_size=self.configuration_manager.patch_size, 30 | dim=dim 31 | ) 32 | 33 | def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): 34 | # we need to disable mirroring here so that no mirroring will be applied in inferene! 35 | rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \ 36 | super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() 37 | mirror_axes = None 38 | self.inference_allowed_mirroring_axes = None 39 | return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes 40 | 41 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerNoMirroring.py: -------------------------------------------------------------------------------- 1 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 2 | 3 | 4 | class nnUNetTrainerNoMirroring(nnUNetTrainer): 5 | def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): 6 | rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \ 7 | super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() 8 | mirror_axes = None 9 | self.inference_allowed_mirroring_axes = None 10 | return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes 11 | 12 | 13 | class nnUNetTrainer_onlyMirror01(nnUNetTrainer): 14 | """ 15 | Only mirrors along spatial axes 0 and 1 for 3D and 0 for 2D 16 | """ 17 | def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): 18 | rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \ 19 | super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() 20 | patch_size = self.configuration_manager.patch_size 21 | dim = len(patch_size) 22 | if dim == 2: 23 | mirror_axes = (0, ) 24 | else: 25 | mirror_axes = (0, 1) 26 | self.inference_allowed_mirroring_axes = mirror_axes 27 | return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes 28 | 29 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/loss/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/variants/loss/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/loss/nnUNetTrainerCELoss.py: -------------------------------------------------------------------------------- 1 | from nnunetv2.training.loss.deep_supervision import DeepSupervisionWrapper 2 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 3 | from nnunetv2.training.loss.robust_ce_loss import RobustCrossEntropyLoss 4 | import numpy as np 5 | 6 | 7 | class nnUNetTrainerCELoss(nnUNetTrainer): 8 | def _build_loss(self): 9 | assert not self.label_manager.has_regions, 'regions not supported by this trainer' 10 | loss = RobustCrossEntropyLoss(weight=None, 11 | ignore_index=self.label_manager.ignore_label if self.label_manager.has_ignore_label else -100) 12 | 13 | deep_supervision_scales = self._get_deep_supervision_scales() 14 | 15 | # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases 16 | # this gives higher resolution outputs more weight in the loss 17 | weights = np.array([1 / (2 ** i) for i in range(len(deep_supervision_scales))]) 18 | 19 | # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1 20 | weights = weights / weights.sum() 21 | # now wrap the loss 22 | loss = DeepSupervisionWrapper(loss, weights) 23 | return loss 24 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/loss/nnUNetTrainerDiceLoss.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from nnunetv2.training.loss.compound_losses import DC_and_BCE_loss, DC_and_CE_loss 5 | from nnunetv2.training.loss.deep_supervision import DeepSupervisionWrapper 6 | from nnunetv2.training.loss.dice import MemoryEfficientSoftDiceLoss 7 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 8 | from nnunetv2.utilities.helpers import softmax_helper_dim1 9 | 10 | 11 | class nnUNetTrainerDiceLoss(nnUNetTrainer): 12 | def _build_loss(self): 13 | loss = MemoryEfficientSoftDiceLoss(**{'batch_dice': self.configuration_manager.batch_dice, 14 | 'do_bg': self.label_manager.has_regions, 'smooth': 1e-5, 'ddp': self.is_ddp}, 15 | apply_nonlin=torch.sigmoid if self.label_manager.has_regions else softmax_helper_dim1) 16 | 17 | deep_supervision_scales = self._get_deep_supervision_scales() 18 | 19 | # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases 20 | # this gives higher resolution outputs more weight in the loss 21 | weights = np.array([1 / (2 ** i) for i in range(len(deep_supervision_scales))]) 22 | 23 | # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1 24 | weights = weights / weights.sum() 25 | # now wrap the loss 26 | loss = DeepSupervisionWrapper(loss, weights) 27 | return loss 28 | 29 | 30 | class nnUNetTrainerDiceCELoss_noSmooth(nnUNetTrainer): 31 | def _build_loss(self): 32 | # set smooth to 0 33 | if self.label_manager.has_regions: 34 | loss = DC_and_BCE_loss({}, 35 | {'batch_dice': self.configuration_manager.batch_dice, 36 | 'do_bg': True, 'smooth': 0, 'ddp': self.is_ddp}, 37 | use_ignore_label=self.label_manager.ignore_label is not None, 38 | dice_class=MemoryEfficientSoftDiceLoss) 39 | else: 40 | loss = DC_and_CE_loss({'batch_dice': self.configuration_manager.batch_dice, 41 | 'smooth': 0, 'do_bg': False, 'ddp': self.is_ddp}, {}, weight_ce=1, weight_dice=1, 42 | ignore_label=self.label_manager.ignore_label, 43 | dice_class=MemoryEfficientSoftDiceLoss) 44 | 45 | deep_supervision_scales = self._get_deep_supervision_scales() 46 | 47 | # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases 48 | # this gives higher resolution outputs more weight in the loss 49 | weights = np.array([1 / (2 ** i) for i in range(len(deep_supervision_scales))]) 50 | 51 | # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1 52 | weights = weights / weights.sum() 53 | # now wrap the loss 54 | loss = DeepSupervisionWrapper(loss, weights) 55 | return loss 56 | 57 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/loss/nnUNetTrainerTopkLoss.py: -------------------------------------------------------------------------------- 1 | from nnunetv2.training.loss.compound_losses import DC_and_topk_loss 2 | from nnunetv2.training.loss.deep_supervision import DeepSupervisionWrapper 3 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 4 | import numpy as np 5 | from nnunetv2.training.loss.robust_ce_loss import TopKLoss 6 | 7 | 8 | class nnUNetTrainerTopk10Loss(nnUNetTrainer): 9 | def _build_loss(self): 10 | assert not self.label_manager.has_regions, 'regions not supported by this trainer' 11 | loss = TopKLoss(ignore_index=self.label_manager.ignore_label if self.label_manager.has_ignore_label else -100, 12 | k=10) 13 | 14 | deep_supervision_scales = self._get_deep_supervision_scales() 15 | 16 | # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases 17 | # this gives higher resolution outputs more weight in the loss 18 | weights = np.array([1 / (2 ** i) for i in range(len(deep_supervision_scales))]) 19 | 20 | # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1 21 | weights = weights / weights.sum() 22 | # now wrap the loss 23 | loss = DeepSupervisionWrapper(loss, weights) 24 | return loss 25 | 26 | 27 | class nnUNetTrainerTopk10LossLS01(nnUNetTrainer): 28 | def _build_loss(self): 29 | assert not self.label_manager.has_regions, 'regions not supported by this trainer' 30 | loss = TopKLoss(ignore_index=self.label_manager.ignore_label if self.label_manager.has_ignore_label else -100, 31 | k=10, label_smoothing=0.1) 32 | 33 | deep_supervision_scales = self._get_deep_supervision_scales() 34 | 35 | # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases 36 | # this gives higher resolution outputs more weight in the loss 37 | weights = np.array([1 / (2 ** i) for i in range(len(deep_supervision_scales))]) 38 | 39 | # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1 40 | weights = weights / weights.sum() 41 | # now wrap the loss 42 | loss = DeepSupervisionWrapper(loss, weights) 43 | return loss 44 | 45 | 46 | class nnUNetTrainerDiceTopK10Loss(nnUNetTrainer): 47 | def _build_loss(self): 48 | assert not self.label_manager.has_regions, 'regions not supported by this trainer' 49 | loss = DC_and_topk_loss({'batch_dice': self.configuration_manager.batch_dice, 50 | 'smooth': 1e-5, 'do_bg': False, 'ddp': self.is_ddp}, 51 | {'k': 10, 52 | 'label_smoothing': 0.0}, 53 | weight_ce=1, weight_dice=1, 54 | ignore_label=self.label_manager.ignore_label) 55 | 56 | deep_supervision_scales = self._get_deep_supervision_scales() 57 | 58 | # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases 59 | # this gives higher resolution outputs more weight in the loss 60 | weights = np.array([1 / (2 ** i) for i in range(len(deep_supervision_scales))]) 61 | 62 | # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1 63 | weights = weights / weights.sum() 64 | # now wrap the loss 65 | loss = DeepSupervisionWrapper(loss, weights) 66 | return loss 67 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/lr_schedule/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/variants/lr_schedule/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/lr_schedule/nnUNetTrainerCosAnneal.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.optim.lr_scheduler import CosineAnnealingLR 3 | 4 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 5 | 6 | 7 | class nnUNetTrainerCosAnneal(nnUNetTrainer): 8 | def configure_optimizers(self): 9 | optimizer = torch.optim.SGD(self.network.parameters(), self.initial_lr, weight_decay=self.weight_decay, 10 | momentum=0.99, nesterov=True) 11 | lr_scheduler = CosineAnnealingLR(optimizer, T_max=self.num_epochs) 12 | return optimizer, lr_scheduler 13 | 14 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/network_architecture/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/variants/network_architecture/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/network_architecture/nnUNetTrainerBN.py: -------------------------------------------------------------------------------- 1 | from dynamic_network_architectures.architectures.unet import ResidualEncoderUNet, PlainConvUNet 2 | from dynamic_network_architectures.building_blocks.helper import convert_dim_to_conv_op, get_matching_batchnorm 3 | from dynamic_network_architectures.initialization.weight_init import init_last_bn_before_add_to_0, InitWeights_He 4 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 5 | from nnunetv2.utilities.plans_handling.plans_handler import ConfigurationManager, PlansManager 6 | from torch import nn 7 | 8 | 9 | class nnUNetTrainerBN(nnUNetTrainer): 10 | @staticmethod 11 | def build_network_architecture(plans_manager: PlansManager, 12 | dataset_json, 13 | configuration_manager: ConfigurationManager, 14 | num_input_channels, 15 | enable_deep_supervision: bool = True) -> nn.Module: 16 | num_stages = len(configuration_manager.conv_kernel_sizes) 17 | 18 | dim = len(configuration_manager.conv_kernel_sizes[0]) 19 | conv_op = convert_dim_to_conv_op(dim) 20 | 21 | label_manager = plans_manager.get_label_manager(dataset_json) 22 | 23 | segmentation_network_class_name = configuration_manager.UNet_class_name 24 | mapping = { 25 | 'PlainConvUNet': PlainConvUNet, 26 | 'ResidualEncoderUNet': ResidualEncoderUNet 27 | } 28 | kwargs = { 29 | 'PlainConvUNet': { 30 | 'conv_bias': True, 31 | 'norm_op': get_matching_batchnorm(conv_op), 32 | 'norm_op_kwargs': {'eps': 1e-5, 'affine': True}, 33 | 'dropout_op': None, 'dropout_op_kwargs': None, 34 | 'nonlin': nn.LeakyReLU, 'nonlin_kwargs': {'inplace': True}, 35 | }, 36 | 'ResidualEncoderUNet': { 37 | 'conv_bias': True, 38 | 'norm_op': get_matching_batchnorm(conv_op), 39 | 'norm_op_kwargs': {'eps': 1e-5, 'affine': True}, 40 | 'dropout_op': None, 'dropout_op_kwargs': None, 41 | 'nonlin': nn.LeakyReLU, 'nonlin_kwargs': {'inplace': True}, 42 | } 43 | } 44 | assert segmentation_network_class_name in mapping.keys(), 'The network architecture specified by the plans file ' \ 45 | 'is non-standard (maybe your own?). Yo\'ll have to dive ' \ 46 | 'into either this ' \ 47 | 'function (get_network_from_plans) or ' \ 48 | 'the init of your nnUNetModule to accomodate that.' 49 | network_class = mapping[segmentation_network_class_name] 50 | 51 | conv_or_blocks_per_stage = { 52 | 'n_conv_per_stage' 53 | if network_class != ResidualEncoderUNet else 'n_blocks_per_stage': configuration_manager.n_conv_per_stage_encoder, 54 | 'n_conv_per_stage_decoder': configuration_manager.n_conv_per_stage_decoder 55 | } 56 | # network class name!! 57 | model = network_class( 58 | input_channels=num_input_channels, 59 | n_stages=num_stages, 60 | features_per_stage=[min(configuration_manager.UNet_base_num_features * 2 ** i, 61 | configuration_manager.unet_max_num_features) for i in range(num_stages)], 62 | conv_op=conv_op, 63 | kernel_sizes=configuration_manager.conv_kernel_sizes, 64 | strides=configuration_manager.pool_op_kernel_sizes, 65 | num_classes=label_manager.num_segmentation_heads, 66 | deep_supervision=enable_deep_supervision, 67 | **conv_or_blocks_per_stage, 68 | **kwargs[segmentation_network_class_name] 69 | ) 70 | model.apply(InitWeights_He(1e-2)) 71 | if network_class == ResidualEncoderUNet: 72 | model.apply(init_last_bn_before_add_to_0) 73 | return model 74 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/optimizer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/variants/optimizer/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/optimizer/nnUNetTrainerAdam.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.optim import Adam, AdamW 3 | 4 | from nnunetv2.training.lr_scheduler.polylr import PolyLRScheduler 5 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 6 | 7 | 8 | class nnUNetTrainerAdam(nnUNetTrainer): 9 | def configure_optimizers(self): 10 | optimizer = AdamW(self.network.parameters(), 11 | lr=self.initial_lr, 12 | weight_decay=self.weight_decay, 13 | amsgrad=True) 14 | # optimizer = torch.optim.SGD(self.network.parameters(), self.initial_lr, weight_decay=self.weight_decay, 15 | # momentum=0.99, nesterov=True) 16 | lr_scheduler = PolyLRScheduler(optimizer, self.initial_lr, self.num_epochs) 17 | return optimizer, lr_scheduler 18 | 19 | 20 | class nnUNetTrainerVanillaAdam(nnUNetTrainer): 21 | def configure_optimizers(self): 22 | optimizer = Adam(self.network.parameters(), 23 | lr=self.initial_lr, 24 | weight_decay=self.weight_decay) 25 | # optimizer = torch.optim.SGD(self.network.parameters(), self.initial_lr, weight_decay=self.weight_decay, 26 | # momentum=0.99, nesterov=True) 27 | lr_scheduler = PolyLRScheduler(optimizer, self.initial_lr, self.num_epochs) 28 | return optimizer, lr_scheduler 29 | 30 | 31 | class nnUNetTrainerVanillaAdam1en3(nnUNetTrainerVanillaAdam): 32 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 33 | device: torch.device = torch.device('cuda')): 34 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 35 | self.initial_lr = 1e-3 36 | 37 | 38 | class nnUNetTrainerVanillaAdam3en4(nnUNetTrainerVanillaAdam): 39 | # https://twitter.com/karpathy/status/801621764144971776?lang=en 40 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 41 | device: torch.device = torch.device('cuda')): 42 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 43 | self.initial_lr = 3e-4 44 | 45 | 46 | class nnUNetTrainerAdam1en3(nnUNetTrainerAdam): 47 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 48 | device: torch.device = torch.device('cuda')): 49 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 50 | self.initial_lr = 1e-3 51 | 52 | 53 | class nnUNetTrainerAdam3en4(nnUNetTrainerAdam): 54 | # https://twitter.com/karpathy/status/801621764144971776?lang=en 55 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 56 | device: torch.device = torch.device('cuda')): 57 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 58 | self.initial_lr = 3e-4 59 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/optimizer/nnUNetTrainerAdan.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from nnunetv2.training.lr_scheduler.polylr import PolyLRScheduler 4 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 5 | from torch.optim.lr_scheduler import CosineAnnealingLR 6 | try: 7 | from adan_pytorch import Adan 8 | except ImportError: 9 | Adan = None 10 | 11 | 12 | class nnUNetTrainerAdan(nnUNetTrainer): 13 | def configure_optimizers(self): 14 | if Adan is None: 15 | raise RuntimeError('This trainer requires adan_pytorch to be installed, install with "pip install adan-pytorch"') 16 | optimizer = Adan(self.network.parameters(), 17 | lr=self.initial_lr, 18 | # betas=(0.02, 0.08, 0.01), defaults 19 | weight_decay=self.weight_decay) 20 | # optimizer = torch.optim.SGD(self.network.parameters(), self.initial_lr, weight_decay=self.weight_decay, 21 | # momentum=0.99, nesterov=True) 22 | lr_scheduler = PolyLRScheduler(optimizer, self.initial_lr, self.num_epochs) 23 | return optimizer, lr_scheduler 24 | 25 | 26 | class nnUNetTrainerAdan1en3(nnUNetTrainerAdan): 27 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 28 | device: torch.device = torch.device('cuda')): 29 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 30 | self.initial_lr = 1e-3 31 | 32 | 33 | class nnUNetTrainerAdan3en4(nnUNetTrainerAdan): 34 | # https://twitter.com/karpathy/status/801621764144971776?lang=en 35 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 36 | device: torch.device = torch.device('cuda')): 37 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 38 | self.initial_lr = 3e-4 39 | 40 | 41 | class nnUNetTrainerAdan1en1(nnUNetTrainerAdan): 42 | # this trainer makes no sense -> nan! 43 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 44 | device: torch.device = torch.device('cuda')): 45 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 46 | self.initial_lr = 1e-1 47 | 48 | 49 | class nnUNetTrainerAdanCosAnneal(nnUNetTrainerAdan): 50 | # def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 51 | # device: torch.device = torch.device('cuda')): 52 | # super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 53 | # self.num_epochs = 15 54 | 55 | def configure_optimizers(self): 56 | if Adan is None: 57 | raise RuntimeError('This trainer requires adan_pytorch to be installed, install with "pip install adan-pytorch"') 58 | optimizer = Adan(self.network.parameters(), 59 | lr=self.initial_lr, 60 | # betas=(0.02, 0.08, 0.01), defaults 61 | weight_decay=self.weight_decay) 62 | # optimizer = torch.optim.SGD(self.network.parameters(), self.initial_lr, weight_decay=self.weight_decay, 63 | # momentum=0.99, nesterov=True) 64 | lr_scheduler = CosineAnnealingLR(optimizer, T_max=self.num_epochs) 65 | return optimizer, lr_scheduler 66 | 67 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/sampling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/variants/sampling/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/training_length/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/training/nnUNetTrainer/variants/training_length/__init__.py -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/training_length/nnUNetTrainer_Xepochs.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 4 | 5 | 6 | class nnUNetTrainer_5epochs(nnUNetTrainer): 7 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 8 | device: torch.device = torch.device('cuda')): 9 | """used for debugging plans etc""" 10 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 11 | self.num_epochs = 5 12 | 13 | 14 | class nnUNetTrainer_1epoch(nnUNetTrainer): 15 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 16 | device: torch.device = torch.device('cuda')): 17 | """used for debugging plans etc""" 18 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 19 | self.num_epochs = 1 20 | 21 | 22 | class nnUNetTrainer_10epochs(nnUNetTrainer): 23 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 24 | device: torch.device = torch.device('cuda')): 25 | """used for debugging plans etc""" 26 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 27 | self.num_epochs = 10 28 | 29 | 30 | class nnUNetTrainer_20epochs(nnUNetTrainer): 31 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 32 | device: torch.device = torch.device('cuda')): 33 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 34 | self.num_epochs = 20 35 | 36 | 37 | class nnUNetTrainer_250epochs(nnUNetTrainer): 38 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 39 | device: torch.device = torch.device('cuda')): 40 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 41 | self.num_epochs = 250 42 | 43 | 44 | class nnUNetTrainer_2000epochs(nnUNetTrainer): 45 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 46 | device: torch.device = torch.device('cuda')): 47 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 48 | self.num_epochs = 2000 49 | 50 | 51 | class nnUNetTrainer_4000epochs(nnUNetTrainer): 52 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 53 | device: torch.device = torch.device('cuda')): 54 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 55 | self.num_epochs = 4000 56 | 57 | 58 | class nnUNetTrainer_8000epochs(nnUNetTrainer): 59 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 60 | device: torch.device = torch.device('cuda')): 61 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 62 | self.num_epochs = 8000 63 | -------------------------------------------------------------------------------- /nnunetv2/training/nnUNetTrainer/variants/training_length/nnUNetTrainer_Xepochs_NoMirroring.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer 4 | 5 | 6 | class nnUNetTrainer_250epochs_NoMirroring(nnUNetTrainer): 7 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 8 | device: torch.device = torch.device('cuda')): 9 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 10 | self.num_epochs = 250 11 | 12 | def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): 13 | rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \ 14 | super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() 15 | mirror_axes = None 16 | self.inference_allowed_mirroring_axes = None 17 | return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes 18 | 19 | 20 | class nnUNetTrainer_2000epochs_NoMirroring(nnUNetTrainer): 21 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 22 | device: torch.device = torch.device('cuda')): 23 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 24 | self.num_epochs = 2000 25 | 26 | def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): 27 | rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \ 28 | super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() 29 | mirror_axes = None 30 | self.inference_allowed_mirroring_axes = None 31 | return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes 32 | 33 | 34 | class nnUNetTrainer_4000epochs_NoMirroring(nnUNetTrainer): 35 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 36 | device: torch.device = torch.device('cuda')): 37 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 38 | self.num_epochs = 4000 39 | 40 | def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): 41 | rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \ 42 | super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() 43 | mirror_axes = None 44 | self.inference_allowed_mirroring_axes = None 45 | return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes 46 | 47 | 48 | class nnUNetTrainer_8000epochs_NoMirroring(nnUNetTrainer): 49 | def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True, 50 | device: torch.device = torch.device('cuda')): 51 | super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device) 52 | self.num_epochs = 8000 53 | 54 | def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): 55 | rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \ 56 | super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() 57 | mirror_axes = None 58 | self.inference_allowed_mirroring_axes = None 59 | return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes 60 | 61 | -------------------------------------------------------------------------------- /nnunetv2/utilities/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/utilities/__init__.py -------------------------------------------------------------------------------- /nnunetv2/utilities/collate_outputs.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | import numpy as np 4 | 5 | 6 | def collate_outputs(outputs: List[dict]): 7 | """ 8 | used to collate default train_step and validation_step outputs. If you want something different then you gotta 9 | extend this 10 | 11 | we expect outputs to be a list of dictionaries where each of the dict has the same set of keys 12 | """ 13 | collated = {} 14 | for k in outputs[0].keys(): 15 | if np.isscalar(outputs[0][k]): 16 | collated[k] = [o[k] for o in outputs] 17 | elif isinstance(outputs[0][k], np.ndarray): 18 | collated[k] = np.vstack([o[k][None] for o in outputs]) 19 | elif isinstance(outputs[0][k], list): 20 | collated[k] = [item for o in outputs for item in o[k]] 21 | else: 22 | raise ValueError(f'Cannot collate input of type {type(outputs[0][k])}. ' 23 | f'Modify collate_outputs to add this functionality') 24 | return collated -------------------------------------------------------------------------------- /nnunetv2/utilities/dataset_name_id_conversion.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from typing import Union 15 | 16 | from nnunetv2.paths import nnUNet_preprocessed, nnUNet_raw, nnUNet_results 17 | from batchgenerators.utilities.file_and_folder_operations import * 18 | import numpy as np 19 | 20 | 21 | def find_candidate_datasets(dataset_id: int): 22 | startswith = "Dataset%03.0d" % dataset_id 23 | if nnUNet_preprocessed is not None and isdir(nnUNet_preprocessed): 24 | candidates_preprocessed = subdirs(nnUNet_preprocessed, prefix=startswith, join=False) 25 | else: 26 | candidates_preprocessed = [] 27 | 28 | if nnUNet_raw is not None and isdir(nnUNet_raw): 29 | candidates_raw = subdirs(nnUNet_raw, prefix=startswith, join=False) 30 | else: 31 | candidates_raw = [] 32 | 33 | candidates_trained_models = [] 34 | if nnUNet_results is not None and isdir(nnUNet_results): 35 | candidates_trained_models += subdirs(nnUNet_results, prefix=startswith, join=False) 36 | 37 | all_candidates = candidates_preprocessed + candidates_raw + candidates_trained_models 38 | unique_candidates = np.unique(all_candidates) 39 | return unique_candidates 40 | 41 | 42 | def convert_id_to_dataset_name(dataset_id: int): 43 | unique_candidates = find_candidate_datasets(dataset_id) 44 | if len(unique_candidates) > 1: 45 | raise RuntimeError("More than one dataset name found for dataset id %d. Please correct that. (I looked in the " 46 | "following folders:\n%s\n%s\n%s" % (dataset_id, nnUNet_raw, nnUNet_preprocessed, nnUNet_results)) 47 | if len(unique_candidates) == 0: 48 | raise RuntimeError(f"Could not find a dataset with the ID {dataset_id}. Make sure the requested dataset ID " 49 | f"exists and that nnU-Net knows where raw and preprocessed data are located " 50 | f"(see Documentation - Installation). Here are your currently defined folders:\n" 51 | f"nnUNet_preprocessed={os.environ.get('nnUNet_preprocessed') if os.environ.get('nnUNet_preprocessed') is not None else 'None'}\n" 52 | f"nnUNet_results={os.environ.get('nnUNet_results') if os.environ.get('nnUNet_results') is not None else 'None'}\n" 53 | f"nnUNet_raw={os.environ.get('nnUNet_raw') if os.environ.get('nnUNet_raw') is not None else 'None'}\n" 54 | f"If something is not right, adapt your environment variables.") 55 | return unique_candidates[0] 56 | 57 | 58 | def convert_dataset_name_to_id(dataset_name: str): 59 | assert dataset_name.startswith("Dataset") 60 | dataset_id = int(dataset_name[7:10]) 61 | return dataset_id 62 | 63 | 64 | def maybe_convert_to_dataset_name(dataset_name_or_id: Union[int, str]) -> str: 65 | if isinstance(dataset_name_or_id, str) and dataset_name_or_id.startswith("Dataset"): 66 | return dataset_name_or_id 67 | if isinstance(dataset_name_or_id, str): 68 | try: 69 | dataset_name_or_id = int(dataset_name_or_id) 70 | except ValueError: 71 | raise ValueError("dataset_name_or_id was a string and did not start with 'Dataset' so we tried to " 72 | "convert it to a dataset ID (int). That failed, however. Please give an integer number " 73 | "('1', '2', etc) or a correct tast name. Your input: %s" % dataset_name_or_id) 74 | return convert_id_to_dataset_name(dataset_name_or_id) -------------------------------------------------------------------------------- /nnunetv2/utilities/ddp_allgather.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from typing import Any, Optional, Tuple 15 | 16 | import torch 17 | from torch import distributed 18 | 19 | 20 | def print_if_rank0(*args): 21 | if distributed.get_rank() == 0: 22 | print(*args) 23 | 24 | 25 | class AllGatherGrad(torch.autograd.Function): 26 | # stolen from pytorch lightning 27 | @staticmethod 28 | def forward( 29 | ctx: Any, 30 | tensor: torch.Tensor, 31 | group: Optional["torch.distributed.ProcessGroup"] = None, 32 | ) -> torch.Tensor: 33 | ctx.group = group 34 | 35 | gathered_tensor = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())] 36 | 37 | torch.distributed.all_gather(gathered_tensor, tensor, group=group) 38 | gathered_tensor = torch.stack(gathered_tensor, dim=0) 39 | 40 | return gathered_tensor 41 | 42 | @staticmethod 43 | def backward(ctx: Any, *grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: 44 | grad_output = torch.cat(grad_output) 45 | 46 | torch.distributed.all_reduce(grad_output, op=torch.distributed.ReduceOp.SUM, async_op=False, group=ctx.group) 47 | 48 | return grad_output[torch.distributed.get_rank()], None 49 | 50 | -------------------------------------------------------------------------------- /nnunetv2/utilities/default_n_proc_DA.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import os 3 | 4 | 5 | def get_allowed_n_proc_DA(): 6 | """ 7 | This function is used to set the number of processes used on different Systems. It is specific to our cluster 8 | infrastructure at DKFZ. You can modify it to suit your needs. Everything is allowed. 9 | 10 | IMPORTANT: if the environment variable nnUNet_n_proc_DA is set it will overwrite anything in this script 11 | (see first line). 12 | 13 | Interpret the output as the number of processes used for data augmentation PER GPU. 14 | 15 | The way it is implemented here is simply a look up table. We know the hostnames, CPU and GPU configurations of our 16 | systems and set the numbers accordingly. For example, a system with 4 GPUs and 48 threads can use 12 threads per 17 | GPU without overloading the CPU (technically 11 because we have a main process as well), so that's what we use. 18 | """ 19 | 20 | if 'nnUNet_n_proc_DA' in os.environ.keys(): 21 | use_this = int(os.environ['nnUNet_n_proc_DA']) 22 | else: 23 | hostname = subprocess.getoutput(['hostname']) 24 | if hostname in ['Fabian', ]: 25 | use_this = 12 26 | elif hostname in ['hdf19-gpu16', 'hdf19-gpu17', 'hdf19-gpu18', 'hdf19-gpu19', 'e230-AMDworkstation']: 27 | use_this = 16 28 | elif hostname.startswith('e230-dgx1'): 29 | use_this = 10 30 | elif hostname.startswith('hdf18-gpu') or hostname.startswith('e132-comp'): 31 | use_this = 16 32 | elif hostname.startswith('e230-dgx2'): 33 | use_this = 6 34 | elif hostname.startswith('e230-dgxa100-'): 35 | use_this = 28 36 | elif hostname.startswith('lsf22-gpu'): 37 | use_this = 28 38 | elif hostname.startswith('hdf19-gpu') or hostname.startswith('e071-gpu'): 39 | use_this = 12 40 | else: 41 | use_this = 12 # default value 42 | 43 | use_this = min(use_this, os.cpu_count()) 44 | return use_this 45 | -------------------------------------------------------------------------------- /nnunetv2/utilities/find_class_by_name.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | import pkgutil 3 | 4 | from batchgenerators.utilities.file_and_folder_operations import * 5 | 6 | 7 | def recursive_find_python_class(folder: str, class_name: str, current_module: str): 8 | tr = None 9 | for importer, modname, ispkg in pkgutil.iter_modules([folder]): 10 | # print(modname, ispkg) 11 | if not ispkg: 12 | m = importlib.import_module(current_module + "." + modname) 13 | if hasattr(m, class_name): 14 | tr = getattr(m, class_name) 15 | break 16 | 17 | if tr is None: 18 | for importer, modname, ispkg in pkgutil.iter_modules([folder]): 19 | if ispkg: 20 | next_current_module = current_module + "." + modname 21 | tr = recursive_find_python_class(join(folder, modname), class_name, current_module=next_current_module) 22 | if tr is not None: 23 | break 24 | return tr -------------------------------------------------------------------------------- /nnunetv2/utilities/helpers.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def softmax_helper_dim0(x: torch.Tensor) -> torch.Tensor: 5 | return torch.softmax(x, 0) 6 | 7 | 8 | def softmax_helper_dim1(x: torch.Tensor) -> torch.Tensor: 9 | return torch.softmax(x, 1) 10 | 11 | 12 | def empty_cache(device: torch.device): 13 | if device.type == 'cuda': 14 | torch.cuda.empty_cache() 15 | elif device.type == 'mps': 16 | from torch import mps 17 | mps.empty_cache() 18 | else: 19 | pass 20 | 21 | 22 | class dummy_context(object): 23 | def __enter__(self): 24 | pass 25 | 26 | def __exit__(self, exc_type, exc_val, exc_tb): 27 | pass 28 | -------------------------------------------------------------------------------- /nnunetv2/utilities/json_export.py: -------------------------------------------------------------------------------- 1 | from collections.abc import Iterable 2 | 3 | import numpy as np 4 | import torch 5 | 6 | 7 | def recursive_fix_for_json_export(my_dict: dict): 8 | # json is stupid. 'cannot serialize object of type bool_/int64/float64'. Come on bro. 9 | keys = list(my_dict.keys()) # cannot iterate over keys() if we change keys.... 10 | for k in keys: 11 | if isinstance(k, (np.int64, np.int32, np.int8, np.uint8)): 12 | tmp = my_dict[k] 13 | del my_dict[k] 14 | my_dict[int(k)] = tmp 15 | del tmp 16 | k = int(k) 17 | 18 | if isinstance(my_dict[k], dict): 19 | recursive_fix_for_json_export(my_dict[k]) 20 | elif isinstance(my_dict[k], np.ndarray): 21 | assert len(my_dict[k].shape) == 1, 'only 1d arrays are supported' 22 | my_dict[k] = fix_types_iterable(my_dict[k], output_type=list) 23 | elif isinstance(my_dict[k], (np.bool_,)): 24 | my_dict[k] = bool(my_dict[k]) 25 | elif isinstance(my_dict[k], (np.int64, np.int32, np.int8, np.uint8)): 26 | my_dict[k] = int(my_dict[k]) 27 | elif isinstance(my_dict[k], (np.float32, np.float64, np.float16)): 28 | my_dict[k] = float(my_dict[k]) 29 | elif isinstance(my_dict[k], list): 30 | my_dict[k] = fix_types_iterable(my_dict[k], output_type=type(my_dict[k])) 31 | elif isinstance(my_dict[k], tuple): 32 | my_dict[k] = fix_types_iterable(my_dict[k], output_type=tuple) 33 | elif isinstance(my_dict[k], torch.device): 34 | my_dict[k] = str(my_dict[k]) 35 | else: 36 | pass # pray it can be serialized 37 | 38 | 39 | def fix_types_iterable(iterable, output_type): 40 | # this sh!t is hacky as hell and will break if you use it for anything outside nnunet. Keep you hands off of this. 41 | out = [] 42 | for i in iterable: 43 | if type(i) in (np.int64, np.int32, np.int8, np.uint8): 44 | out.append(int(i)) 45 | elif isinstance(i, dict): 46 | recursive_fix_for_json_export(i) 47 | out.append(i) 48 | elif type(i) in (np.float32, np.float64, np.float16): 49 | out.append(float(i)) 50 | elif type(i) in (np.bool_,): 51 | out.append(bool(i)) 52 | elif isinstance(i, str): 53 | out.append(i) 54 | elif isinstance(i, Iterable): 55 | # print('recursive call on', i, type(i)) 56 | out.append(fix_types_iterable(i, type(i))) 57 | else: 58 | out.append(i) 59 | return output_type(out) 60 | -------------------------------------------------------------------------------- /nnunetv2/utilities/label_handling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/utilities/label_handling/__init__.py -------------------------------------------------------------------------------- /nnunetv2/utilities/network_initialization.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | 4 | class InitWeights_He(object): 5 | def __init__(self, neg_slope=1e-2): 6 | self.neg_slope = neg_slope 7 | 8 | def __call__(self, module): 9 | if isinstance(module, nn.Conv3d) or isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.ConvTranspose3d): 10 | module.weight = nn.init.kaiming_normal_(module.weight, a=self.neg_slope) 11 | if module.bias is not None: 12 | module.bias = nn.init.constant_(module.bias, 0) 13 | -------------------------------------------------------------------------------- /nnunetv2/utilities/plans_handling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kent0n-Li/nnSAM/eb5ccf7f555641e00457eeb3f936afe0a80f56f6/nnunetv2/utilities/plans_handling/__init__.py -------------------------------------------------------------------------------- /nnunetv2/utilities/tensor_utilities.py: -------------------------------------------------------------------------------- 1 | from typing import Union, List, Tuple 2 | 3 | import numpy as np 4 | import torch 5 | 6 | 7 | def sum_tensor(inp: torch.Tensor, axes: Union[np.ndarray, Tuple, List], keepdim: bool = False) -> torch.Tensor: 8 | axes = np.unique(axes).astype(int) 9 | if keepdim: 10 | for ax in axes: 11 | inp = inp.sum(int(ax), keepdim=True) 12 | else: 13 | for ax in sorted(axes, reverse=True): 14 | inp = inp.sum(int(ax)) 15 | return inp 16 | -------------------------------------------------------------------------------- /nnunetv2/utilities/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 HIP Applied Computer Vision Lab, Division of Medical Image Computing, German Cancer Research Center 2 | # (DKFZ), Heidelberg, Germany 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | from typing import Union 16 | 17 | from batchgenerators.utilities.file_and_folder_operations import * 18 | import numpy as np 19 | import re 20 | 21 | 22 | def get_identifiers_from_splitted_dataset_folder(folder: str, file_ending: str): 23 | files = subfiles(folder, suffix=file_ending, join=False) 24 | # all files must be .nii.gz and have 4 digit channel index 25 | crop = len(file_ending) + 5 26 | files = [i[:-crop] for i in files] 27 | # only unique image ids 28 | files = np.unique(files) 29 | return files 30 | 31 | 32 | def create_lists_from_splitted_dataset_folder(folder: str, file_ending: str, identifiers: List[str] = None) -> List[List[str]]: 33 | """ 34 | does not rely on dataset.json 35 | """ 36 | if identifiers is None: 37 | identifiers = get_identifiers_from_splitted_dataset_folder(folder, file_ending) 38 | files = subfiles(folder, suffix=file_ending, join=False, sort=True) 39 | list_of_lists = [] 40 | for f in identifiers: 41 | p = re.compile(f + "_\d\d\d\d" + file_ending) 42 | list_of_lists.append([join(folder, i) for i in files if p.fullmatch(i)]) 43 | return list_of_lists -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "nnunetv2" 3 | version = "2.1.2" 4 | requires-python = ">=3.9" 5 | description = "nnU-Net is a framework for out-of-the box image segmentation." 6 | readme = "readme.md" 7 | license = { file = "LICENSE" } 8 | authors = [ 9 | { name = "Fabian Isensee", email = "f.isensee@dkfz-heidelberg.de"}, 10 | { name = "Helmholtz Imaging Applied Computer Vision Lab" } 11 | ] 12 | classifiers = [ 13 | "Development Status :: 5 - Production/Stable", 14 | "Intended Audience :: Developers", 15 | "Intended Audience :: Science/Research", 16 | "Intended Audience :: Healthcare Industry", 17 | "Programming Language :: Python :: 3", 18 | "License :: OSI Approved :: Apache Software License", 19 | "Topic :: Scientific/Engineering :: Artificial Intelligence", 20 | "Topic :: Scientific/Engineering :: Image Recognition", 21 | "Topic :: Scientific/Engineering :: Medical Science Apps.", 22 | ] 23 | keywords = [ 24 | 'deep learning', 25 | 'image segmentation', 26 | 'semantic segmentation', 27 | 'medical image analysis', 28 | 'medical image segmentation', 29 | 'nnU-Net', 30 | 'nnunet' 31 | ] 32 | dependencies = [ 33 | "torch>=2.0.0", 34 | "acvl-utils>=0.2", 35 | "dynamic-network-architectures>=0.2", 36 | "tqdm", 37 | "dicom2nifti", 38 | "scipy", 39 | "batchgenerators>=0.25", 40 | "numpy", 41 | "scikit-learn", 42 | "scikit-image>=0.19.3", 43 | "SimpleITK>=2.2.1", 44 | "pandas", 45 | "graphviz", 46 | 'tifffile', 47 | 'requests', 48 | "nibabel", 49 | "matplotlib", 50 | "seaborn", 51 | "imagecodecs", 52 | "yacs" 53 | ] 54 | 55 | [project.urls] 56 | homepage = "https://github.com/MIC-DKFZ/nnUNet" 57 | repository = "https://github.com/MIC-DKFZ/nnUNet" 58 | 59 | [project.scripts] 60 | nnUNetv2_plan_and_preprocess = "nnunetv2.experiment_planning.plan_and_preprocess_entrypoints:plan_and_preprocess_entry" 61 | nnUNetv2_extract_fingerprint = "nnunetv2.experiment_planning.plan_and_preprocess_entrypoints:extract_fingerprint_entry" 62 | nnUNetv2_plan_experiment = "nnunetv2.experiment_planning.plan_and_preprocess_entrypoints:plan_experiment_entry" 63 | nnUNetv2_preprocess = "nnunetv2.experiment_planning.plan_and_preprocess_entrypoints:preprocess_entry" 64 | nnUNetv2_train = "nnunetv2.run.run_training:run_training_entry" 65 | nnUNetv2_predict_from_modelfolder = "nnunetv2.inference.predict_from_raw_data:predict_entry_point_modelfolder" 66 | nnUNetv2_predict = "nnunetv2.inference.predict_from_raw_data:predict_entry_point" 67 | nnUNetv2_convert_old_nnUNet_dataset = "nnunetv2.dataset_conversion.convert_raw_dataset_from_old_nnunet_format:convert_entry_point" 68 | nnUNetv2_find_best_configuration = "nnunetv2.evaluation.find_best_configuration:find_best_configuration_entry_point" 69 | nnUNetv2_determine_postprocessing = "nnunetv2.postprocessing.remove_connected_components:entry_point_determine_postprocessing_folder" 70 | nnUNetv2_apply_postprocessing = "nnunetv2.postprocessing.remove_connected_components:entry_point_apply_postprocessing" 71 | nnUNetv2_ensemble = "nnunetv2.ensembling.ensemble:entry_point_ensemble_folders" 72 | nnUNetv2_accumulate_crossval_results = "nnunetv2.evaluation.find_best_configuration:accumulate_crossval_results_entry_point" 73 | nnUNetv2_plot_overlay_pngs = "nnunetv2.utilities.overlay_plots:entry_point_generate_overlay" 74 | nnUNetv2_download_pretrained_model_by_url = "nnunetv2.model_sharing.entry_points:download_by_url" 75 | nnUNetv2_install_pretrained_model_from_zip = "nnunetv2.model_sharing.entry_points:install_from_zip_entry_poin" 76 | nnUNetv2_export_model_to_zip = "nnunetv2.model_sharing.entry_points:export_pretrained_model_entr" 77 | nnUNetv2_move_plans_between_datasets = "nnunetv2.experiment_planning.plans_for_pretraining.move_plans_between_datasets:entry_point_move_plans_between_datasets" 78 | nnUNetv2_evaluate_folder = "nnunetv2.evaluation.evaluate_predictions:evaluate_folder_entry_point" 79 | nnUNetv2_evaluate_simple = "nnunetv2.evaluation.evaluate_predictions:evaluate_simple_entry_point" 80 | nnUNetv2_convert_MSD_dataset = "nnunetv2.dataset_conversion.convert_MSD_dataset:entry_point" 81 | 82 | [project.optional-dependencies] 83 | dev = [ 84 | "black", 85 | "ruff", 86 | "pre-commit" 87 | ] -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | if __name__ == "__main__": 4 | setuptools.setup() 5 | --------------------------------------------------------------------------------