├── nnunet ├── run │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ └── default_configuration.cpython-37.pyc │ ├── default_configuration.py │ └── run_training.py ├── evaluation │ ├── __init__.py │ ├── model_selection │ │ ├── __init__.py │ │ ├── collect_all_fold0_results_and_summarize_in_one_csv.py │ │ ├── figure_out_what_to_submit.py │ │ ├── summarize_results_in_one_json.py │ │ ├── ensemble.py │ │ └── summarize_results_with_plans.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── metrics.cpython-37.pyc │ │ └── evaluator.cpython-37.pyc │ ├── collect_results_files.py │ ├── add_mean_dice_to_json.py │ ├── surface_dice.py │ └── add_dummy_task_with_mean_over_all_tasks.py ├── inference │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ └── segmentation_export.cpython-37.pyc │ ├── ensemble_predictions.py │ ├── segmentation_export.py │ └── predict_simple.py ├── training │ ├── __init__.py │ ├── dataloading │ │ ├── __init__.py │ │ └── __pycache__ │ │ │ ├── __init__.cpython-37.pyc │ │ │ └── dataset_loading.cpython-37.pyc │ ├── cascade_stuff │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-37.pyc │ │ │ └── predict_next_stage.cpython-37.pyc │ │ └── predict_next_stage.py │ ├── data_augmentation │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-37.pyc │ │ │ ├── custom_transforms.cpython-37.pyc │ │ │ ├── pyramid_augmentations.cpython-37.pyc │ │ │ └── default_data_augmentation.cpython-37.pyc │ │ ├── custom_transforms.py │ │ └── pyramid_augmentations.py │ ├── loss_functions │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-37.pyc │ │ │ ├── TopK_loss.cpython-37.pyc │ │ │ ├── dice_loss.cpython-37.pyc │ │ │ ├── LovaszSoftmax.cpython-37.pyc │ │ │ └── ND_Crossentropy.cpython-37.pyc │ │ ├── TopK_loss.py │ │ ├── GDL.py │ │ ├── ND_Crossentropy.py │ │ └── LovaszSoftmax.py │ ├── network_training │ │ ├── __init__.py │ │ ├── nnUNet_variants │ │ │ ├── __init__.py │ │ │ ├── nnUNetTrainerCE.py │ │ │ ├── nnUNetTrainerNoMirroring.py │ │ │ └── nnUNetTrainerNoDA.py │ │ └── __pycache__ │ │ │ ├── __init__.cpython-37.pyc │ │ │ ├── nnUNetTrainer.cpython-37.pyc │ │ │ ├── network_trainer.cpython-37.pyc │ │ │ └── nnUNetTrainerCascadeFullRes.cpython-37.pyc │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ └── model_restore.cpython-37.pyc │ └── model_restore.py ├── utilities │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── to_torch.cpython-37.pyc │ │ ├── nd_softmax.cpython-37.pyc │ │ ├── one_hot_encoding.cpython-37.pyc │ │ └── tensor_utilities.cpython-37.pyc │ ├── one_hot_encoding.py │ ├── to_torch.py │ ├── nd_softmax.py │ ├── tensor_utilities.py │ └── online_evaluation_metrics.py ├── preprocessing │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── cropping.cpython-37.pyc │ │ └── preprocessing.cpython-37.pyc │ ├── lumbosacral_joint_sampling.py │ └── cropping.py ├── experiment_planning │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── common_utils.cpython-37.pyc │ │ ├── configuration.cpython-37.pyc │ │ ├── DatasetAnalyzer.cpython-37.pyc │ │ ├── summarize_plans.cpython-37.pyc │ │ ├── find_classes_in_slice.cpython-37.pyc │ │ ├── plan_and_preprocess_task.cpython-37.pyc │ │ ├── experiment_planner_baseline_2DUNet.cpython-37.pyc │ │ └── experiment_planner_baseline_3DUNet.cpython-37.pyc │ ├── configuration.py │ ├── find_classes_in_slice.py │ ├── summarize_plans.py │ └── common_utils.py ├── network_architecture │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── generic_UNet.cpython-37.pyc │ │ ├── initialization.cpython-37.pyc │ │ └── neural_network.cpython-37.pyc │ └── initialization.py ├── dataset_conversion │ ├── __init__.py │ ├── readme.md │ └── JstPelvisSegmentation_5label.py ├── __pycache__ │ ├── paths.cpython-37.pyc │ └── __init__.cpython-37.pyc ├── __init__.py ├── paths.py └── runs.py ├── plans.pkl ├── images ├── post.png ├── dataset.png ├── visual.png └── ctplevic.png ├── splits_final.pkl ├── __pycache__ ├── utils.cpython-37.pyc └── postprocessing.cpython-37.pyc ├── requirements.txt ├── utils.py ├── save_evaluation_results2csv.py ├── save_evaluation_results2csv_Manu.py ├── postprocessing.py └── evaluation.py /nnunet/run/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /plans.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/plans.pkl -------------------------------------------------------------------------------- /nnunet/evaluation/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/inference/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/training/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/utilities/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /images/post.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/images/post.png -------------------------------------------------------------------------------- /nnunet/preprocessing/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /splits_final.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/splits_final.pkl -------------------------------------------------------------------------------- /images/dataset.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/images/dataset.png -------------------------------------------------------------------------------- /images/visual.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/images/visual.png -------------------------------------------------------------------------------- /nnunet/experiment_planning/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/network_architecture/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/training/dataloading/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /images/ctplevic.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/images/ctplevic.png -------------------------------------------------------------------------------- /nnunet/dataset_conversion/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | from . import * -------------------------------------------------------------------------------- /nnunet/evaluation/model_selection/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/training/cascade_stuff/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/training/data_augmentation/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/training/loss_functions/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/training/network_training/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /nnunet/training/network_training/nnUNet_variants/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import * -------------------------------------------------------------------------------- /__pycache__/utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/__pycache__/utils.cpython-37.pyc -------------------------------------------------------------------------------- /__pycache__/postprocessing.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/__pycache__/postprocessing.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/__pycache__/paths.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/__pycache__/paths.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/run/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/run/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/evaluation/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/evaluation/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/evaluation/__pycache__/metrics.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/evaluation/__pycache__/metrics.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/inference/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/inference/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/utilities/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/utilities/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/utilities/__pycache__/to_torch.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/utilities/__pycache__/to_torch.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/evaluation/__pycache__/evaluator.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/evaluation/__pycache__/evaluator.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/preprocessing/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/preprocessing/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/preprocessing/__pycache__/cropping.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/preprocessing/__pycache__/cropping.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/__pycache__/model_restore.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/__pycache__/model_restore.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/utilities/__pycache__/nd_softmax.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/utilities/__pycache__/nd_softmax.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/preprocessing/__pycache__/preprocessing.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/preprocessing/__pycache__/preprocessing.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/run/__pycache__/default_configuration.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/run/__pycache__/default_configuration.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/utilities/__pycache__/one_hot_encoding.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/utilities/__pycache__/one_hot_encoding.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/utilities/__pycache__/tensor_utilities.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/utilities/__pycache__/tensor_utilities.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/experiment_planning/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/experiment_planning/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/inference/__pycache__/segmentation_export.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/inference/__pycache__/segmentation_export.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/network_architecture/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/network_architecture/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/dataloading/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/dataloading/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch>1.0.0 2 | tqdm 3 | dicom2nifti 4 | scikit-image>=0.14 5 | medpy 6 | scipy 7 | batchgenerators>=0.19.4 8 | numpy 9 | sklearn 10 | SimpleITK 11 | pandas 12 | -------------------------------------------------------------------------------- /nnunet/experiment_planning/__pycache__/common_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/experiment_planning/__pycache__/common_utils.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/cascade_stuff/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/cascade_stuff/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/loss_functions/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/loss_functions/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/experiment_planning/__pycache__/configuration.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/experiment_planning/__pycache__/configuration.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/network_architecture/__pycache__/generic_UNet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/network_architecture/__pycache__/generic_UNet.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/loss_functions/__pycache__/TopK_loss.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/loss_functions/__pycache__/TopK_loss.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/loss_functions/__pycache__/dice_loss.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/loss_functions/__pycache__/dice_loss.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/network_training/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/network_training/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/experiment_planning/__pycache__/DatasetAnalyzer.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/experiment_planning/__pycache__/DatasetAnalyzer.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/experiment_planning/__pycache__/summarize_plans.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/experiment_planning/__pycache__/summarize_plans.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/network_architecture/__pycache__/initialization.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/network_architecture/__pycache__/initialization.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/network_architecture/__pycache__/neural_network.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/network_architecture/__pycache__/neural_network.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/data_augmentation/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/data_augmentation/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/dataloading/__pycache__/dataset_loading.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/dataloading/__pycache__/dataset_loading.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/loss_functions/__pycache__/LovaszSoftmax.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/loss_functions/__pycache__/LovaszSoftmax.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/loss_functions/__pycache__/ND_Crossentropy.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/loss_functions/__pycache__/ND_Crossentropy.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/network_training/__pycache__/nnUNetTrainer.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/network_training/__pycache__/nnUNetTrainer.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/experiment_planning/__pycache__/find_classes_in_slice.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/experiment_planning/__pycache__/find_classes_in_slice.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/cascade_stuff/__pycache__/predict_next_stage.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/cascade_stuff/__pycache__/predict_next_stage.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/network_training/__pycache__/network_trainer.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/network_training/__pycache__/network_trainer.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/experiment_planning/__pycache__/plan_and_preprocess_task.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/experiment_planning/__pycache__/plan_and_preprocess_task.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/data_augmentation/__pycache__/custom_transforms.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/data_augmentation/__pycache__/custom_transforms.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/data_augmentation/__pycache__/pyramid_augmentations.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/data_augmentation/__pycache__/pyramid_augmentations.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/data_augmentation/__pycache__/default_data_augmentation.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/data_augmentation/__pycache__/default_data_augmentation.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/experiment_planning/__pycache__/experiment_planner_baseline_2DUNet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/experiment_planning/__pycache__/experiment_planner_baseline_2DUNet.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/experiment_planning/__pycache__/experiment_planner_baseline_3DUNet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/experiment_planning/__pycache__/experiment_planner_baseline_3DUNet.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/training/network_training/__pycache__/nnUNetTrainerCascadeFullRes.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MIRACLE-Center/CTPelvic1K/HEAD/nnunet/training/network_training/__pycache__/nnUNetTrainerCascadeFullRes.cpython-37.pyc -------------------------------------------------------------------------------- /nnunet/utilities/one_hot_encoding.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | 5 | 6 | def to_one_hot(seg, all_seg_labels=None): 7 | if all_seg_labels is None: 8 | all_seg_labels = np.unique(seg) 9 | result = np.zeros((len(all_seg_labels), *seg.shape), dtype=seg.dtype) 10 | for i, l in enumerate(all_seg_labels): 11 | result[i][seg == l] = 1 12 | return result 13 | -------------------------------------------------------------------------------- /nnunet/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | print("Please cite the following paper when using CTPelvic1K dataset:\n\nPengbo, Liu, et al. \"Deep Learning to Segment " 3 | "Pelvic Bones: Large-scale CT Datasets and Baseline Models.\" arXiv preprint arXiv: (2021).") 4 | print("\nIf you have questions or suggestions, feel free to open an issue at https://github.com/ICT-MIRACLE-lab/CTPelvic1K") 5 | 6 | from . import * -------------------------------------------------------------------------------- /nnunet/utilities/to_torch.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import torch 4 | 5 | 6 | def maybe_to_torch(d): 7 | if isinstance(d, list): 8 | d = [maybe_to_torch(i) if not isinstance(i, torch.Tensor) else i for i in d] 9 | elif not isinstance(d, torch.Tensor): 10 | d = torch.from_numpy(d).float() 11 | return d 12 | 13 | 14 | def to_cuda(data, non_blocking=True, gpu_id=0): 15 | if isinstance(data, list): 16 | data = [i.cuda(gpu_id, non_blocking=non_blocking) for i in data] 17 | else: 18 | data = data.cuda(gpu_id, non_blocking=True) 19 | return data 20 | -------------------------------------------------------------------------------- /nnunet/utilities/nd_softmax.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import torch 4 | 5 | 6 | def softmax_helper(x): 7 | """ 8 | same as F.softmax 9 | """ 10 | rpt = [1 for _ in range(len(x.size()))] 11 | rpt[1] = x.size(1) 12 | x_max = x.max(1, keepdim=True)[0].repeat(*rpt) 13 | e_x = torch.exp(x - x_max) 14 | return e_x / e_x.sum(1, keepdim=True).repeat(*rpt) 15 | 16 | if __name__ == '__main__': 17 | x=torch.rand(1,3,2,2) 18 | print(x) 19 | y=softmax_helper(x) 20 | print('-'*33) 21 | print(y) 22 | print('-'*33) 23 | import torch.nn as nn 24 | import torch.nn.functional as F 25 | 26 | print(F.softmax(x,dim=1)) -------------------------------------------------------------------------------- /nnunet/training/network_training/nnUNet_variants/nnUNetTrainerCE.py: -------------------------------------------------------------------------------- 1 | from nnunet.training.loss_functions.ND_Crossentropy import CrossentropyND 2 | from nnunet.training.network_training.nnUNetTrainer import nnUNetTrainer 3 | 4 | 5 | class nnUNetTrainerCE(nnUNetTrainer): 6 | def __init__(self, plans_file, fold, output_folder=None, dataset_directory=None, batch_dice=True, stage=None, 7 | unpack_data=True, deterministic=True, fp16=False): 8 | super(nnUNetTrainerCE, self).__init__(plans_file, fold, output_folder, dataset_directory, batch_dice, stage, 9 | unpack_data, deterministic, fp16) 10 | self.loss = CrossentropyND() 11 | -------------------------------------------------------------------------------- /nnunet/training/loss_functions/TopK_loss.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from nnunet.training.loss_functions.ND_Crossentropy import CrossentropyND 4 | 5 | 6 | class TopKLoss(CrossentropyND): 7 | """ 8 | Network has to have NO LINEARITY! 9 | """ 10 | def __init__(self, weight=None, ignore_index=-100, k=10): 11 | self.k = k 12 | super(TopKLoss, self).__init__(weight, False, ignore_index, reduce=False) 13 | 14 | def forward(self, inp, target): 15 | target = target[:, 0].long() 16 | res = super(TopKLoss, self).forward(inp, target) 17 | num_voxels = np.prod(res.shape, dtype=np.int64) 18 | res, _ = torch.topk(res.view((-1, )), int(num_voxels * self.k / 100), sorted=False) 19 | return res.mean() 20 | -------------------------------------------------------------------------------- /nnunet/experiment_planning/configuration.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | # experiment planning 4 | 5 | # a class in a patient will be set to background if it has less than X times the volume of the minimum volume of 6 | # that class in the training data 7 | MIN_SIZE_PER_CLASS_FACTOR = 0.5 8 | TARGET_SPACING_PERCENTILE = 50 9 | 10 | FEATUREMAP_MIN_EDGE_LENGTH_BOTTLENECK = 4 11 | FEATUREMAP_MIN_EDGE_LENGTH_BOTTLENECK2 = 6 12 | RESAMPLING_SEPARATE_Z_ANISOTROPY_THRESHOLD = 3 # z is defined as the axis with the highest spacing, also used to 13 | # determine whether we use 2d or 3d data augmentation 14 | 15 | HOW_MUCH_OF_A_PATIENT_MUST_THE_NETWORK_SEE_AT_STAGE0 = 4 # 1/4 of a patient 16 | 17 | batch_size_covers_max_percent_of_dataset = 0.05 # all samples in the batch together cannot cover more than 5% of the entire dataset 18 | dataset_min_batch_size_cap = 2 # if the dataset size dictates a very small batch size, do not make that smaller than 3 (if architecture dictates smaller batch size then use the smaller one of these two) 19 | 20 | 21 | 22 | -------------------------------------------------------------------------------- /nnunet/utilities/tensor_utilities.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | import torch 5 | 6 | 7 | def sum_tensor(inp, axes, keepdim=False): 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 | 17 | 18 | def mean_tensor(inp, axes, keepdim=False): 19 | axes = np.unique(axes).astype(int) 20 | if keepdim: 21 | for ax in axes: 22 | inp = inp.mean(int(ax), keepdim=True) 23 | else: 24 | for ax in sorted(axes, reverse=True): 25 | inp = inp.mean(int(ax)) 26 | return inp 27 | 28 | 29 | def flip(x, dim): 30 | """ 31 | flips the tensor at dimension dim (mirroring!) 32 | :param x: 33 | :param dim: 34 | :return: 35 | """ 36 | indices = [slice(None)] * x.dim() 37 | indices[dim] = torch.arange(x.size(dim) - 1, -1, -1, 38 | dtype=torch.long, device=x.device) 39 | return x[tuple(indices)] 40 | -------------------------------------------------------------------------------- /nnunet/network_architecture/initialization.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | from torch import nn 4 | 5 | 6 | class InitWeights_He(object): 7 | def __init__(self, neg_slope=1e-2): 8 | self.neg_slope = neg_slope ### todo: not be used in __call__ 9 | 10 | def __call__(self, module): 11 | if isinstance(module, nn.Conv3d) or isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.ConvTranspose3d): 12 | module.weight = nn.init.kaiming_normal_(module.weight, a=self.neg_slope) ### I HAVE changed this a corrosponding to up todo 13 | if module.bias is not None: 14 | module.bias = nn.init.constant_(module.bias, 0) 15 | 16 | 17 | class InitWeights_XavierUniform(object): 18 | def __init__(self, gain=1): 19 | self.gain = gain 20 | 21 | def __call__(self, module): 22 | if isinstance(module, nn.Conv3d) or isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.ConvTranspose3d): 23 | module.weight = nn.init.xavier_uniform_(module.weight, self.gain) 24 | if module.bias is not None: 25 | module.bias = nn.init.constant_(module.bias, 0) 26 | -------------------------------------------------------------------------------- /nnunet/evaluation/collect_results_files.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | from batchgenerators.utilities.file_and_folder_operations import subdirs, subfiles 4 | 5 | 6 | def crawl_and_copy(current_folder, out_folder, prefix="fabian_", suffix="ummary.json"): 7 | """ 8 | This script will run recursively through all subfolders of current_folder and copy all files that end with 9 | suffix with some automatically generated prefix into out_folder 10 | :param current_folder: 11 | :param out_folder: 12 | :param prefix: 13 | :return: 14 | """ 15 | s = subdirs(current_folder, join=False) 16 | f = subfiles(current_folder, join=False) 17 | f = [i for i in f if i.endswith(suffix)] 18 | if current_folder.find("fold0") != -1: 19 | for fl in f: 20 | shutil.copy(os.path.join(current_folder, fl), os.path.join(out_folder, prefix+fl)) 21 | for su in s: 22 | if prefix == "": 23 | add = su 24 | else: 25 | add = "__" + su 26 | crawl_and_copy(os.path.join(current_folder, su), out_folder, prefix=prefix+add) 27 | 28 | 29 | if __name__ == "__main__": 30 | from nnunet.paths import network_training_output_dir 31 | output_folder = "/home/fabian/PhD/results/nnUNetV2/leaderboard" 32 | crawl_and_copy(network_training_output_dir, output_folder) 33 | from nnunet.evaluation.add_mean_dice_to_json import run_in_folder 34 | run_in_folder(output_folder) 35 | -------------------------------------------------------------------------------- /nnunet/utilities/online_evaluation_metrics.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | import torch 5 | 6 | 7 | def hard_dice(output, target): 8 | if isinstance(output, torch.Tensor): 9 | output = output.detach().cpu().numpy() 10 | if isinstance(target, torch.Tensor): 11 | target = target.detach().cpu().numpy() 12 | target = target[:, 0] 13 | # target is not one hot encoded, output is 14 | # target must be the CPU segemtnation, not tensor. output is pytorch tensor 15 | num_classes = output.shape[1] 16 | output = output.argmax(1) 17 | foreground_classes = np.arange(1, num_classes) 18 | all_tp = [] 19 | all_fp = [] 20 | all_fn = [] 21 | all_fg_dc = [] 22 | for s in range(target.shape[0]): 23 | tp = [] 24 | fp = [] 25 | fn = [] 26 | for c in foreground_classes: 27 | t_is_c = target[s] == c 28 | o_is_c = output[s] == c 29 | t_is_not_c = target[s] != c 30 | o_is_not_c = output[s] != c 31 | tp.append(np.sum(o_is_c & t_is_c)) 32 | fp.append(np.sum(o_is_c & t_is_not_c)) 33 | fn.append(np.sum(o_is_not_c & t_is_c)) 34 | foreground_dice = [2 * i / (2 * i + j + k + 1e-8) for i, j, k in zip(tp, fp, fn)] 35 | all_tp.append(tp) 36 | all_fp.append(fp) 37 | all_fn.append(fn) 38 | all_fg_dc.append(foreground_dice) 39 | return all_fg_dc, all_tp, all_fp, all_fn 40 | -------------------------------------------------------------------------------- /nnunet/evaluation/add_mean_dice_to_json.py: -------------------------------------------------------------------------------- 1 | import json 2 | import numpy as np 3 | from batchgenerators.utilities.file_and_folder_operations import subfiles 4 | from collections import OrderedDict 5 | 6 | 7 | def foreground_mean(filename): 8 | with open(filename, 'r') as f: 9 | res = json.load(f) 10 | class_ids = np.array([int(i) for i in res['results']['mean'].keys() if (i != 'mean')]) 11 | class_ids = class_ids[class_ids != 0] 12 | class_ids = class_ids[class_ids != -1] 13 | class_ids = class_ids[class_ids != 99] 14 | 15 | tmp = res['results']['mean'].get('99') 16 | if tmp is not None: 17 | _ = res['results']['mean'].pop('99') 18 | 19 | metrics = res['results']['mean']['1'].keys() 20 | res['results']['mean']["mean"] = OrderedDict() 21 | for m in metrics: 22 | foreground_values = [res['results']['mean'][str(i)][m] for i in class_ids] 23 | res['results']['mean']["mean"][m] = np.nanmean(foreground_values) 24 | with open(filename, 'w') as f: 25 | json.dump(res, f, indent=4, sort_keys=True) 26 | 27 | def run_in_folder(folder): 28 | json_files = subfiles(folder, True, None, ".json", True) 29 | json_files = [i for i in json_files if not i.split("/")[-1].startswith(".") and not i.endswith("_globalMean.json")] # stupid mac 30 | for j in json_files: 31 | foreground_mean(j) 32 | 33 | if __name__ == "__main__": 34 | folder = "/media/fabian/Results/nnUNetOutput_final/summary_jsons" 35 | run_in_folder(folder) -------------------------------------------------------------------------------- /nnunet/experiment_planning/find_classes_in_slice.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | import pickle 5 | from collections import OrderedDict 6 | 7 | 8 | def add_classes_in_slice_info(args): 9 | """ 10 | We need this for 2D dataloader with oversampling. As of now it will detect slices that contain specific classes 11 | at run time, meaning it needs to iterate over an entire patient just to extract one slice. That is obviously bad, 12 | so we are doing this once beforehand and just give the dataloader the info it needs in the patients pkl file. 13 | 14 | """ 15 | npz_file, pkl_file, all_classes = args 16 | seg_map = np.load(npz_file)['data'][-1] 17 | with open(pkl_file, 'rb') as f: 18 | props = pickle.load(f) 19 | #if props.get('classes_in_slice_per_axis') is not None: 20 | print(pkl_file) 21 | # this will be a dict of dict where the first dict encodes the axis along which a slice is extracted in its keys. 22 | # The second dict (value of first dict) will have all classes as key and as values a list of all slice ids that 23 | # contain this class 24 | classes_in_slice = OrderedDict() 25 | for axis in range(3): 26 | other_axes = tuple([i for i in range(3) if i != axis]) 27 | classes_in_slice[axis] = OrderedDict() 28 | for c in all_classes: 29 | valid_slices = np.where(np.sum(seg_map == c, axis=other_axes) > 0)[0] 30 | classes_in_slice[axis][c] = valid_slices 31 | 32 | number_of_voxels_per_class = OrderedDict() 33 | for c in all_classes: 34 | number_of_voxels_per_class[c] = np.sum(seg_map == c) 35 | 36 | props['classes_in_slice_per_axis'] = classes_in_slice 37 | props['number_of_voxels_per_class'] = number_of_voxels_per_class 38 | 39 | with open(pkl_file, 'wb') as f: 40 | pickle.dump(props, f) 41 | -------------------------------------------------------------------------------- /nnunet/training/network_training/nnUNet_variants/nnUNetTrainerNoMirroring.py: -------------------------------------------------------------------------------- 1 | from nnunet.training.network_training.nnUNetTrainer import nnUNetTrainer 2 | 3 | 4 | class nnUNetTrainerNoMirroring(nnUNetTrainer): 5 | def validate(self, do_mirroring=True, use_train_mode=False, tiled=True, step=2, save_softmax=True, 6 | use_gaussian=True, compute_global_dice=True, overwrite=True, validation_folder_name='validation_raw'): 7 | if do_mirroring: 8 | print("WARNING! do_mirroring was True but we cannot do that because we trained without mirroring. " 9 | "do_mirroring was set to False") 10 | do_mirroring = False 11 | return super().validate(do_mirroring, use_train_mode, tiled, step, save_softmax, use_gaussian, 12 | compute_global_dice, overwrite, validation_folder_name) 13 | 14 | def setup_DA_params(self): 15 | super().setup_DA_params() 16 | self.data_aug_params["do_mirror"] = False 17 | # you can also use self.data_aug_params["mirror_axes"] to set axes for mirroring. 18 | # Default is self.data_aug_params["mirror_axes"] = (0, 1, 2) 19 | # 0, 1, 2 are the first, second and thirs spatial axes. 20 | 21 | def predict_preprocessed_data_return_softmax(self, data, do_mirroring, num_repeats, use_train_mode, batch_size, 22 | mirror_axes, tiled, tile_in_z, step, min_size, use_gaussian, 23 | all_in_gpu=False): 24 | 25 | if do_mirroring: 26 | print("WARNING! do_mirroring was True but we cannot do that because we trained without mirroring. " 27 | "do_mirroring was set to False") 28 | return super().predict_preprocessed_data_return_softmax(data, do_mirroring, num_repeats, use_train_mode, 29 | batch_size, mirror_axes, tiled, tile_in_z, step, 30 | min_size, use_gaussian) 31 | -------------------------------------------------------------------------------- /nnunet/evaluation/surface_dice.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | import numpy as np 5 | from medpy.metric.binary import __surface_distances 6 | 7 | 8 | def normalized_surface_dice(a: np.ndarray, b: np.ndarray, threshold: float, spacing: tuple = None, connectivity=1): 9 | """ 10 | This implementation differs from the official surface dice implementation! These two are not comparable!!!!! 11 | 12 | The normalized surface dice is symmetric, so it should not matter wheter a or b is the reference image 13 | 14 | This implementation natively supports 2D and 3D images. Whether other dimensions are supported depends in the 15 | __surface_distances implementation in medpy 16 | 17 | :param a: image 1, must have the same shape as b 18 | :param b: image 2, must have the same shape as a 19 | :param threshold: distances below this threshold will be counted as true positives. Threshold is in mm, not voxels! 20 | (if spacing = (1, 1(, 1)) then one voxel=1mm so the threshold is effectively in voxels) 21 | must be a tuple of len dimension(a) 22 | :param spacing: how many mm is one voxel in reality? Can be left at None, we then assume an isotropic spacing of 1mm 23 | :param connectivity: see scipy.ndimage.generate_binary_structure for more information. I suggest you leave that 24 | one alone 25 | :return: 26 | """ 27 | assert all([i == j for i, j in zip(a.shape, b.shape)]), "a and b must have the same shape. a.shape= %s, " \ 28 | "b.shape= %s" % (str(a.shape), str(b.shape)) 29 | if spacing is None: 30 | spacing = tuple([1 for _ in range(len(a.shape))]) 31 | a_to_b = __surface_distances(a, b, spacing, connectivity) 32 | b_to_a = __surface_distances(b, a, spacing, connectivity) 33 | 34 | tp_a = np.sum(a_to_b <= threshold) 35 | tp_b = np.sum(b_to_a <= threshold) 36 | 37 | fp = np.sum(a_to_b > threshold) 38 | fn = np.sum(b_to_a > threshold) 39 | 40 | dc = (tp_a + tp_b) / (tp_a + tp_b + fp + fn + 1e-8) # 1e-8 just so that we don't get div by 0 41 | return dc 42 | 43 | -------------------------------------------------------------------------------- /nnunet/dataset_conversion/readme.md: -------------------------------------------------------------------------------- 1 | # Dataset conversion instructions 2 | 3 | ## How to convert non-Decathlon datasets for nnU-Net 4 | In this file we provide a description on how you need to convert your dataset to make it compatible with nnU-Net. 5 | 6 | For the purpose of this 7 | manual, we refer to your dataset as `TaskXX_MY_DATASET`. Hereby, `XX` is is a dual-digit number and `MY_DATASET` can 8 | be anything you want. Like `Task22_ipcai2021` in this experiment. 9 | 10 | We follow the folder structure of the Medical Segmentation Decathlon (MSD). In the `splitted_4d_output_dir` (see `paths.py`) 11 | , create a subfolder 12 | called `TaskXX_MY_DATASET`. In that subfolder, create the following three directories: `imagesTr`, `labelsTr`(, `imagesTs`). 13 | These are for training images, training labels (and test images), respectively. 14 | 15 | Just like the MSD we use the nifti (.nii.gz) file format. There is a crucial difference though. 16 | While the MSD provides 4D niftis where the first axis is for the modality, we prefer to split each training case into 17 | separate 3D Niftis. So what was a File `patientID.nii.gz` containing an image of shape (4, 160, 190, 160) now is four files with shape 18 | (160, 190, 160) each. These four files should be named `patientID_0000.nii.gz`, `patientID_0001.nii.gz`, `patientID_0002.nii.gz`, 19 | `patientID_0003.nii.gz`, where the ending 4-digit represents the modality. `patientID` can be anything you want. Make 20 | sure that the same modality is assigned the same digit for all patients (for example if you have T1 and T2 then you 21 | need to make sure T1 is always 0000 and T2 always 0001). If you data has only one modality, like this experiment, just name your cases `patientID_0000.nii.gz`. 22 | 23 | Copy all training cases (just the data, not the labels) into the `imagesTr` subfolder. 24 | 25 | Copy training labels into the `labelsTr` subfolder. Labels are always 3D niftis and should be named `patientID.nii.gz` 26 | where `patientiID` is identical to the identifier used for the corresponding raw data. 27 | 28 | 29 | Finally you need to create a `dataset.json` file that you place in the `TaskXX_MY_DATASET` root folder. This file was 30 | always provided in the MSD challenge and is therefore a requirement for nnU-Net. Have a look at (for example) 31 | `LiverTumorSegmentationChallenge.py` to see what it needs to look like. Important: The list stored in 'training' contains images and labels are 32 | both called `patientID.nii.gz` (dropping the 0000 ending we used here) for consistency! 33 | 34 | -------------------------------------------------------------------------------- /nnunet/training/loss_functions/GDL.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from nnunet.training.loss_functions.dice_loss import get_tp_fp_fn 3 | from nnunet.utilities.tensor_utilities import sum_tensor 4 | from torch import nn 5 | 6 | 7 | class GDL(nn.Module): 8 | def __init__(self, apply_nonlin=None, batch_dice=False, do_bg=True, smooth=1., 9 | square=False, square_volumes=False): 10 | """ 11 | square_volumes will square the weight term. The paper recommends square_volumes=True; I don't (just an intuition) 12 | """ 13 | super(GDL, self).__init__() 14 | 15 | self.square_volumes = square_volumes 16 | self.square = square 17 | self.do_bg = do_bg 18 | self.batch_dice = batch_dice 19 | self.apply_nonlin = apply_nonlin 20 | self.smooth = smooth 21 | 22 | def forward(self, x, y, loss_mask=None): 23 | shp_x = x.shape 24 | shp_y = y.shape 25 | 26 | if self.batch_dice: 27 | axes = [0] + list(range(2, len(shp_x))) 28 | else: 29 | axes = list(range(2, len(shp_x))) 30 | 31 | if len(shp_x) != len(shp_y): 32 | y = y.view((shp_y[0], 1, *shp_y[1:])) 33 | 34 | if all([i == j for i, j in zip(x.shape, y.shape)]): 35 | # if this is the case then gt is probably already a one hot encoding 36 | y_onehot = y 37 | else: 38 | gt = y.long() 39 | y_onehot = torch.zeros(shp_x) 40 | if x.device.type == "cuda": 41 | y_onehot = y_onehot.cuda(x.device.index) 42 | y_onehot.scatter_(1, gt, 1) 43 | 44 | if self.apply_nonlin is not None: 45 | x = self.apply_nonlin(x) 46 | 47 | if not self.do_bg: 48 | x = x[:, 1:] 49 | y_onehot = y_onehot[:, 1:] 50 | 51 | tp, fp, fn = get_tp_fp_fn(x, y_onehot, axes, loss_mask, self.square) 52 | 53 | # GDL weight computation, we use 1/V 54 | volumes = sum_tensor(y_onehot, axes) 55 | 56 | if self.square_volumes: 57 | volumes = volumes ** 2 58 | 59 | # apply weights 60 | tp = tp / volumes 61 | fp = fp / volumes 62 | fn = fn / volumes 63 | 64 | # sum over classes 65 | if self.batch_dice: 66 | axis = 0 67 | else: 68 | axis = 1 69 | 70 | tp = tp.sum(axis, keepdim=False) 71 | fp = fp.sum(axis, keepdim=False) 72 | fn = fn.sum(axis, keepdim=False) 73 | 74 | # compute dice 75 | dc = (2 * tp + self.smooth) / (2 * tp + fp + fn + self.smooth) 76 | 77 | dc = dc.mean() 78 | 79 | return -dc 80 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import SimpleITK as sitk 3 | 4 | import pickle as pkl 5 | import matplotlib.pyplot as plt 6 | import time 7 | import os 8 | import cv2 9 | 10 | def sdf_func(segImg, name): 11 | """ 12 | segImg is a sitk Image 13 | """ 14 | Sdf = sitk.SignedMaurerDistanceMap(segImg, squaredDistance=False) 15 | # Sdf = sitk.Sigmoid(-Sdf, 50, 0, 1, 0) # alpha, beta, max, min 16 | Sdf = sitk.Sigmoid(-Sdf, 2, 0, 1, 0) # alpha, beta, max, min ### 2 for boundary weight 17 | seg = sitk.GetArrayFromImage(Sdf) 18 | # seg[seg > 0.4999] = 1 # convert sdf back to numpy array, and clip 0.5 above to 1 (inside) 19 | # heat_map = seg + 0.5 # putGaussianMapF(mask, sigma=50.0) 20 | seg = (seg-seg.min())/(seg.max()-seg.min()) 21 | print(name, seg.shape, seg.max(), seg.min()) 22 | return seg#heat_map 23 | 24 | def gatherfiles(path, prefix=None, midfix=None, postfix=None, extname=True): 25 | files = os.listdir(path) 26 | if not prefix is None: 27 | files = [i for i in files if i.startswith(prefix)] 28 | if not midfix is None: 29 | files = [i for i in files if midfix in i] 30 | if not postfix is None: 31 | files = [i for i in files if i.endswith(postfix)] 32 | if extname: 33 | return files 34 | else: 35 | files = [os.path.splitext(i)[0] for i in files] 36 | return files 37 | 38 | def load_pkl(file): 39 | with open(file, 'rb') as f: 40 | a = pkl.load(f) 41 | return a 42 | 43 | 44 | def save_pkl(info, file): 45 | with open(file, 'wb') as f: 46 | pkl.dump(info, f) 47 | 48 | 49 | def _change_label(label, idx_before, idx_after): 50 | label[label == idx_before] = idx_after 51 | return label 52 | 53 | 54 | def _Series_dicom_reader(path): 55 | Reader = sitk.ImageSeriesReader() 56 | name = Reader.GetGDCMSeriesFileNames(path) 57 | Reader.SetFileNames(name) 58 | Image = Reader.Execute() 59 | 60 | image = sitk.GetArrayFromImage(Image) 61 | Spa = Image.GetSpacing() 62 | Ori = Image.GetOrigin() 63 | Dir = Image.GetDirection() 64 | return Image, image, (Spa, Ori, Dir) 65 | 66 | 67 | def _sitk_Image_reader(path): 68 | Image = sitk.ReadImage(path) 69 | image = sitk.GetArrayFromImage(Image) 70 | Spa = Image.GetSpacing() 71 | Ori = Image.GetOrigin() 72 | Dir = Image.GetDirection() 73 | return Image, image, (Spa, Ori, Dir) 74 | 75 | 76 | def _sitk_image_writer(image, meta, path): 77 | Image = sitk.GetImageFromArray(image) 78 | if meta is None: 79 | pass 80 | else: 81 | Image.SetSpacing(meta[0]) 82 | Image.SetOrigin(meta[1]) 83 | Image.SetDirection(meta[2]) 84 | sitk.WriteImage(Image, path) -------------------------------------------------------------------------------- /nnunet/evaluation/model_selection/collect_all_fold0_results_and_summarize_in_one_csv.py: -------------------------------------------------------------------------------- 1 | 2 | from nnunet.evaluation.model_selection.summarize_results_in_one_json import summarize 3 | from nnunet.paths import network_training_output_dir 4 | from batchgenerators.utilities.file_and_folder_operations import * 5 | 6 | if __name__ == "__main__": 7 | summary_output_folder = join(network_training_output_dir, "summary_jsons_fold0") 8 | maybe_mkdir_p(summary_output_folder) 9 | summarize(range(50), output_dir=summary_output_folder, folds=(0,)) 10 | 11 | results_csv = join(network_training_output_dir, "summary_fold0.csv") 12 | 13 | summary_files = subfiles(summary_output_folder, suffix='.json', join=False) 14 | 15 | with open(results_csv, 'w') as f: 16 | for s in summary_files: 17 | if s.find("ensemble") == -1: 18 | task, network, trainer, plans, validation_folder = s.split("__") 19 | else: 20 | n1, n2 = s.split("--") 21 | n1 = n1[n1.find("ensemble_") + len("ensemble_") :] 22 | task = s.split("__")[0] 23 | network = "ensemble" 24 | trainer = n1 25 | plans = n2 26 | validation_folder = "none" 27 | validation_folder = validation_folder[:-len('.json')] 28 | results = load_json(join(summary_output_folder, s))['results']['mean']['mean']['Dice'] 29 | f.write("%s,%s,%s,%s,%s,%02.4f\n" % (task, 30 | network, trainer, validation_folder, plans, results)) 31 | 32 | summary_output_folder = join(network_training_output_dir, "summary_jsons") 33 | maybe_mkdir_p(summary_output_folder) 34 | summarize(['all'], output_dir=summary_output_folder) 35 | 36 | results_csv = join(network_training_output_dir, "summary_allFolds.csv") 37 | 38 | summary_files = subfiles(summary_output_folder, suffix='.json', join=False) 39 | 40 | with open(results_csv, 'w') as f: 41 | for s in summary_files: 42 | if s.find("ensemble") == -1: 43 | task, network, trainer, plans, validation_folder = s.split("__") 44 | else: 45 | n1, n2 = s.split("--") 46 | n1 = n1[n1.find("ensemble_") + len("ensemble_") :] 47 | task = s.split("__")[0] 48 | network = "ensemble" 49 | trainer = n1 50 | plans = n2 51 | validation_folder = "none" 52 | validation_folder = validation_folder[:-len('.json')] 53 | results = load_json(join(summary_output_folder, s))['results']['mean']['mean']['Dice'] 54 | f.write("%s,%s,%s,%s,%s,%02.4f\n" % (task, 55 | network, trainer, validation_folder, plans, results)) 56 | 57 | -------------------------------------------------------------------------------- /nnunet/evaluation/add_dummy_task_with_mean_over_all_tasks.py: -------------------------------------------------------------------------------- 1 | import json 2 | import numpy as np 3 | from batchgenerators.utilities.file_and_folder_operations import subfiles 4 | import os 5 | from collections import OrderedDict 6 | 7 | folder = "/home/fabian/drives/E132-Projekte/Projects/2018_MedicalDecathlon/Leaderboard" 8 | task_descriptors = ['2D final 2', 9 | '2D final, less pool, dc and topK, fold0', 10 | '2D final pseudo3d 7, fold0', 11 | '2D final, less pool, dc and ce, fold0', 12 | '3D stage0 final 2, fold0', 13 | '3D fullres final 2, fold0'] 14 | task_ids_with_no_stage0 = ["Task01_BrainTumour", "Task04_Hippocampus", "Task05_Prostate"] 15 | 16 | mean_scores = OrderedDict() 17 | for t in task_descriptors: 18 | mean_scores[t] = OrderedDict() 19 | 20 | json_files = subfiles(folder, True, None, ".json", True) 21 | json_files = [i for i in json_files if not i.split("/")[-1].startswith(".")] # stupid mac 22 | for j in json_files: 23 | with open(j, 'r') as f: 24 | res = json.load(f) 25 | task = res['task'] 26 | if task != "Task99_ALL": 27 | name = res['name'] 28 | if name in task_descriptors: 29 | if task not in list(mean_scores[name].keys()): 30 | mean_scores[name][task] = res['results']['mean']['mean'] 31 | else: 32 | raise RuntimeError("duplicate task %s for description %s" % (task, name)) 33 | 34 | for t in task_ids_with_no_stage0: 35 | mean_scores["3D stage0 final 2, fold0"][t] = mean_scores["3D fullres final 2, fold0"][t] 36 | 37 | a = set() 38 | for i in mean_scores.keys(): 39 | a = a.union(list(mean_scores[i].keys())) 40 | 41 | for i in mean_scores.keys(): 42 | try: 43 | for t in list(a): 44 | assert t in mean_scores[i].keys(), "did not find task %s for experiment %s" % (t, i) 45 | new_res = OrderedDict() 46 | new_res['name'] = i 47 | new_res['author'] = "Fabian" 48 | new_res['task'] = "Task99_ALL" 49 | new_res['results'] = OrderedDict() 50 | new_res['results']['mean'] = OrderedDict() 51 | new_res['results']['mean']['mean'] = OrderedDict() 52 | tasks = list(mean_scores[i].keys()) 53 | metrics = mean_scores[i][tasks[0]].keys() 54 | for m in metrics: 55 | foreground_values = [mean_scores[i][n][m] for n in tasks] 56 | new_res['results']['mean']["mean"][m] = np.nanmean(foreground_values) 57 | output_fname = i.replace(" ", "_") + "_globalMean.json" 58 | with open(os.path.join(folder, output_fname), 'w') as f: 59 | json.dump(new_res, f) 60 | except AssertionError: 61 | print("could not process experiment %s" % i) 62 | print("did not find task %s for experiment %s" % (t, i)) 63 | 64 | -------------------------------------------------------------------------------- /nnunet/paths.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | """ 4 | Put your personal paths in here. This file will shortly be added to gitignore so that your personal paths will not be tracked 5 | """ 6 | import os 7 | from batchgenerators.utilities.file_and_folder_operations import maybe_mkdir_p, join 8 | 9 | Server_Base_Path = os.path.join(os.environ['HOME'],'all_data') 10 | 11 | 12 | # You need to set the following folders: base, preprocessing_output_dir and network_training_output_dir. See below for details. 13 | 14 | # do not modify these unless you know what you are doing 15 | my_output_identifier = "CTPelvic1K" 16 | default_plans_identifier = "nnUNetPlans" 17 | default_data_identifier = 'nnUNet' 18 | 19 | try: 20 | # base is the folder where the raw data is stored. You just need to set base only, the others will be created 21 | # automatically (they are subfolders of base). 22 | # Here I use environment variables to set the base folder. Environment variables allow me to use the same code on 23 | # different systems (and our compute cluster). You can replace this line with something like: 24 | # base = "/path/to/my/folder" 25 | 26 | # base = os.environ['nnUNet_base'] 27 | base = f'{Server_Base_Path}/nnUNet' 28 | raw_dataset_dir = join(base, "nnUNet_raw") 29 | splitted_4d_output_dir = join(base, "nnUNet_raw_splitted") 30 | cropped_output_dir = join(base, "nnUNet_raw_cropped") 31 | maybe_mkdir_p(splitted_4d_output_dir) 32 | maybe_mkdir_p(raw_dataset_dir) 33 | maybe_mkdir_p(cropped_output_dir) 34 | except KeyError: 35 | cropped_output_dir = splitted_4d_output_dir = raw_dataset_dir = base = None 36 | 37 | # preprocessing_output_dir is where the preprocessed data is stored. If you run a training I very strongly recommend 38 | # this is a SSD! 39 | try: 40 | # Here I use environment variables to set the folder. Environment variables allow me to use the same code on 41 | # different systems (and our compute cluster). You can replace this line with something like: 42 | # preprocessing_output_dir = "/path/to/my/folder_with_preprocessed_data" 43 | 44 | # preprocessing_output_dir = os.environ['nnUNet_preprocessed'] 45 | preprocessing_output_dir = f'{Server_Base_Path}/nnUNet/nnUNet_processed' 46 | except KeyError: 47 | preprocessing_output_dir = None 48 | 49 | # This is where the trained model parameters are stored 50 | try: 51 | # Here I use environment variables to set the folder. Environment variables allow me to use the same code on 52 | # different systems (and our compute cluster). You can replace this line with something like: 53 | # network_training_output_dir = "/path/to/my/folder_with_results" 54 | 55 | # network_training_output_dir = os.path.join(os.environ['RESULTS_FOLDER'], my_output_identifier) 56 | network_training_output_dir = os.path.join(f'{Server_Base_Path}/nnUNet/nnUNet_results_folder', my_output_identifier) 57 | maybe_mkdir_p(network_training_output_dir) 58 | except KeyError: 59 | network_training_output_dir = None 60 | print("RESULTS_FOLDER was not in your environment variables, network_training_output_dir could not be determined. " 61 | "Please go to nnunet/paths.py and manually set network_training_output_dir. You can ignore this warning if " 62 | "you are using nnunet only as a toolkit and don't intend to run network trainings") 63 | -------------------------------------------------------------------------------- /nnunet/experiment_planning/summarize_plans.py: -------------------------------------------------------------------------------- 1 | from batchgenerators.utilities.file_and_folder_operations import * 2 | from nnunet.paths import preprocessing_output_dir 3 | 4 | 5 | # This file is intended to double check nnUNets design choices. It is intended to be used for developent purposes only 6 | def summarize_plans(file): 7 | plans = load_pickle(file) 8 | 9 | for i in range(len(plans['plans_per_stage'])): 10 | print("stage: ", i) 11 | [print(k,': ',v) for k,v in plans['plans_per_stage'][i].items()] 12 | print("") 13 | 14 | 15 | def write_plans_to_file(f, plans_file): 16 | a = load_pickle(plans_file) 17 | stages = list(a['plans_per_stage'].keys()) 18 | stages.sort() 19 | for stage in stages: 20 | patch_size_in_mm = [i * j for i, j in zip(a['plans_per_stage'][stages[stage]]['patch_size'], 21 | a['plans_per_stage'][stages[stage]]['current_spacing'])] 22 | median_patient_size_in_mm = [i * j for i, j in zip(a['plans_per_stage'][stages[stage]]['median_patient_size_in_voxels'], 23 | a['plans_per_stage'][stages[stage]]['current_spacing'])] 24 | f.write(plans_file.split("/")[-2]) 25 | f.write(";%s" % plans_file.split("/")[-1]) 26 | f.write(";%d" % stage) 27 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['batch_size'])) 28 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['num_pool_per_axis'])) 29 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['patch_size'])) 30 | f.write(";%s" % str([str("%03.2f" % i) for i in patch_size_in_mm])) 31 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['median_patient_size_in_voxels'])) 32 | f.write(";%s" % str([str("%03.2f" % i) for i in median_patient_size_in_mm])) 33 | f.write(";%s" % str([str("%03.2f" % i) for i in a['plans_per_stage'][stages[stage]]['current_spacing']])) 34 | f.write(";%s" % str([str("%03.2f" % i) for i in a['plans_per_stage'][stages[stage]]['original_spacing']])) 35 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['pool_op_kernel_sizes'])) 36 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['conv_kernel_sizes'])) 37 | f.write(";%s" % str(a['data_identifier'])) 38 | f.write("\n") 39 | 40 | 41 | if __name__ == "__main__": 42 | base_dir = preprocessing_output_dir 43 | task_dirs = [i for i in subdirs(preprocessing_output_dir, join=False, prefix="Task") if i.find("BrainTumor") == -1 and i.find("MSSeg") == -1] 44 | print("found %d tasks" % len(task_dirs)) 45 | """for t in task_dirs: 46 | print(t, "\n") 47 | tmp = join(preprocessing_output_dir, t) 48 | print("##### 2D #####") 49 | summarize_plans(join(tmp, my_plans_identifier + "_plans_2D.pkl")) 50 | print("##### 3D #####") 51 | summarize_plans(join(tmp, my_plans_identifier + "_plans_3D.pkl")) 52 | print("-------------------------------------------------------\n")""" 53 | 54 | with open("2019_02_06_plans_summary.csv", 'w') as f: 55 | f.write("task;plans_file;stage;batch_size;num_pool_per_axis;patch_size;patch_size(mm);median_patient_size_in_voxels;median_patient_size_in_mm;current_spacing;original_spacing;pool_op_kernel_sizes;conv_kernel_sizes\n") 56 | for t in task_dirs: 57 | print(t) 58 | tmp = join(preprocessing_output_dir, t) 59 | plans_files = [i for i in subfiles(tmp, suffix=".pkl", join=False) if i.find("_plans_") != -1 and i.find("Dgx2") == -1] 60 | for p in plans_files: 61 | write_plans_to_file(f, join(tmp, p)) 62 | 63 | 64 | -------------------------------------------------------------------------------- /nnunet/run/default_configuration.py: -------------------------------------------------------------------------------- 1 | import nnunet 2 | from nnunet.paths import network_training_output_dir, preprocessing_output_dir, default_plans_identifier 3 | from batchgenerators.utilities.file_and_folder_operations import * 4 | from nnunet.experiment_planning.summarize_plans import summarize_plans 5 | from nnunet.training.model_restore import recursive_find_trainer 6 | 7 | 8 | def get_configuration_from_output_folder(folder): 9 | folder = folder[len(network_training_output_dir):] 10 | if folder.startswith("/"): 11 | folder = folder[1:] 12 | 13 | configuration, task, trainer_and_plans_identifier = folder.split("/") 14 | trainer, plans_identifier = trainer_and_plans_identifier.split("__") 15 | return configuration, task, trainer, plans_identifier 16 | 17 | 18 | def get_output_folder(configuration, task, trainer, plans_identifier): 19 | return join(network_training_output_dir, configuration, task, trainer + "__" + plans_identifier) 20 | 21 | 22 | def get_default_configuration(network, task, network_trainer, plans_identifier=default_plans_identifier, 23 | search_in=(nnunet.__path__[0], 24 | "training", 25 | "network_training"), 26 | base_module='nnunet.training.network_training'): 27 | 28 | assert network in ['2d', '3d_lowres', '3d_fullres', '3d_cascade_fullres'], \ 29 | "network can only be one of the following: \'3d\', \'3d_lowres\', \'3d_fullres\', \'3d_cascade_fullres\'" 30 | 31 | dataset_directory = join(preprocessing_output_dir, task) 32 | 33 | if network == '2d': 34 | plans_file = join(preprocessing_output_dir, task, plans_identifier + "_plans_2D.pkl") 35 | else: 36 | plans_file = join(preprocessing_output_dir, task, plans_identifier + "_plans_3D.pkl") 37 | plans = load_pickle(plans_file) 38 | possible_stages = list(plans['plans_per_stage'].keys()) 39 | if (network == '3d_cascade_fullres' or network == "3d_lowres") and len(possible_stages) == 1:### !!!!! 40 | raise RuntimeError("3d_lowres/3d_cascade_fullres only applies if there is more than one stage. This task does " 41 | "not require the cascade. Run 3d_fullres instead") 42 | 43 | if network == '2d' or network == "3d_lowres": 44 | stage = 0 45 | else: 46 | stage = possible_stages[-1] 47 | 48 | trainer_class = recursive_find_trainer([join(*search_in)], network_trainer, current_module=base_module) 49 | 50 | print('\n~·~~·~~·~~·~~·~~·~~·~~·~ get default configuration ~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~\n') 51 | 52 | output_folder_name = join(network_training_output_dir, network, task, network_trainer + "__" + plans_identifier) 53 | 54 | print("Following configuration: ") 55 | summarize_plans(plans_file) 56 | 57 | if (network == '2d' or len(possible_stages) > 1) and not network == '3d_lowres': 58 | batch_dice = True 59 | print("Loss: BATCH dice + other losses") 60 | else: 61 | batch_dice = False 62 | print("Loss: Simple dice + other losses") 63 | 64 | print('network training output dir: ', network_training_output_dir) 65 | print('network: ', network) 66 | print('task: ', task) 67 | print("stage: ", stage) 68 | print("class of my trainer is: ", trainer_class) 69 | print("Training/validation data folder: ", join(dataset_directory, plans['data_identifier']+f'_stage{stage}')) 70 | print('~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~~·~\n') 71 | return plans_file, output_folder_name, dataset_directory, batch_dice, stage, trainer_class 72 | -------------------------------------------------------------------------------- /nnunet/training/network_training/nnUNet_variants/nnUNetTrainerNoDA.py: -------------------------------------------------------------------------------- 1 | import matplotlib 2 | from batchgenerators.utilities.file_and_folder_operations import maybe_mkdir_p, join 3 | from nnunet.network_architecture.neural_network import SegmentationNetwork 4 | from nnunet.training.data_augmentation.default_data_augmentation import get_no_augmentation 5 | from nnunet.training.dataloading.dataset_loading import unpack_dataset, DataLoader3D, DataLoader2D 6 | from nnunet.training.network_training.nnUNetTrainer import nnUNetTrainer 7 | from torch import nn 8 | 9 | matplotlib.use("agg") 10 | 11 | 12 | class nnUNetTrainerNoDA(nnUNetTrainer): 13 | def get_basic_generators(self): 14 | self.load_dataset() 15 | self.do_split() 16 | 17 | if self.threeD: 18 | dl_tr = DataLoader3D(self.dataset_tr, self.patch_size, self.patch_size, self.batch_size, 19 | False, oversample_foreground_percent=self.oversample_foreground_percent 20 | , pad_mode="constant", pad_sides=self.pad_all_sides) 21 | dl_val = DataLoader3D(self.dataset_val, self.patch_size, self.patch_size, self.batch_size, False, 22 | oversample_foreground_percent=self.oversample_foreground_percent, 23 | pad_mode="constant", pad_sides=self.pad_all_sides) 24 | else: 25 | dl_tr = DataLoader2D(self.dataset_tr, self.patch_size, self.patch_size, self.batch_size, 26 | transpose=self.plans.get('transpose_forward'), 27 | oversample_foreground_percent=self.oversample_foreground_percent 28 | , pad_mode="constant", pad_sides=self.pad_all_sides) 29 | dl_val = DataLoader2D(self.dataset_val, self.patch_size, self.patch_size, self.batch_size, 30 | transpose=self.plans.get('transpose_forward'), 31 | oversample_foreground_percent=self.oversample_foreground_percent, 32 | pad_mode="constant", pad_sides=self.pad_all_sides) 33 | return dl_tr, dl_val 34 | 35 | def initialize(self, training=True, force_load_plans=False): 36 | """ 37 | For prediction of test cases just set training=False, this will prevent loading of training data and 38 | training batchgenerator initialization 39 | :param training: 40 | :return: 41 | """ 42 | 43 | maybe_mkdir_p(self.output_folder) 44 | 45 | if force_load_plans or (self.plans is None): 46 | self.load_plans_file() 47 | 48 | self.process_plans(self.plans) 49 | 50 | self.setup_DA_params() 51 | 52 | self.folder_with_preprocessed_data = join(self.dataset_directory, self.plans['data_identifier'] + 53 | "_stage%d" % self.stage) 54 | if training: 55 | self.dl_tr, self.dl_val = self.get_basic_generators() 56 | if self.unpack_data: 57 | print("unpacking dataset") 58 | unpack_dataset(self.folder_with_preprocessed_data) 59 | print("done") 60 | else: 61 | print("INFO: Not unpacking data! Training may be slow due to that. Pray you are not using 2d or you " 62 | "will wait all winter for your model to finish!") 63 | self.tr_gen, self.val_gen = get_no_augmentation(self.dl_tr, self.dl_val, 64 | self.data_aug_params[ 65 | 'patch_size_for_spatialtransform'], 66 | self.data_aug_params) 67 | self.print_to_log_file("TRAINING KEYS:\n %s" % (str(self.dataset_tr.keys())), 68 | also_print_to_console=False) 69 | self.print_to_log_file("VALIDATION KEYS:\n %s" % (str(self.dataset_val.keys())), 70 | also_print_to_console=False) 71 | else: 72 | pass 73 | self.initialize_network_optimizer_and_scheduler() 74 | assert isinstance(self.network, (SegmentationNetwork, nn.DataParallel)) 75 | self.was_initialized = True 76 | self.data_aug_params['mirror_axes'] = () 77 | -------------------------------------------------------------------------------- /nnunet/inference/ensemble_predictions.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | from nnunet.inference.segmentation_export import save_segmentation_nifti_from_softmax 4 | from batchgenerators.utilities.file_and_folder_operations import * 5 | import numpy as np 6 | from multiprocessing import Pool 7 | 8 | 9 | def merge_files(args): 10 | files, properties_file, out_file, only_keep_largest_connected_component, min_region_size_per_class, override = args 11 | if override or not isfile(out_file): 12 | softmax = [np.load(f)['softmax'][None] for f in files] 13 | softmax = np.vstack(softmax) 14 | softmax = np.mean(softmax, 0) 15 | props = load_pickle(properties_file) 16 | save_segmentation_nifti_from_softmax(softmax, out_file, props, 1, None, None, None) 17 | 18 | 19 | def merge(folders, output_folder, threads, override=True): 20 | maybe_mkdir_p(output_folder) 21 | 22 | patient_ids = [subfiles(i, suffix=".npz", join=False) for i in folders] 23 | patient_ids = [i for j in patient_ids for i in j] 24 | patient_ids = [i[:-4] for i in patient_ids] 25 | patient_ids = np.unique(patient_ids) 26 | 27 | for f in folders: 28 | assert all([isfile(join(f, i + ".npz")) for i in patient_ids]), "Not all patient npz are available in " \ 29 | "all folders" 30 | assert all([isfile(join(f, i + ".pkl")) for i in patient_ids]), "Not all patient pkl are available in " \ 31 | "all folders" 32 | 33 | files = [] 34 | property_files = [] 35 | out_files = [] 36 | for p in patient_ids: 37 | files.append([join(f, p + ".npz") for f in folders]) 38 | property_files.append(join(folders[0], p + ".pkl")) 39 | out_files.append(join(output_folder, p + ".nii.gz")) 40 | 41 | plans = load_pickle(join(folders[0], "plans.pkl")) 42 | 43 | only_keep_largest_connected_component, min_region_size_per_class = plans['keep_only_largest_region'], \ 44 | plans['min_region_size_per_class'] 45 | p = Pool(threads) 46 | p.map(merge_files, zip(files, property_files, out_files, [only_keep_largest_connected_component] * len(out_files), 47 | [min_region_size_per_class] * len(out_files), [override] * len(out_files))) 48 | p.close() 49 | p.join() 50 | 51 | 52 | if __name__ == "__main__": 53 | import argparse 54 | parser = argparse.ArgumentParser(description="This requires that all folders to be merged use the same " 55 | "postprocessing function " 56 | "(nnunet.utilities.postprocessing.postprocess_segmentation). " 57 | "This will be the case if the corresponding " 58 | "models were trained with nnUNetTrainer or nnUNetTrainerCascadeFullRes" 59 | "but may not be the case if you added models of your own that use a " 60 | "different postprocessing. This script also requires a plans file to" 61 | "be present in all of the folders (if they are not present you can " 62 | "take them from the respective model training output folders. " 63 | "Parameters for the postprocessing " 64 | "will be taken from the plans file. If the folders were created by " 65 | "predict_folder.py then the plans file will have been copied " 66 | "automatically (if --save_npz is specified)") 67 | parser.add_argument('-f', '--folders', nargs='+', help="list of folders to merge. All folders must contain npz " 68 | "files", required=True) 69 | parser.add_argument('-o', '--output_folder', help="where to save the results", required=True, type=str) 70 | parser.add_argument('-t', '--threads', help="number of threads used to saving niftis", required=False, default=2, 71 | type=int) 72 | 73 | args = parser.parse_args() 74 | 75 | folders = args.folders 76 | threads = args.threads 77 | output_folder = args.output_folder 78 | 79 | merge(folders, output_folder, threads, override=True) 80 | -------------------------------------------------------------------------------- /nnunet/training/data_augmentation/custom_transforms.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from batchgenerators.transforms import AbstractTransform 3 | 4 | 5 | class RemoveKeyTransform(AbstractTransform): 6 | def __init__(self, key_to_remove): 7 | self.key_to_remove = key_to_remove 8 | 9 | def __call__(self, **data_dict): 10 | _ = data_dict.pop(self.key_to_remove, None) 11 | return data_dict 12 | 13 | 14 | class MaskTransform(AbstractTransform): 15 | def __init__(self, dct_for_where_it_was_used, mask_idx_in_seg=1, set_outside_to=0, data_key="data", seg_key="seg"): 16 | """ 17 | data[mask < 0] = 0 18 | Sets everything outside the mask to 0. CAREFUL! outside is defined as < 0, not =0 (in the Mask)!!! 19 | 20 | :param dct_for_where_it_was_used: 21 | :param mask_idx_in_seg: 22 | :param set_outside_to: 23 | :param data_key: 24 | :param seg_key: 25 | """ 26 | self.dct_for_where_it_was_used = dct_for_where_it_was_used 27 | self.seg_key = seg_key 28 | self.data_key = data_key 29 | self.set_outside_to = set_outside_to 30 | self.mask_idx_in_seg = mask_idx_in_seg 31 | 32 | def __call__(self, **data_dict): 33 | seg = data_dict.get(self.seg_key) 34 | if seg is None or seg.shape[1] < self.mask_idx_in_seg: 35 | raise Warning("mask not found, seg may be missing or seg[:, mask_idx_in_seg] may not exist") 36 | data = data_dict.get(self.data_key) 37 | for b in range(data.shape[0]): 38 | mask = seg[b, self.mask_idx_in_seg] 39 | for c in range(data.shape[1]): 40 | if self.dct_for_where_it_was_used[c]: 41 | data[b, c][mask < 0] = self.set_outside_to 42 | data_dict[self.data_key] = data 43 | return data_dict 44 | 45 | 46 | def convert_3d_to_2d_generator(data_dict): 47 | shp = data_dict['data'].shape 48 | data_dict['data'] = data_dict['data'].reshape((shp[0], shp[1] * shp[2], shp[3], shp[4])) 49 | data_dict['orig_shape_data'] = shp 50 | shp = data_dict['seg'].shape 51 | data_dict['seg'] = data_dict['seg'].reshape((shp[0], shp[1] * shp[2], shp[3], shp[4])) 52 | data_dict['orig_shape_seg'] = shp 53 | return data_dict 54 | 55 | 56 | def convert_2d_to_3d_generator(data_dict): 57 | shp = data_dict['orig_shape_data'] 58 | current_shape = data_dict['data'].shape 59 | data_dict['data'] = data_dict['data'].reshape((shp[0], shp[1], shp[2], current_shape[-2], current_shape[-1])) 60 | shp = data_dict['orig_shape_seg'] 61 | current_shape_seg = data_dict['seg'].shape 62 | data_dict['seg'] = data_dict['seg'].reshape((shp[0], shp[1], shp[2], current_shape_seg[-2], current_shape_seg[-1])) 63 | return data_dict 64 | 65 | 66 | class Convert3DTo2DTransform(AbstractTransform): 67 | def __init__(self): 68 | pass 69 | 70 | def __call__(self, **data_dict): 71 | return convert_3d_to_2d_generator(data_dict) 72 | 73 | 74 | class Convert2DTo3DTransform(AbstractTransform): 75 | def __init__(self): 76 | pass 77 | 78 | def __call__(self, **data_dict): 79 | return convert_2d_to_3d_generator(data_dict) 80 | 81 | 82 | class ConvertSegmentationToRegionsTransform(AbstractTransform): 83 | def __init__(self, regions, seg_key="seg", output_key="seg", seg_channel=0): 84 | """ 85 | regions are tuple of tuples where each inner tuple holds the class indices that are merged into one region, example: 86 | regions= ((1, 2), (2, )) will result in 2 regions: one covering the region of labels 1&2 and the other just 2 87 | :param regions: 88 | :param seg_key: 89 | :param output_key: 90 | """ 91 | self.seg_channel = seg_channel 92 | self.output_key = output_key 93 | self.seg_key = seg_key 94 | self.regions = regions 95 | 96 | def __call__(self, **data_dict): 97 | seg = data_dict.get(self.seg_key) 98 | num_regions = len(self.regions) 99 | if seg is not None: 100 | seg_shp = seg.shape 101 | output_shape = list(seg_shp) 102 | output_shape[1] = num_regions 103 | region_output = np.zeros(output_shape, dtype=seg.dtype) 104 | for b in range(seg_shp[0]): 105 | for r in range(num_regions): 106 | for l in self.regions[r]: 107 | region_output[b, r][seg[b, self.seg_channel] == l] = 1 108 | data_dict[self.output_key] = region_output 109 | return data_dict 110 | -------------------------------------------------------------------------------- /nnunet/dataset_conversion/JstPelvisSegmentation_5label.py: -------------------------------------------------------------------------------- 1 | from collections import OrderedDict 2 | import SimpleITK as sitk 3 | from batchgenerators.utilities.file_and_folder_operations import * 4 | from multiprocessing import Pool 5 | import numpy as np 6 | from scipy.ndimage import label 7 | import shutil 8 | import argparse 9 | 10 | 11 | def export_segmentations(indir, outdir): 12 | niftis = subfiles(indir, suffix='nii.gz', join=False) 13 | for n in niftis: 14 | identifier = str(n.split("_")[-1][:-7]) 15 | outfname = join(outdir, "test-segmentation-%s.nii" % identifier) 16 | img = sitk.ReadImage(join(indir, n)) 17 | sitk.WriteImage(img, outfname) 18 | 19 | 20 | def export_segmentations_postprocess(indir, outdir): 21 | maybe_mkdir_p(outdir) 22 | niftis = subfiles(indir, suffix='nii.gz', join=False) 23 | for n in niftis: 24 | print("\n", n) 25 | identifier = str(n.split("_")[-1][:-7]) 26 | outfname = join(outdir, "test-segmentation-%s.nii" % identifier) 27 | img = sitk.ReadImage(join(indir, n)) 28 | img_npy = sitk.GetArrayFromImage(img) 29 | lmap, num_objects = label((img_npy > 0).astype(int)) 30 | sizes = [] 31 | for o in range(1, num_objects + 1): 32 | sizes.append((lmap == o).sum()) 33 | mx = np.argmax(sizes) + 1 34 | print(sizes) 35 | img_npy[lmap != mx] = 0 36 | img_new = sitk.GetImageFromArray(img_npy) 37 | img_new.CopyInformation(img) 38 | sitk.WriteImage(img_new, outfname) 39 | 40 | 41 | if __name__ == "__main__": 42 | parser = argparse.ArgumentParser() 43 | parser.add_argument("--train_dir", type=str, default=None, 44 | help="path to dataset CTPelvic1K") 45 | parser.add_argument("--output_dir", type=str, default=None, 46 | help="random seed (default: 42)") 47 | opts = parser.parse_args() 48 | """ 49 | dataset format: 50 | image: *_data.nii.gz 51 | label: *_mask_4label.nii.gz 52 | test: 53 | image: *_data.nii.gz 54 | Data will be converted into a unified format used in nnunet 55 | """ 56 | train_dir = opts.train_dir 57 | output_folder = opts.output_dir 58 | test_dir = "/path/to/testing dataset" 59 | 60 | img_dir = join(output_folder, "imagesTr") 61 | lab_dir = join(output_folder, "labelsTr") 62 | img_dir_te = join(output_folder, "imagesTs") 63 | maybe_mkdir_p(img_dir) 64 | maybe_mkdir_p(lab_dir) 65 | maybe_mkdir_p(img_dir_te) 66 | 67 | 68 | def load_save_train(args): 69 | data_file, seg_file = args 70 | pat_id = data_file.split("/")[-1] 71 | pat_id = "train_" + pat_id.split("-")[-1][:-12] 72 | 73 | shutil.copy(data_file, join(img_dir, pat_id + ".nii.gz")) 74 | 75 | shutil.copy(seg_file, join(lab_dir, pat_id + ".nii.gz")) 76 | return pat_id 77 | def load_save_test(args): 78 | data_file = args 79 | pat_id = data_file.split("/")[-1] 80 | pat_id = "test_" + pat_id.split("-")[-1][:-12] 81 | 82 | shutil.copy(data_file, join(img_dir_te, pat_id + ".nii.gz")) 83 | return pat_id 84 | 85 | 86 | nii_files_tr_data = subfiles(train_dir, True, None, "_data.nii.gz", True) 87 | nii_files_tr_seg = subfiles(train_dir, True, None, "_mask_4label.nii.gz", True) 88 | 89 | nii_files_ts = subfiles(test_dir, True, None, "_data.nii.gz", True) 90 | 91 | p = Pool(16) 92 | train_ids = p.map(load_save_train, zip(nii_files_tr_data, nii_files_tr_seg)) 93 | test_ids = p.map(load_save_test, nii_files_ts) 94 | p.close() 95 | p.join() 96 | 97 | json_dict = OrderedDict() 98 | json_dict['name'] = "CTPelvic1K_4label" 99 | json_dict['description'] = "CTPelvic1K_4label" 100 | json_dict['tensorImageSize'] = "3D" 101 | json_dict['reference'] = "see challenge website" 102 | json_dict['licence'] = "see challenge website" 103 | json_dict['release'] = "0.1" 104 | json_dict['modality'] = { 105 | "0": "CT" 106 | } 107 | 108 | json_dict['labels'] = { 109 | "0": "background", 110 | "1": "sacrum", 111 | "2": "right_hip", 112 | "3": "left_hip", 113 | "4": "lumbar_vertebra" 114 | } 115 | 116 | json_dict['numTraining'] = len(train_ids) 117 | 118 | json_dict['training'] = [{'image': "./imagesTr/%s.nii.gz" % i, 119 | "label": "./labelsTr/%s.nii.gz" % i} for i in train_ids] 120 | 121 | json_dict['test'] = ["./imagesTs/%s.nii.gz" % i for i in test_ids] 122 | 123 | with open(os.path.join(output_folder, "dataset.json"), 'w') as f: 124 | json.dump(json_dict, f, indent=4, sort_keys=True) 125 | -------------------------------------------------------------------------------- /nnunet/preprocessing/lumbosacral_joint_sampling.py: -------------------------------------------------------------------------------- 1 | """ 2 | make a dataset for lumbosacral joint prediction. 3 | three class: background, sacrum, lumbar - 0,1,2 4 | 5 | """ 6 | import numpy as np 7 | from multiprocessing import Pool 8 | from functools import partial 9 | import cv2 10 | import os 11 | import pickle as pkl 12 | 13 | def get_reasonable_crops_of_lumbar(gt3d, patch_size, stage): 14 | d, h, w = gt3d.shape 15 | print((d,h,w)) 16 | 17 | lambar = np.where(gt3d == 4, 1, 0) 18 | sacral = np.where(gt3d == 1, 1, 0) 19 | upside_down = False 20 | d_lambar = np.argwhere(lambar > 0)[:, 0] 21 | d_sacral = np.argwhere(sacral > 0)[:, 0] 22 | if d_sacral.mean() > d_lambar.mean(): 23 | upside_down = True 24 | print(upside_down) 25 | 26 | if lambar.sum()==0: 27 | if stage==0: 28 | c_s = 50 29 | elif stage==1: 30 | c_s = 100 31 | if upside_down: 32 | return [(0, int(patch_size[0])), 33 | (int(h//2-c_s), int(h//2+c_s)), 34 | (int(w//2-c_s), int(w//2+c_s))] 35 | else: 36 | return [(int(d-patch_size[0]), int(d-1)), 37 | (int(h // 2 - c_s), int(h // 2 + c_s)), 38 | (int(w // 2 - c_s), int(w // 2 + c_s))] 39 | 40 | if upside_down: 41 | lambar_boundary_d = sorted(set(d_lambar))[-5]+4 42 | else: 43 | lambar_boundary_d = sorted(set(d_lambar))[4]-4 44 | 45 | h_lambar = np.argwhere(lambar > 0)[:, 1] 46 | w_lambar = np.argwhere(lambar > 0)[:, 2] 47 | lambar_center_h = (max(h_lambar)+min(h_lambar))//2 48 | lambar_center_w = (max(w_lambar)+min(w_lambar))//2 49 | print(max(h_lambar), min(h_lambar)) 50 | print(max(w_lambar), min(w_lambar)) 51 | print(lambar_boundary_d, lambar_center_h, lambar_center_w) 52 | 53 | valid_crop = [(int(max(0, lambar_boundary_d-patch_size[0]//2)), int(min(d, lambar_boundary_d + patch_size[0]//2))), 54 | (int(max(0, lambar_center_h-patch_size[1]//2)), int(min(h, lambar_center_h + patch_size[1]//2))), 55 | (int(max(0, lambar_center_w-patch_size[2]//2)), int(min(w, lambar_center_w + patch_size[2]//2)))] 56 | 57 | return valid_crop 58 | 59 | def _main_3d_one_case(name, path, crop_size, stage, img_save_path): 60 | print(name) 61 | try: 62 | array = np.load(path + '/' + name + '.npy') 63 | except Exception as _: 64 | array = np.load(path + '/' + name + '.npz')['data'] 65 | 66 | with open(path + '/' + name + '.pkl', 'rb') as f: 67 | info = pkl.load(f) 68 | 69 | reasonal = get_reasonable_crops_of_lumbar(array[-1, :, :, :], patch_size=crop_size, stage=stage) 70 | info['Lumbosacral_Region'] = reasonal 71 | 72 | # backup raw pkl file 73 | os.rename(path + '/' + name + '.pkl', path + '/' + name + '_backup.pkl') 74 | # save new pkl file 75 | with open(path + '/' + name + '.pkl', 'wb') as f: 76 | pkl.dump(info, f) 77 | print(reasonal) 78 | 79 | # check reasonal 80 | saveimg = array[0][(reasonal[0][0] + reasonal[0][1]) // 2, reasonal[1][0]:reasonal[1][1], 81 | reasonal[2][0]:reasonal[2][1]] 82 | saveimg = (saveimg - saveimg.min()) / (saveimg.max() - saveimg.min()) * 255 83 | cv2.imwrite(img_save_path + '/' + name + '.png', saveimg) 84 | def main_3d(base_path, check_save_path): 85 | base_path = base_path 86 | stage = 1 87 | 88 | path = f'{base_path}/nnUNet_stage{stage}' 89 | plans_path = f'{base_path}/nnUNetPlans_plans_3D.pkl' 90 | 91 | names = os.listdir(path) 92 | names = [i[:-4] for i in names if i.endswith('.npy')] 93 | 94 | print(len(names),'files to process...') 95 | 96 | with open(plans_path, "rb") as f: 97 | plans_info = pkl.load(f) 98 | 99 | crop_size = plans_info["plans_per_stage"][stage]["patch_size"] # array([ 96, 160, 128]) 100 | crop_size[2] = 180 # enlarge 101 | 102 | assert stage==1 103 | 104 | img_save_path = check_save_path 105 | if not os.path.exists(img_save_path): 106 | os.makedirs(img_save_path) 107 | 108 | pool = Pool() 109 | func = partial(_main_3d_one_case, path=path, crop_size=crop_size, stage=stage, img_save_path=img_save_path) 110 | _ = pool.map(func, names) 111 | pool.close() 112 | pool.join() 113 | 114 | if __name__ == '__main__': 115 | import argparse 116 | parser = argparse.ArgumentParser() 117 | parser.add_argument("--processed_path", type=str, default=None, 118 | help="path to dataset CTPelvic1K") 119 | parser.add_argument("--check_save_path", type=str, default=None, 120 | help="random seed (default: 42)") 121 | opts = parser.parse_args() 122 | 123 | main_3d(base_path=opts.processed_path, 124 | check_save_path=opts.check_save_path) 125 | 126 | -------------------------------------------------------------------------------- /save_evaluation_results2csv.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | import pickle as pkl 3 | import numpy as np 4 | import os 5 | 6 | 7 | if __name__ == '__main__': 8 | base_dir = os.environ['HOME'] 9 | eval_reslult_pkl_path = base_dir + '/all_data/nnUNet/rawdata/ipcai2021_ALL_Test/' \ 10 | 'Task22_ipcai2021_T__nnUNet_without_mirror_IPCAI2021_deeps_exclusion__nnUNet_without_mirror_IPCAI2021_deeps_exclusion__fold0_3dcascadefullres_pred/' \ 11 | 'evaluation_mcr__2000_False.pkl' 12 | 13 | print(eval_reslult_pkl_path) 14 | with open(eval_reslult_pkl_path,'rb') as f: 15 | eval_reslult = pkl.load(f) # dict of names and quality sub-dict 16 | 17 | datasets = ['dataset1','dataset2','dataset3','dataset4','dataset5','dataset6'] 18 | for dataset in datasets: 19 | names = [] 20 | bone1_Hausdorff = [] 21 | bone1_Dice = [] 22 | bone2_Hausdorff = [] 23 | bone2_Dice = [] 24 | bone3_Hausdorff = [] 25 | bone3_Dice = [] 26 | bone4_Hausdorff = [] 27 | bone4_Dice = [] 28 | whole_Hausdorff = [] 29 | whole_Dice = [] 30 | mean_Hausdorff = [] 31 | mean_Dice = [] 32 | weight_Hausdorff = [] 33 | weight_Dice = [] 34 | for na, quality in eval_reslult.items(): 35 | if not dataset in na: 36 | continue 37 | names.append(na) 38 | bone1_Hausdorff.append(quality[1]['Hausdorff']) 39 | bone1_Dice.append(quality[1]['dice']) 40 | bone2_Hausdorff.append(quality[2]['Hausdorff']) 41 | bone2_Dice.append(quality[2]['dice']) 42 | bone3_Hausdorff.append(quality[3]['Hausdorff']) 43 | bone3_Dice.append(quality[3]['dice']) 44 | bone4_Hausdorff.append(quality[4]['Hausdorff']) 45 | bone4_Dice.append(quality[4]['dice']) 46 | whole_Hausdorff.append(quality['whole']['Hausdorff']) 47 | whole_Dice.append(quality['whole']['dice']) 48 | mean_Hausdorff.append(quality['mean_hausdorff']) 49 | mean_Dice.append(quality['mean_dice']) 50 | weight_Hausdorff.append(quality['weighted_mean_hausdorff']) 51 | weight_Dice.append(quality['weighted_mean_dice']) 52 | print(dataset, len(names)) 53 | print( 54 | 'bone1_Dice:', np.array(bone1_Dice).mean(),'\n', 55 | 'bone1_Hausdorff:', np.array(bone1_Hausdorff).mean(),'\n', 56 | 'bone2_Dice:', np.array(bone2_Dice).mean(),'\n', 57 | 'bone2_Hausdorff:', np.array(bone2_Hausdorff).mean(),'\n', 58 | 'bone3_Dice:', np.array(bone3_Dice).mean(),'\n', 59 | 'bone3_Hausdorff:', np.array(bone3_Hausdorff).mean(),'\n', 60 | 'bone4_Dice:', np.array(bone4_Dice).mean(),'\n', 61 | 'bone4_Hausdorff:', np.array(bone4_Hausdorff).mean(),'\n', 62 | 'whole_Dice:', np.array(whole_Dice).mean(),'\n', 63 | 'whole_Hausdorff:', np.array(whole_Hausdorff).mean(),'\n', 64 | 'mean_Dice:', np.array(mean_Dice).mean(),'\n', 65 | 'mean_Hausdorff:', np.array(mean_Hausdorff).mean(),'\n', 66 | 'weight_Dice:', np.array(weight_Dice).mean(),'\n', 67 | 'weight_Hausdorff:', np.array(weight_Hausdorff).mean(),'\n', 68 | ) 69 | assert len(names)==len(bone1_Dice) 70 | assert len(names)==len(bone1_Hausdorff) 71 | assert len(names)==len(bone2_Dice) 72 | assert len(names)==len(bone2_Hausdorff) 73 | assert len(names)==len(bone3_Dice) 74 | assert len(names)==len(bone3_Hausdorff) 75 | assert len(names)==len(bone4_Dice) 76 | assert len(names)==len(bone4_Hausdorff) 77 | assert len(names)==len(whole_Dice) 78 | assert len(names)==len(whole_Hausdorff) 79 | assert len(names)==len(mean_Dice) 80 | assert len(names)==len(mean_Hausdorff) 81 | assert len(names)==len(weight_Dice) 82 | assert len(names)==len(weight_Hausdorff) 83 | 84 | results = {'names':names, 85 | 'bone1_Dice': bone1_Dice, 86 | 'bone1_Hausdorff': bone1_Hausdorff, 87 | 'bone2_Dice': bone2_Dice, 88 | 'bone2_Hausdorff': bone2_Hausdorff, 89 | 'bone3_Dice': bone3_Dice, 90 | 'bone3_Hausdorff': bone3_Hausdorff, 91 | 'bone4_Dice': bone4_Dice, 92 | 'bone4_Hausdorff': bone4_Hausdorff, 93 | 'whole_Dice': whole_Dice, 94 | 'whole_Hausdorff':whole_Hausdorff, 95 | 'mean_Dice':mean_Dice, 96 | 'mean_Hausdorff':mean_Hausdorff, 97 | 'weight_Dice':weight_Dice, 98 | 'weight_Hausdorff':weight_Hausdorff 99 | } 100 | results_pd = pd.DataFrame(results) 101 | save_csv_path = eval_reslult_pkl_path.replace('.pkl','_{}.csv'.format(dataset)) 102 | results_pd.to_csv(save_csv_path) 103 | 104 | 105 | 106 | 107 | -------------------------------------------------------------------------------- /nnunet/evaluation/model_selection/figure_out_what_to_submit.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import nnunet 4 | from batchgenerators.utilities.file_and_folder_operations import * 5 | from nnunet.paths import network_training_output_dir 6 | import numpy as np 7 | from nnunet.evaluation.add_mean_dice_to_json import foreground_mean 8 | from subprocess import call 9 | import SimpleITK as sitk 10 | from nnunet.run.default_configuration import get_output_folder 11 | 12 | 13 | def copy_nifti_and_convert_to_uint8(args): 14 | source_file, target_file = args 15 | i = sitk.ReadImage(source_file) 16 | j = sitk.GetImageFromArray(sitk.GetArrayFromImage(i).astype(np.uint8)) 17 | j.CopyInformation(i) 18 | sitk.WriteImage(j, target_file) 19 | 20 | 21 | if __name__ == "__main__": 22 | # This script was hacked together at some point and is ugly af. TODO this needs to be redone properly 23 | import argparse 24 | parser = argparse.ArgumentParser(usage="This is intended to identify the best model based on the five fold " 25 | "cross-validation. Running this script requires alle models to have been run " 26 | "already. This script will summarize the results of the five folds of all " 27 | "models in one json each for easy interpretability") 28 | parser.add_argument("-m", '--models', nargs="+", required=False, default=['2d', '3d_lowres', '3d_fullres', '3d_cascade_fullres']) 29 | parser.add_argument("-t", '--task_ids', nargs="+", required=False, default=list(range(100))) 30 | 31 | args = parser.parse_args() 32 | tasks = args.task_ids 33 | models = args.models 34 | 35 | out_dir_all_json = join(network_training_output_dir, "summary_jsons") 36 | 37 | json_files = [i for i in subfiles(out_dir_all_json, suffix=".json", join=True) if i.find("ensemble") == -1] 38 | 39 | # do mean over foreground 40 | for j in json_files: 41 | foreground_mean(j) 42 | 43 | # for each task, run ensembling using all combinations of two models 44 | for t in tasks: 45 | t = int(t) 46 | json_files_task = [i for i in subfiles(out_dir_all_json, prefix="Task%02.0d_" % t) if i.find("ensemble") == -1] 47 | if len(json_files_task) > 0: 48 | task_name = json_files_task[0].split("/")[-1].split("__")[0] 49 | print(task_name) 50 | 51 | for i in range(len(json_files_task) - 1): 52 | for j in range(i+1, len(json_files_task)): 53 | # networks are stored as 54 | # task__configuration__trainer__plans 55 | network1 = json_files_task[i].split("/")[-1].split("__") 56 | network1[-1] = network1[-1].split(".")[0] 57 | task, configuration, trainer, plans_identifier, _ = network1 58 | network1_folder = get_output_folder(configuration, task, trainer, plans_identifier) 59 | name1 = configuration + "__" + trainer + "__" + plans_identifier 60 | 61 | network2 = json_files_task[j].split("/")[-1].split("__") 62 | network2[-1] = network2[-1].split(".")[0] 63 | task, configuration, trainer, plans_identifier, _ = network2 64 | network2_folder = get_output_folder(configuration, task, trainer, plans_identifier) 65 | name2 = configuration + "__" + trainer + "__" + plans_identifier 66 | 67 | if np.argsort((name1, name2))[0] == 1: 68 | name1, name2 = name2, name1 69 | network1_folder, network2_folder = network2_folder, network1_folder 70 | 71 | output_folder = join(network_training_output_dir, "ensembles", task_name, "ensemble_" + name1 + "--" + name2) 72 | # now ensemble 73 | print(network1_folder, network2_folder) 74 | p = call(["python", join(nnunet.__path__[0], "evaluation/model_selection/ensemble.py"), network1_folder, network2_folder, output_folder, task_name]) 75 | 76 | # now rerun adding the mean foreground dice 77 | json_files = subfiles(out_dir_all_json, suffix=".json", join=True) 78 | 79 | # do mean over foreground 80 | for j in json_files: 81 | foreground_mean(j) 82 | 83 | # now load all json for each task and find best 84 | with open(join(network_training_output_dir, "use_this_for_test.csv"), 'w') as f: 85 | for t in tasks: 86 | t = int(t) 87 | json_files_task = subfiles(out_dir_all_json, prefix="Task%02.0d_" % t) 88 | if len(json_files_task) > 0: 89 | task_name = json_files_task[0].split("/")[-1].split("__")[0] 90 | print(task_name) 91 | mean_dice = [] 92 | for j in json_files_task: 93 | js = load_json(j) 94 | mean_dice.append(js['results']['mean']['mean']['Dice']) 95 | best = np.argsort(mean_dice)[::-1][0] 96 | j = json_files_task[best].split("/")[-1] 97 | 98 | print("%s: submit model %s" % (task_name, j)) 99 | f.write("%s,%s\n" % (task_name, j)) 100 | -------------------------------------------------------------------------------- /nnunet/evaluation/model_selection/summarize_results_in_one_json.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | from collections import OrderedDict 4 | from nnunet.evaluation.add_mean_dice_to_json import foreground_mean 5 | from batchgenerators.utilities.file_and_folder_operations import * 6 | from nnunet.paths import network_training_output_dir 7 | import numpy as np 8 | 9 | 10 | def summarize(tasks, models=('2d', '3d_lowres', '3d_fullres', '3d_cascade_fullres'), 11 | output_dir=join(network_training_output_dir, "summary_jsons"), folds=(0, 1, 2, 3, 4)): 12 | maybe_mkdir_p(output_dir) 13 | 14 | if len(tasks) == 1 and tasks[0] == "all": 15 | tasks = list(range(100)) 16 | else: 17 | tasks = [int(i) for i in tasks] 18 | 19 | for model in models: 20 | for t in tasks: 21 | t = int(t) 22 | if not isdir(join(network_training_output_dir, model)): 23 | continue 24 | task_name = subfolders(join(network_training_output_dir, model), prefix="Task%02.0d" % t, join=False) 25 | if len(task_name) != 1: 26 | print("did not find unique output folder for network %s and task %s" % (model, t)) 27 | continue 28 | task_name = task_name[0] 29 | out_dir_task = join(network_training_output_dir, model, task_name) 30 | 31 | model_trainers = subdirs(out_dir_task, join=False) 32 | for trainer in model_trainers: 33 | if trainer.startswith("fold"): 34 | continue 35 | out_dir = join(out_dir_task, trainer) 36 | 37 | validation_folders = [] 38 | for fld in folds: 39 | d = join(out_dir, "fold%d"%fld) 40 | if not isdir(d): 41 | d = join(out_dir, "fold_%d"%fld) 42 | if not isdir(d): 43 | break 44 | validation_folders += subfolders(d, prefix="validation", join=False) 45 | 46 | for v in validation_folders: 47 | ok = True 48 | metrics = OrderedDict() 49 | for fld in folds: 50 | d = join(out_dir, "fold%d"%fld) 51 | if not isdir(d): 52 | d = join(out_dir, "fold_%d"%fld) 53 | if not isdir(d): 54 | ok = False 55 | break 56 | validation_folder = join(d, v) 57 | 58 | if not isfile(join(validation_folder, "summary.json")): 59 | print("summary.json missing for net %s task %s fold %d" % (model, task_name, fld)) 60 | ok = False 61 | break 62 | 63 | metrics_tmp = load_json(join(validation_folder, "summary.json"))["results"]["mean"] 64 | for l in metrics_tmp.keys(): 65 | if metrics.get(l) is None: 66 | metrics[l] = OrderedDict() 67 | for m in metrics_tmp[l].keys(): 68 | if metrics[l].get(m) is None: 69 | metrics[l][m] = [] 70 | metrics[l][m].append(metrics_tmp[l][m]) 71 | if ok: 72 | for l in metrics.keys(): 73 | for m in metrics[l].keys(): 74 | assert len(metrics[l][m]) == len(folds) 75 | metrics[l][m] = np.mean(metrics[l][m]) 76 | json_out = OrderedDict() 77 | json_out["results"] = OrderedDict() 78 | json_out["results"]["mean"] = metrics 79 | json_out["task"] = task_name 80 | json_out["description"] = model + " " + task_name + " all folds summary" 81 | json_out["name"] = model + " " + task_name + " all folds summary" 82 | json_out["experiment_name"] = model 83 | save_json(json_out, join(out_dir, "summary_allFolds__%s.json" % v)) 84 | save_json(json_out, join(output_dir, "%s__%s__%s__%s.json" % (task_name, model, trainer, v))) 85 | foreground_mean(join(out_dir, "summary_allFolds__%s.json" % v)) 86 | foreground_mean(join(output_dir, "%s__%s__%s__%s.json" % (task_name, model, trainer, v))) 87 | 88 | 89 | if __name__ == "__main__": 90 | import argparse 91 | parser = argparse.ArgumentParser(usage="This is intended to identify the best model based on the five fold " 92 | "cross-validation. Running this script requires alle models to have been run " 93 | "already. This script will summarize the results of the five folds of all " 94 | "models in one json each for easy interpretability") 95 | parser.add_argument("-t", '--task_ids', nargs="+", required=True, help="task id. can be 'all'") 96 | parser.add_argument("-f", '--folds', nargs="+", required=False, type=int, default=[0, 1, 2, 3, 4]) 97 | parser.add_argument("-m", '--models', nargs="+", required=False, default=['2d', '3d_lowres', '3d_fullres', '3d_cascade_fullres']) 98 | 99 | args = parser.parse_args() 100 | tasks = args.task_ids 101 | models = args.models 102 | 103 | folds = args.folds 104 | summarize(tasks, models, folds=folds) 105 | 106 | -------------------------------------------------------------------------------- /save_evaluation_results2csv_Manu.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | import pickle as pkl 3 | import numpy as np 4 | import shutil 5 | import os 6 | 7 | def func(): 8 | fold = 21 9 | base_dir = os.environ['HOME'] 10 | eval_reslult_pkl_path = base_dir + '/all_data/nnUNet/rawdata/ipcai2021_M_Test/' \ 11 | f'Task22_ipcai2021_T__nnUNet_without_mirror_IPCAI2021_deeps_exclusion__nnUNet_without_mirror_IPCAI2021_deeps_exclusion__fold{fold}_3dcascadefullres_pred/' \ 12 | 'evaluation_oldsdf_0.25__2000_False.pkl' 13 | 14 | names_M = base_dir + '/all_data/nnUNet/nnUNet_processed/Task22_ipcai2021/splits_final.pkl' 15 | 16 | print(eval_reslult_pkl_path) 17 | with open(eval_reslult_pkl_path, 'rb') as f: 18 | eval_reslult = pkl.load(f) # dict of names and quality sub-dict 19 | with open(names_M, 'rb') as f: 20 | names2eval = pkl.load(f) 21 | 22 | print(eval_reslult.keys()) 23 | 24 | M_d = {} 25 | M_d['SIEMENS'] = 13 26 | M_d['GE'] = 14 27 | M_d['Philips'] = 15 28 | M_d['TOSHIBA'] = 16 29 | M_d['allM'] = 21 30 | 31 | Manus = ['SIEMENS', 'GE', 'Philips', 'TOSHIBA', 'allM'] 32 | for manu in Manus: 33 | names = [] 34 | bone1_Hausdorff = [] 35 | bone1_Dice = [] 36 | bone2_Hausdorff = [] 37 | bone2_Dice = [] 38 | bone3_Hausdorff = [] 39 | bone3_Dice = [] 40 | bone4_Hausdorff = [] 41 | bone4_Dice = [] 42 | whole_Hausdorff = [] 43 | whole_Dice = [] 44 | mean_Hausdorff = [] 45 | mean_Dice = [] 46 | weight_Hausdorff = [] 47 | weight_Dice = [] 48 | for na in names2eval[M_d[manu]]['test']: 49 | key = na.replace('train_', '') 50 | quality = eval_reslult[key] 51 | 52 | names.append(key) 53 | bone1_Hausdorff.append(quality[1]['Hausdorff']) 54 | bone1_Dice.append(quality[1]['dice']) 55 | bone2_Hausdorff.append(quality[2]['Hausdorff']) 56 | bone2_Dice.append(quality[2]['dice']) 57 | bone3_Hausdorff.append(quality[3]['Hausdorff']) 58 | bone3_Dice.append(quality[3]['dice']) 59 | bone4_Hausdorff.append(quality[4]['Hausdorff']) 60 | bone4_Dice.append(quality[4]['dice']) 61 | whole_Hausdorff.append(quality['whole']['Hausdorff']) 62 | whole_Dice.append(quality['whole']['dice']) 63 | mean_Hausdorff.append(quality['mean_hausdorff']) 64 | mean_Dice.append(quality['mean_dice']) 65 | weight_Hausdorff.append(quality['weighted_mean_hausdorff']) 66 | weight_Dice.append(quality['weighted_mean_dice']) 67 | print(manu, len(names)) 68 | print( 69 | 'bone1_Dice:', np.array(bone1_Dice).mean(), '\n', 70 | 'bone1_Hausdorff:', np.array(bone1_Hausdorff).mean(), '\n', 71 | 'bone2_Dice:', np.array(bone2_Dice).mean(), '\n', 72 | 'bone2_Hausdorff:', np.array(bone2_Hausdorff).mean(), '\n', 73 | 'bone3_Dice:', np.array(bone3_Dice).mean(), '\n', 74 | 'bone3_Hausdorff:', np.array(bone3_Hausdorff).mean(), '\n', 75 | 'bone4_Dice:', np.array(bone4_Dice).mean(), '\n', 76 | 'bone4_Hausdorff:', np.array(bone4_Hausdorff).mean(), '\n', 77 | 'whole_Dice:', np.array(whole_Dice).mean(), '\n', 78 | 'whole_Hausdorff:', np.array(whole_Hausdorff).mean(), '\n', 79 | 'mean_Dice:', np.array(mean_Dice).mean(), '\n', 80 | 'mean_Hausdorff:', np.array(mean_Hausdorff).mean(), '\n', 81 | 'weight_Dice:', np.array(weight_Dice).mean(), '\n', 82 | 'weight_Hausdorff:', np.array(weight_Hausdorff).mean(), '\n', 83 | ) 84 | assert len(names) == len(bone1_Dice) 85 | assert len(names) == len(bone1_Hausdorff) 86 | assert len(names) == len(bone2_Dice) 87 | assert len(names) == len(bone2_Hausdorff) 88 | assert len(names) == len(bone3_Dice) 89 | assert len(names) == len(bone3_Hausdorff) 90 | assert len(names) == len(bone4_Dice) 91 | assert len(names) == len(bone4_Hausdorff) 92 | assert len(names) == len(whole_Dice) 93 | assert len(names) == len(whole_Hausdorff) 94 | assert len(names) == len(mean_Dice) 95 | assert len(names) == len(mean_Hausdorff) 96 | assert len(names) == len(weight_Dice) 97 | assert len(names) == len(weight_Hausdorff) 98 | 99 | results = {'names': names, 100 | 'bone1_Dice': bone1_Dice, 101 | 'bone1_Hausdorff': bone1_Hausdorff, 102 | 'bone2_Dice': bone2_Dice, 103 | 'bone2_Hausdorff': bone2_Hausdorff, 104 | 'bone3_Dice': bone3_Dice, 105 | 'bone3_Hausdorff': bone3_Hausdorff, 106 | 'bone4_Dice': bone4_Dice, 107 | 'bone4_Hausdorff': bone4_Hausdorff, 108 | 'whole_Dice': whole_Dice, 109 | 'whole_Hausdorff': whole_Hausdorff, 110 | 'mean_Dice': mean_Dice, 111 | 'mean_Hausdorff': mean_Hausdorff, 112 | 'weight_Dice': weight_Dice, 113 | 'weight_Hausdorff': weight_Hausdorff 114 | } 115 | results_pd = pd.DataFrame(results) 116 | save_csv_path = eval_reslult_pkl_path.replace('.pkl', '_{}.csv'.format(manu)) 117 | results_pd.to_csv(save_csv_path) 118 | 119 | if __name__ == '__main__': 120 | 121 | func() -------------------------------------------------------------------------------- /nnunet/runs.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | sys.path.append(os.path.dirname(os.getcwd())) 4 | from nnunet.paths import my_output_identifier 5 | 6 | ######################################################################################################################## 7 | # Experiment preparing # 8 | ######################################################################################################################## 9 | home_dir = os.environ['HOME'] 10 | train_dir = os.path.join(home_dir,'all_data/nnUNet/rawdata/Task11_CTPelvic1K') 11 | output_dir = os.path.join(home_dir, 'all_data/nnUNet/nnUNet_raw/Task11_CTPelvic1K') 12 | # command = f'python dataset_conversion/JstPelvisSegmentation_5label.py --train_dir {train_dir} --output_dir {output_dir}' 13 | # command = 'python experiment_planning/plan_and_preprocess_task.py -t Task11_CTPelvic1K -pl 20 -pf 20' 14 | 15 | processed_path = os.path.join(home_dir, 'all_data/nnUNet/nnUNet_processed/Task11_CTPelvic1K') 16 | check_save_path = os.path.join(home_dir, 'all_data/nnUNet/nnUNet_processed/Task11_CTPelvic1K/Task11_check') 17 | # command = f'python preprocessing/lumbosacral_joint_sampling.py --processed_path {processed_path} --check_save_path {check_save_path}' 18 | 19 | 20 | 21 | ######################################################################################################################## 22 | # Experiment running # 23 | ######################################################################################################################## 24 | # TASK = 'Task11_CTPelvic1K' 25 | TASK = 'Task22_ipcai2021' 26 | FOLD = 0 27 | GPU = 0 28 | """ 29 | Training 30 | """ 31 | # command = f'python run/run_training.py 2d nnUNetTrainer {TASK} {FOLD} --gpu {GPU}' # TASK fold gpu_idx 32 | # command = f'python run/run_training.py 3d_fullres nnUNetTrainer {TASK} {FOLD} --gpu {GPU}' 33 | # command = f'python run/run_training.py 3d_lowres nnUNetTrainer {TASK} {FOLD} --gpu {GPU}' 34 | # command = f'python run/run_training.py 3d_cascade_fullres nnUNetTrainerCascadeFullRes {TASK} {FOLD} --gpu {GPU}' 35 | 36 | 37 | """ 38 | Validation: add " --validation_only" "--valbest" 39 | pay attention: do_mirroring has been changed to "False" 40 | """ 41 | # command = f'python run/run_training.py 2d nnUNetTrainer {TASK} {FOLD} --gpu {GPU} --validation_only --valbest' 42 | # command = f'python run/run_training.py 3d_fullres nnUNetTrainer {TASK} {FOLD} --gpu {GPU} --validation_only --valbest' 43 | # command = f'python run/run_training.py 3d_lowres nnUNetTrainer {TASK} {FOLD} --gpu {GPU} --validation_only --valbest' 44 | # command = f'python run/run_training.py 3d_cascade_fullres nnUNetTrainerCascadeFullRes {TASK} {FOLD} --gpu {GPU} --validation_only --valbest' 45 | 46 | """ 47 | Testing on data never seen 48 | """ 49 | test_data_path = os.path.join(home_dir, 'all_data/nnUNet/rawdata/ipcai2021_ALL_Test') 50 | 51 | # command = f'python inference/predict_simple.py ' \ 52 | # f'-i {test_data_path} ' \ 53 | # f'-o {test_data_path}/{TASK}__{my_output_identifier}__fold{FOLD}_2d_pred ' \ 54 | # f'-t {TASK} ' \ 55 | # f'-tr nnUNetTrainer ' \ 56 | # f'-m 2d ' \ 57 | # f'-f {FOLD} ' \ 58 | # f'--num_threads_preprocessing 12 '\ 59 | # f'--num_threads_nifti_save 6 '\ 60 | # f'--gpu {GPU}' 61 | # 62 | # command = f'python inference/predict_simple.py ' \ 63 | # f'-i {test_data_path} ' \ 64 | # f'-o {test_data_path}/{TASK}__{my_output_identifier}__fold{FOLD}_3dfullres_pred ' \ 65 | # f'-t {TASK} ' \ 66 | # f'-tr nnUNetTrainer ' \ 67 | # f'-m 3d_fullres ' \ 68 | # f'-f {FOLD} ' \ 69 | # f'--gpu {GPU}' 70 | # 71 | # 72 | # command = f'python inference/predict_simple.py ' \ 73 | # f'-i {test_data_path} ' \ 74 | # f'-o {test_data_path}/{TASK}__{my_output_identifier}__fold{FOLD}_3dlowres_pred ' \ 75 | # f'-t {TASK} ' \ 76 | # f'-tr nnUNetTrainer ' \ 77 | # f'-m 3d_lowres ' \ 78 | # f'-f {FOLD} ' \ 79 | # f'--gpu {GPU} ' \ 80 | # f'--overwrite_existing 0' 81 | # 82 | # my_task_lowres = TASK 83 | # my_output_identifier_lowres = 'CTPelvic1K' #your low_res experiment\'s "my_output_identifier" in path 84 | # command = f'python inference/predict_simple.py ' \ 85 | # f'-i {test_data_path} ' \ 86 | # f'-o {test_data_path}/{TASK}__{my_output_identifier_lowres}__{my_output_identifier}__fold{FOLD}_3dcascadefullres_pred ' \ 87 | # f'-t {TASK} ' \ 88 | # f'-tr nnUNetTrainerCascadeFullRes ' \ 89 | # f'-m 3d_cascade_fullres ' \ 90 | # f'-f {FOLD} ' \ 91 | # f'-l {test_data_path}/{my_task_lowres}__{my_output_identifier_lowres}__fold{FOLD}_3dlowres_pred ' \ 92 | # f'--gpu {GPU} ' \ 93 | # f'--overwrite_existing 0' 94 | 95 | ######################################################################################################################## 96 | # evaluation # 97 | ######################################################################################################################## 98 | # command = 'python ../evaluation.py' 99 | # command = 'python ../save_evaluation_results2csv.py' 100 | # command = 'python ../save_evaluation_results2csv_Manu.py' 101 | 102 | if __name__ == '__main__': 103 | print('\n'*2) 104 | print(command,'\n'*2) 105 | os.system(command) 106 | -------------------------------------------------------------------------------- /nnunet/evaluation/model_selection/ensemble.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | from multiprocessing.pool import Pool 4 | import shutil 5 | import numpy as np 6 | from nnunet.evaluation.evaluator import aggregate_scores 7 | from nnunet.inference.segmentation_export import save_segmentation_nifti_from_softmax 8 | from batchgenerators.utilities.file_and_folder_operations import * 9 | from nnunet.paths import network_training_output_dir, preprocessing_output_dir, default_plans_identifier 10 | import argparse 11 | 12 | 13 | def merge(args): 14 | file1, file2, properties_file, out_file = args 15 | if not isfile(out_file): 16 | res1 = np.load(file1)['softmax'] 17 | res2 = np.load(file2)['softmax'] 18 | props = load_pickle(properties_file) 19 | mn = np.mean((res1, res2), 0) 20 | save_segmentation_nifti_from_softmax(mn, out_file, props, 1, None, None, None) 21 | 22 | 23 | if __name__ == "__main__": 24 | parser = argparse.ArgumentParser(usage="This is intended to ensemble training images (from cross-validation) only. Use" 25 | "inference/ensemble_predictions.py instead") 26 | parser.add_argument("training_output_folder1") 27 | parser.add_argument("training_output_folder2") 28 | parser.add_argument("output_folder") 29 | parser.add_argument("task") # we need to know this for gt_segmentations 30 | 31 | args = parser.parse_args() 32 | 33 | training_output_folder1 = args.training_output_folder1 34 | training_output_folder2 = args.training_output_folder2 35 | output_folder = args.output_folder 36 | task = args.task 37 | 38 | # only_keep_largest_connected_component is the same for all stages 39 | dataset_directory = join(preprocessing_output_dir, task) 40 | plans = load_pickle(join(preprocessing_output_dir, task, default_plans_identifier + "_plans_2D.pkl")) # we need this only for the labels 41 | 42 | files1 = [] 43 | files2 = [] 44 | property_files = [] 45 | out_files = [] 46 | gt_segmentations = [] 47 | 48 | folder_with_gt_segs = join(dataset_directory, "gt_segmentations") 49 | folder_where_some_pkl_are = join(dataset_directory, "nnUNet_2D_stage0") # we can use this because npz are already 50 | # in the correct shape and we need the original geometry to restore the niftis 51 | 52 | folds = np.arange(5) 53 | 54 | for f in folds: 55 | validation_folder_net1 = join(training_output_folder1, "fold_%d" % f, "validation") 56 | validation_folder_net2 = join(training_output_folder2, "fold_%d" % f, "validation") 57 | patient_identifiers1 = subfiles(validation_folder_net1, False, None, 'npz', True) 58 | patient_identifiers2 = subfiles(validation_folder_net2, False, None, 'npz', True) 59 | # we don't do postprocessing anymore so there should not be any of that noPostProcess 60 | patient_identifiers1_nii = [i for i in subfiles(validation_folder_net1, False, None, suffix='nii.gz', sort=True) if not i.endswith("noPostProcess.nii.gz") and not i.endswith('_postprocessed.nii.gz')] 61 | patient_identifiers2_nii = [i for i in subfiles(validation_folder_net2, False, None, suffix='nii.gz', sort=True) if not i.endswith("noPostProcess.nii.gz") and not i.endswith('_postprocessed.nii.gz')] 62 | assert all([i[:-4] == j[:-7] for i, j in zip(patient_identifiers1, patient_identifiers1_nii)]), "npz seem to be missing. run validation with save_softmax=True" 63 | assert all([i[:-4] == j[:-7] for i, j in zip(patient_identifiers2, patient_identifiers2_nii)]), "npz seem to be missing. run validation with save_softmax=True" 64 | 65 | all_patient_identifiers = patient_identifiers1 66 | for p in patient_identifiers2: 67 | if p not in all_patient_identifiers: 68 | all_patient_identifiers.append(p) 69 | 70 | # assert these patients exist for both methods 71 | assert all([isfile(join(validation_folder_net1, i)) for i in all_patient_identifiers]) 72 | assert all([isfile(join(validation_folder_net2, i)) for i in all_patient_identifiers]) 73 | 74 | maybe_mkdir_p(output_folder) 75 | 76 | for p in all_patient_identifiers: 77 | files1.append(join(validation_folder_net1, p)) 78 | files2.append(join(validation_folder_net2, p)) 79 | property_files.append(join(folder_where_some_pkl_are, p)[:-3] + "pkl") 80 | out_files.append(join(output_folder, p[:-4] + ".nii.gz")) 81 | gt_segmentations.append(join(folder_with_gt_segs, p[:-4] + ".nii.gz")) 82 | 83 | p = Pool(8) 84 | p.map(merge, zip(files1, files2, property_files, out_files)) 85 | p.close() 86 | p.join() 87 | 88 | """for args in zip(files1, files2, property_files, out_files, [only_keep_largest_connected_component] * len(files1)): 89 | print(args[0], args[1]) 90 | merge(args)""" 91 | 92 | if not isfile(join(output_folder, "summary_allFolds.json")) and len(out_files) > 0: 93 | out_dir_all_json = join(network_training_output_dir, "summary_jsons") 94 | # now evaluate if all these gt files exist 95 | aggregate_scores(tuple(zip(out_files, gt_segmentations)), labels=plans['all_classes'], 96 | json_output_file=join(output_folder, "summary_allFolds.json"), json_task=task, 97 | json_name=task + "__" + output_folder.split("/")[-1], num_threads=4) 98 | json_out = load_json(join(output_folder, "summary_allFolds.json")) 99 | json_out["experiment_name"] = output_folder.split("/")[-1] 100 | save_json(json_out, join(output_folder, "summary_allFolds.json")) 101 | shutil.copy(join(output_folder, "summary_allFolds.json"), join(out_dir_all_json, "%s__%s.json" % (task, output_folder.split("/")[-1]))) 102 | -------------------------------------------------------------------------------- /nnunet/evaluation/model_selection/summarize_results_with_plans.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | from batchgenerators.utilities.file_and_folder_operations import * 4 | import os 5 | from nnunet.evaluation.model_selection.summarize_results_in_one_json import summarize 6 | from nnunet.paths import network_training_output_dir 7 | import numpy as np 8 | 9 | 10 | def list_to_string(l, delim=","): 11 | st = "%03.3f" % l[0] 12 | for i in l[1:]: 13 | st += delim + "%03.3f" % i 14 | return st 15 | 16 | 17 | def write_plans_to_file(f, plans_file, stage=0, do_linebreak_at_end=True, override_name=None): 18 | a = load_pickle(plans_file) 19 | stages = list(a['plans_per_stage'].keys()) 20 | stages.sort() 21 | patch_size_in_mm = [i * j for i, j in zip(a['plans_per_stage'][stages[stage]]['patch_size'], 22 | a['plans_per_stage'][stages[stage]]['current_spacing'])] 23 | median_patient_size_in_mm = [i * j for i, j in zip(a['plans_per_stage'][stages[stage]]['median_patient_size_in_voxels'], 24 | a['plans_per_stage'][stages[stage]]['current_spacing'])] 25 | if override_name is None: 26 | f.write(plans_file.split("/")[-2] + "__" + plans_file.split("/")[-1]) 27 | else: 28 | f.write(override_name) 29 | f.write(";%d" % stage) 30 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['batch_size'])) 31 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['num_pool_per_axis'])) 32 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['patch_size'])) 33 | f.write(";%s" % list_to_string(patch_size_in_mm)) 34 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['median_patient_size_in_voxels'])) 35 | f.write(";%s" % list_to_string(median_patient_size_in_mm)) 36 | f.write(";%s" % list_to_string(a['plans_per_stage'][stages[stage]]['current_spacing'])) 37 | f.write(";%s" % list_to_string(a['plans_per_stage'][stages[stage]]['original_spacing'])) 38 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['pool_op_kernel_sizes'])) 39 | f.write(";%s" % str(a['plans_per_stage'][stages[stage]]['conv_kernel_sizes'])) 40 | if do_linebreak_at_end: 41 | f.write("\n") 42 | 43 | 44 | if __name__ == "__main__": 45 | summarize((1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 24, 27), output_dir=join(network_training_output_dir, "summary_fold0"), folds=(0,)) 46 | base_dir = os.environ['RESULTS_FOLDER'] 47 | nnunets = ['nnUNetV2', 'nnUNetV2_zspacing'] 48 | task_ids = list(range(99)) 49 | with open("summary.csv", 'w') as f: 50 | f.write("identifier;stage;batch_size;num_pool_per_axis;patch_size;patch_size(mm);median_patient_size_in_voxels;median_patient_size_in_mm;current_spacing;original_spacing;pool_op_kernel_sizes;conv_kernel_sizes;patient_dc;global_dc\n") 51 | for i in task_ids: 52 | for nnunet in nnunets: 53 | try: 54 | summary_folder = join(base_dir, nnunet, "summary_fold0") 55 | if isdir(summary_folder): 56 | summary_files = subfiles(summary_folder, join=False, prefix="Task%02.0d_" % i, suffix=".json", sort=True) 57 | for s in summary_files: 58 | tmp = s.split("__") 59 | trainer = tmp[2] 60 | 61 | expected_output_folder = join(base_dir, nnunet, tmp[1], tmp[0], tmp[2].split(".")[0]) 62 | name = tmp[0] + "__" + nnunet + "__" + tmp[1] + "__" + tmp[2].split(".")[0] 63 | global_dice_json = join(base_dir, nnunet, tmp[1], tmp[0], tmp[2].split(".")[0], "fold_0", "validation_tiledTrue_doMirror_True", "global_dice.json") 64 | 65 | if not isdir(expected_output_folder) or len(tmp) > 3: 66 | if len(tmp) == 2: 67 | continue 68 | expected_output_folder = join(base_dir, nnunet, tmp[1], tmp[0], tmp[2] + "__" + tmp[3].split(".")[0]) 69 | name = tmp[0] + "__" + nnunet + "__" + tmp[1] + "__" + tmp[2] + "__" + tmp[3].split(".")[0] 70 | global_dice_json = join(base_dir, nnunet, tmp[1], tmp[0], tmp[2] + "__" + tmp[3].split(".")[0], "fold_0", "validation_tiledTrue_doMirror_True", "global_dice.json") 71 | 72 | assert isdir(expected_output_folder), "expected output dir not found" 73 | plans_file = join(expected_output_folder, "plans.pkl") 74 | assert isfile(plans_file) 75 | 76 | plans = load_pickle(plans_file) 77 | num_stages = len(plans['plans_per_stage']) 78 | if num_stages > 1 and tmp[1] == "3d_fullres": 79 | stage = 1 80 | elif (num_stages == 1 and tmp[1] == "3d_fullres") or tmp[1] == "3d_lowres": 81 | stage = 0 82 | else: 83 | print("skipping", s) 84 | continue 85 | 86 | g_dc = load_json(global_dice_json) 87 | mn_glob_dc = np.mean(list(g_dc.values())) 88 | 89 | write_plans_to_file(f, plans_file, stage, False, name) 90 | # now read and add result to end of line 91 | results = load_json(join(summary_folder, s)) 92 | mean_dc = results['results']['mean']['mean']['Dice'] 93 | f.write(";%03.3f" % mean_dc) 94 | f.write(";%03.3f\n" % mn_glob_dc) 95 | print(name, mean_dc) 96 | except Exception as e: 97 | print(e) 98 | -------------------------------------------------------------------------------- /nnunet/run/run_training.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | sys.path.append(os.path.dirname(os.getcwd())) 4 | 5 | import argparse 6 | from batchgenerators.utilities.file_and_folder_operations import * 7 | from nnunet.run.default_configuration import get_default_configuration 8 | from nnunet.paths import default_plans_identifier 9 | from nnunet.training.cascade_stuff.predict_next_stage import predict_next_stage 10 | from nnunet.training.network_training.nnUNetTrainer import nnUNetTrainer 11 | from nnunet.training.network_training.nnUNetTrainerCascadeFullRes import nnUNetTrainerCascadeFullRes 12 | 13 | if __name__ == "__main__": 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument("network") ### '2d', '3d_lowres', '3d_fullres', '3d_cascade_fullres' 16 | parser.add_argument("network_trainer") #### nnUNetTrainer 17 | parser.add_argument("task") #### Task 18 | parser.add_argument("fold", help='0, 1, ..., 5 or \'all\'') #### 19 | parser.add_argument("--ndet", required=False, default=False, action="store_true", 20 | help="Per default training is deterministic, " 21 | "nondeterministic allows cudnn.benchmark which will can give up to 20%% performance. " 22 | "Set this to do nondeterministic training") ### ....????? what this means ..... 23 | 24 | parser.add_argument("-val", "--validation_only", help="use this if you want to only run the validation", action="store_true") 25 | parser.add_argument("-c", "--continue_training", help="use this if you want to continue a training", action="store_true") 26 | parser.add_argument("-p", help="plans identifier", default=default_plans_identifier, required=False) 27 | parser.add_argument("-u", "--unpack_data", required=False, default=1, type=int, help="Leave it as 1, development only") 28 | parser.add_argument("--npz", required=False, default=False, action="store_true", help="if set then nnUNet will " 29 | "export npz files of " 30 | "predicted segmentations " 31 | "in the vlaidation as well. " ### !!!! 32 | "This is needed to run the " 33 | "ensembling step so unless " 34 | "you are developing nnUNet " 35 | "you should enable this") 36 | 37 | parser.add_argument("--find_lr", required=False, default=False, action="store_true", help="not used here, just for fun") ### ...... 38 | parser.add_argument("--valbest", required=False, default=False, action="store_true", help="hands off. This is not intended to be used") 39 | parser.add_argument("--fp16", required=False, default=False, action="store_true", help="enable fp16 training. Makes sense for 2d only! (and only on supported hardware!)") 40 | parser.add_argument("--gpu", default='0', type=str) 41 | args = parser.parse_args() 42 | 43 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu 44 | 45 | 46 | network = args.network 47 | network_trainer = args.network_trainer 48 | task = args.task 49 | fold = args.fold 50 | deterministic = not args.ndet 51 | 52 | validation_only = args.validation_only 53 | plans_identifier = args.p 54 | find_lr = args.find_lr 55 | unpack = args.unpack_data 56 | valbest = args.valbest 57 | fp16 = args.fp16 58 | 59 | if unpack == 0: 60 | unpack = False 61 | elif unpack == 1: 62 | unpack = True 63 | else: 64 | raise ValueError("Unexpected value for -u/--unpack_data: %s. Use 1 or 0." % str(unpack)) 65 | 66 | if fold == 'all': 67 | pass 68 | else: 69 | fold = int(fold) 70 | 71 | plans_file, output_folder_name, dataset_directory, batch_dice, stage, trainer_class = \ 72 | get_default_configuration(network, task, network_trainer, plans_identifier) 73 | 74 | if trainer_class is None: 75 | raise RuntimeError("Could not find trainer class in nnunet.training.network_training") 76 | 77 | if network == "3d_cascade_fullres": 78 | assert issubclass(trainer_class, nnUNetTrainerCascadeFullRes), \ 79 | "If running 3d_cascade_fullres then your trainer class must be derived from nnunetTrainerCascadeFullRes" 80 | else: 81 | assert issubclass(trainer_class, nnUNetTrainer), "network_trainer was found but is not derived from nnunetTrainer" ### ....argument 'network_trainer' ... 82 | 83 | trainer = trainer_class(plans_file, fold, 84 | output_folder=output_folder_name, 85 | dataset_directory=dataset_directory, 86 | batch_dice=batch_dice, 87 | stage=stage, 88 | unpack_data=unpack, 89 | deterministic=deterministic, 90 | fp16=fp16, 91 | network_dims=network) 92 | trainer.initialize(not validation_only) 93 | 94 | if find_lr: 95 | trainer.find_lr() 96 | else: 97 | if not validation_only: 98 | if args.continue_training: 99 | trainer.load_latest_checkpoint() 100 | trainer.run_training() 101 | 102 | ### validation 103 | elif not valbest: 104 | trainer.load_latest_checkpoint(train=False) 105 | 106 | if valbest: 107 | trainer.load_best_checkpoint(train=False) 108 | val_folder = "validation_best_epoch" 109 | else: 110 | val_folder = "validation" 111 | 112 | ### we set do_mirroring False because of asymmetry of anatomy. 113 | trainer.validate(save_softmax=args.npz, validation_folder_name=val_folder, do_mirroring=False, use_gaussian=True) 114 | 115 | if network == '3d_lowres': 116 | trainer.load_best_checkpoint(False) 117 | print("predicting segmentations for the next stage of the cascade") 118 | print() 119 | predict_next_stage(trainer, 120 | join(dataset_directory, trainer.plans['data_identifier'] + "_stage%d" % 1), 121 | fold) 122 | -------------------------------------------------------------------------------- /nnunet/training/cascade_stuff/predict_next_stage.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from batchgenerators.utilities.file_and_folder_operations import * 3 | import argparse 4 | from nnunet.preprocessing.preprocessing import resample_data_or_seg 5 | from batchgenerators.utilities.file_and_folder_operations import maybe_mkdir_p 6 | import nnunet 7 | from nnunet.run.default_configuration import get_default_configuration 8 | from multiprocessing import Pool 9 | 10 | from nnunet.training.model_restore import recursive_find_trainer 11 | from nnunet.training.network_training.nnUNetTrainer import nnUNetTrainer 12 | 13 | 14 | def resample_and_save(predicted, target_shape, output_file): 15 | predicted_new_shape = resample_data_or_seg(predicted, target_shape, False, order=1, do_separate_z=False, cval=0) 16 | seg_new_shape = predicted_new_shape.argmax(0) 17 | np.savez_compressed(output_file, data=seg_new_shape.astype(np.uint8)) 18 | 19 | 20 | def predict_next_stage(trainer, stage_to_be_predicted_folder, fold): 21 | output_folder = join(pardir(trainer.output_folder), f"pred_next_stage_{fold}") 22 | maybe_mkdir_p(output_folder) 23 | 24 | process_manager = Pool(2) 25 | results = [] 26 | 27 | print(len(trainer.dataset_val.keys())) 28 | for pat in trainer.dataset_val.keys(): 29 | 30 | if os.path.exists(output_folder+'/'+pat+"_segFromPrevStage.npz"): 31 | print(pat,'has been predicted!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!continue!!!!!!!!!') 32 | continue 33 | 34 | 35 | 36 | print(pat) 37 | data_file = trainer.dataset_val[pat]['data_file'] 38 | try: 39 | data_preprocessed = np.load(data_file)['data'][:-1] 40 | except Exception as _: 41 | data_preprocessed = np.load(data_file.replace('.npz','.npy'))[:-1] 42 | 43 | predicted = trainer.predict_preprocessed_data_return_softmax(data_preprocessed, False, # do_mirroring 44 | 1, False, 1, 45 | trainer.data_aug_params['mirror_axes'], 46 | True, True, 2, trainer.patch_size, True) 47 | data_file_nofolder = data_file.split("/")[-1] 48 | data_file_nextstage = join(stage_to_be_predicted_folder, data_file_nofolder) 49 | try: 50 | data_nextstage = np.load(data_file_nextstage)['data'] 51 | except Exception as _: 52 | data_nextstage = np.load(data_file_nextstage.replace('.npz','.npy')) 53 | 54 | target_shp = data_nextstage.shape[1:] 55 | output_file = join(output_folder, data_file_nextstage.split("/")[-1][:-4] + "_segFromPrevStage.npz") 56 | results.append(process_manager.starmap_async(resample_and_save, [(predicted, target_shp, output_file)])) 57 | 58 | _ = [i.get() for i in results] 59 | 60 | ### train 61 | print(len(trainer.dataset_tr.keys())) 62 | for pat in trainer.dataset_tr.keys(): 63 | 64 | if os.path.exists(output_folder+'/'+pat+"_segFromPrevStage.npz"): 65 | print(pat,'has been predicted!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!continue!!!!!!!!!') 66 | continue 67 | 68 | print(pat) 69 | data_file = trainer.dataset_tr[pat]['data_file'] 70 | try: 71 | data_preprocessed = np.load(data_file)['data'][:-1] 72 | except Exception as _: 73 | data_preprocessed = np.load(data_file.replace('.npz','.npy'))[:-1] 74 | 75 | predicted = trainer.predict_preprocessed_data_return_softmax(data_preprocessed, False, # do_mirroring 76 | 1, False, 1, 77 | trainer.data_aug_params['mirror_axes'], 78 | True, True, 2, trainer.patch_size, True) 79 | data_file_nofolder = data_file.split("/")[-1] 80 | data_file_nextstage = join(stage_to_be_predicted_folder, data_file_nofolder) 81 | try: 82 | data_nextstage = np.load(data_file_nextstage)['data'] 83 | except Exception as _: 84 | data_nextstage = np.load(data_file_nextstage.replace('.npz','.npy')) 85 | 86 | target_shp = data_nextstage.shape[1:] 87 | output_file = join(output_folder, data_file_nextstage.split("/")[-1][:-4] + "_segFromPrevStage.npz") 88 | results.append(process_manager.starmap_async(resample_and_save, [(predicted, target_shp, output_file)])) 89 | 90 | _ = [i.get() for i in results] 91 | 92 | 93 | if __name__ == "__main__": 94 | """ 95 | RUNNING THIS SCRIPT MANUALLY IS USUALLY NOT NECESSARY. USE THE run_training.py FILE! 96 | 97 | This script is intended for predicting all the low resolution predictions of 3d_lowres for the next stage of the 98 | cascade. It needs to run once for each fold so that the segmentation is only generated for the validation set 99 | and not on the data the network was trained on. Run it with 100 | python predict_next_stage TRAINERCLASS TASK FOLD""" 101 | 102 | parser = argparse.ArgumentParser() 103 | parser.add_argument("network_trainer") 104 | parser.add_argument("task") 105 | parser.add_argument("fold", type=int) 106 | 107 | args = parser.parse_args() 108 | 109 | trainerclass = args.network_trainer 110 | task = args.task 111 | fold = args.fold 112 | 113 | plans_file, folder_with_preprocessed_data, output_folder_name, dataset_directory, batch_dice, stage = \ 114 | get_default_configuration("3d_lowres", task) 115 | 116 | trainer_class = recursive_find_trainer([join(nnunet.__path__[0], "training", "network_training")], trainerclass, 117 | "nnunet.training.network_training") 118 | 119 | if trainer_class is None: 120 | raise RuntimeError("Could not find trainer class in nnunet.training.network_training") 121 | else: 122 | assert issubclass(trainer_class, 123 | nnUNetTrainer), "network_trainer was found but is not derived from nnunetTrainer" 124 | 125 | trainer = trainer_class(plans_file, fold, folder_with_preprocessed_data,output_folder=output_folder_name, 126 | dataset_directory=dataset_directory, batch_dice=batch_dice, stage=stage) 127 | 128 | trainer.initialize(False) 129 | trainer.load_dataset() 130 | trainer.do_split() 131 | trainer.load_best_checkpoint(train=False) 132 | 133 | stage_to_be_predicted_folder = join(dataset_directory, trainer.plans['data_identifier'] + "_stage%d" % 1) 134 | output_folder = join(pardir(trainer.output_folder), "pred_next_stage") 135 | maybe_mkdir_p(output_folder) 136 | 137 | predict_next_stage(trainer, stage_to_be_predicted_folder) 138 | 139 | -------------------------------------------------------------------------------- /nnunet/inference/segmentation_export.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | from copy import deepcopy 4 | import numpy as np 5 | import SimpleITK as sitk 6 | from nnunet.preprocessing.preprocessing import get_lowres_axis, get_do_separate_z, resample_data_or_seg 7 | from batchgenerators.utilities.file_and_folder_operations import * 8 | 9 | 10 | def save_segmentation_nifti_from_softmax(segmentation_softmax, out_fname, dct, order=1, region_class_order=None, 11 | seg_postprogess_fn=None, seg_postprocess_args=None, resampled_npz_fname=None, 12 | non_postprocessed_fname=None, force_separate_z=None): 13 | """ 14 | This is a utility for writing segmentations to nifto and npz. 15 | 16 | It requires the data to have been preprocessed by GenericPreprocessor because it depends on the property dictionary 17 | output (dct) to know the geometry of the original data. 18 | 19 | segmentation_softmax does not have to have the same size in pixels as the original data, it will be resampled to match that. 20 | This is generally useful because the spacings our networks operate on are most of the time not the native spacings of the image data. 21 | 22 | If seg_postprogess_fn is not None then seg_postprogess_fnseg_postprogess_fn(segmentation, *seg_postprocess_args) 23 | will be called before nifto export 24 | 25 | There is a problem with python process communication that prevents us from communicating obejcts 26 | larger than 2 GB between processes (basically when the length of the pickle string that will be sent is 27 | communicated by the multiprocessing.Pipe object then the placeholder (\%i I think) does not allow for long 28 | enough strings (lol). This could be fixed by changing i to l (for long) but that would require manually 29 | patching system python code.) We circumvent that problem here by saving softmax_pred to a npy file that will 30 | then be read (and finally deleted) by the Process. save_segmentation_nifti_from_softmax can take either 31 | filename or np.ndarray for segmentation_softmax and will handle this automatically 32 | :param segmentation_softmax: 33 | :param out_fname: 34 | :param dct: 35 | :param order: 36 | :param region_class_order: 37 | :param seg_postprogess_fn: 38 | :param seg_postprocess_args: 39 | :param resampled_npz_fname: 40 | :param non_postprocessed_fname: 41 | :param force_separate_z: if None then we dynamically decide how to resample along z, if True/False then always 42 | /never resample along z separately. Do not touch unless you know what you are doing 43 | :return: 44 | """ 45 | if isinstance(segmentation_softmax, str): 46 | assert isfile(segmentation_softmax), "If isinstance(segmentation_softmax, str) then " \ 47 | "isfile(segmentation_softmax) must be True" 48 | del_file = deepcopy(segmentation_softmax) 49 | segmentation_softmax = np.load(segmentation_softmax) 50 | os.remove(del_file) 51 | 52 | # first resample, then put result into bbox of cropping, then save 53 | current_shape = segmentation_softmax.shape 54 | shape_original_after_cropping = dct.get('size_after_cropping') 55 | shape_original_before_cropping = dct.get('original_size_of_raw_data') 56 | # current_spacing = dct.get('spacing_after_resampling') 57 | # original_spacing = dct.get('original_spacing') 58 | 59 | if np.any([i != j for i, j in zip(np.array(current_shape[1:]), np.array(shape_original_after_cropping))]): 60 | if force_separate_z is None: 61 | if get_do_separate_z(dct.get('original_spacing')): 62 | do_separate_z = True 63 | lowres_axis = get_lowres_axis(dct.get('original_spacing')) 64 | elif get_do_separate_z(dct.get('spacing_after_resampling')): 65 | do_separate_z = True 66 | lowres_axis = get_lowres_axis(dct.get('spacing_after_resampling')) 67 | else: 68 | do_separate_z = False 69 | lowres_axis = None 70 | else: 71 | do_separate_z = force_separate_z 72 | if do_separate_z: 73 | lowres_axis = get_lowres_axis(dct.get('original_spacing')) 74 | else: 75 | lowres_axis = None 76 | 77 | print("separate z:",do_separate_z, "lowres axis", lowres_axis) 78 | seg_old_spacing = resample_data_or_seg(segmentation_softmax, shape_original_after_cropping, is_seg=False, 79 | axis=lowres_axis, order=order, do_separate_z=do_separate_z, cval=0) 80 | #seg_old_spacing = resize_softmax_output(segmentation_softmax, shape_original_after_cropping, order=order) 81 | else: 82 | seg_old_spacing = segmentation_softmax 83 | 84 | if resampled_npz_fname is not None: 85 | np.savez_compressed(resampled_npz_fname, softmax=seg_old_spacing.astype(np.float16)) 86 | save_pickle(dct, resampled_npz_fname[:-4] + ".pkl") 87 | 88 | if region_class_order is None: 89 | seg_old_spacing = seg_old_spacing.argmax(0) 90 | else: 91 | seg_old_spacing_final = np.zeros(seg_old_spacing.shape[1:]) 92 | for i, c in enumerate(region_class_order): 93 | seg_old_spacing_final[seg_old_spacing[i] > 0.5] = c 94 | seg_old_spacing = seg_old_spacing_final 95 | 96 | bbox = dct.get('crop_bbox') 97 | 98 | if bbox is not None: 99 | seg_old_size = np.zeros(shape_original_before_cropping) 100 | for c in range(3): 101 | bbox[c][1] = np.min((bbox[c][0] + seg_old_spacing.shape[c], shape_original_before_cropping[c])) 102 | seg_old_size[bbox[0][0]:bbox[0][1], 103 | bbox[1][0]:bbox[1][1], 104 | bbox[2][0]:bbox[2][1]] = seg_old_spacing 105 | else: 106 | seg_old_size = seg_old_spacing 107 | 108 | if seg_postprogess_fn is not None: 109 | seg_old_size_postprocessed = seg_postprogess_fn(np.copy(seg_old_size), *seg_postprocess_args) 110 | else: 111 | seg_old_size_postprocessed = seg_old_size 112 | 113 | seg_resized_itk = sitk.GetImageFromArray(seg_old_size_postprocessed.astype(np.uint8)) 114 | seg_resized_itk.SetSpacing(dct['itk_spacing']) 115 | seg_resized_itk.SetOrigin(dct['itk_origin']) 116 | seg_resized_itk.SetDirection(dct['itk_direction']) 117 | sitk.WriteImage(seg_resized_itk, out_fname) 118 | 119 | if (non_postprocessed_fname is not None) and (seg_postprogess_fn is not None): 120 | seg_resized_itk = sitk.GetImageFromArray(seg_old_size.astype(np.uint8)) 121 | seg_resized_itk.SetSpacing(dct['itk_spacing']) 122 | seg_resized_itk.SetOrigin(dct['itk_origin']) 123 | seg_resized_itk.SetDirection(dct['itk_direction']) 124 | sitk.WriteImage(seg_resized_itk, non_postprocessed_fname) 125 | print('prediction saved....') 126 | 127 | -------------------------------------------------------------------------------- /postprocessing.py: -------------------------------------------------------------------------------- 1 | from skimage.measure import label, regionprops 2 | import numpy as np 3 | import SimpleITK as sitk 4 | from utils import _sitk_Image_reader, _sitk_image_writer 5 | import os 6 | from multiprocessing import Pool 7 | from functools import partial 8 | 9 | ###---###---###---###---###---###---###---###---###---###---###---###---###---###---###---###---###---###---###---###--- 10 | def gatherfiles(path, prefix=None, midfix=None, postfix=None, extname=True): 11 | files = os.listdir(path) 12 | if not prefix is None: 13 | files = [i for i in files if i.startswith(prefix)] 14 | if not midfix is None: 15 | files = [i for i in files if midfix in i] 16 | if not postfix is None: 17 | files = [i for i in files if i.endswith(postfix)] 18 | if extname: 19 | return files 20 | else: 21 | files = [os.path.splitext(i)[0] for i in files] 22 | return files 23 | 24 | def sdf_func(segImg): 25 | """ 26 | segImg is a sitk Image 27 | """ 28 | Sdf = sitk.SignedMaurerDistanceMap(segImg, squaredDistance=False) 29 | # Sdf = sitk.Sigmoid(-Sdf, 50, 0, 1, 0) # alpha, beta, max, min 30 | Sdf = sitk.Sigmoid(-Sdf, 10, 0, 1, 0) # alpha, beta, max, min 31 | seg = sitk.GetArrayFromImage(Sdf) 32 | seg[seg > 0.4999] = 1 # convert sdf back to numpy array, and clip 0.5 above to 1 (inside) 33 | return seg 34 | 35 | def raw_sdf_func(segImg): 36 | Sdf = sitk.SignedMaurerDistanceMap(segImg, insideIsPositive=False, squaredDistance=False) 37 | seg = sitk.GetArrayFromImage(Sdf) 38 | 39 | return seg 40 | 41 | def newsdf_post_processor(pred, main_region_th = 100000, sdf_th = 35, region_th = 2000): 42 | pred_test = pred.copy() 43 | mask_whole = np.zeros_like(pred_test) 44 | for anot in range(1, pred.max() + 1): 45 | print('i', anot) 46 | pred_single = np.zeros_like(pred_test) 47 | pred_single[pred_test == anot] = 1 48 | 49 | connected_label = label(pred_single, connectivity=pred_single.ndim) 50 | props = regionprops(connected_label) 51 | sorted_Props = sorted(props, key=lambda e: e.__getitem__('area'), reverse=True) 52 | 53 | mask_single = np.zeros_like(pred_test) 54 | index = None 55 | for idx_r, i in enumerate(range(len(sorted_Props))): 56 | print('ii', i) 57 | if sorted_Props[i]['area'] > main_region_th or idx_r == 0: 58 | print(sorted_Props[i]['area']) 59 | mask_single[connected_label == sorted_Props[i]['label']] = 1 60 | else: 61 | # only keep region bigger than main_region_th, here. 62 | index = i 63 | break 64 | if index == None: 65 | mask_whole[mask_single > 0] = 1 66 | continue 67 | 68 | ### second stage 69 | sdf_distance_mask_single = raw_sdf_func(sitk.GetImageFromArray(mask_single)) 70 | sdf_mask_single = np.zeros_like(pred) 71 | sdf_mask_single[sdf_distance_mask_single < sdf_th] = 1 72 | sdf_mask_single = sdf_mask_single.astype('uint16') 73 | 74 | # if True: 75 | # sdf_mask_single_2save = sdf_mask_single.copy() * 5 76 | # sdf_mask_single_2save[pred > 0] = pred[pred > 0] 77 | # path = '/home1/pbliu/all_data/nnUNet/IL_test_tmp/2017_05163189_SongSiBao_crop_mask_4label.nii.gz' 78 | # _, _, meta = _sitk_Image_reader(path) 79 | # print('name:', '_visual{}.nii.gz'.format(anot)) 80 | # _sitk_image_writer(sdf_mask_single_2save, meta, path.replace('.nii.gz', '_visual{}.nii.gz'.format(anot))) 81 | # else: 82 | # pass 83 | 84 | for i in range(index, len(sorted_Props)): 85 | print('iii', i) 86 | if sorted_Props[i]['area'] < region_th: 87 | break 88 | else: 89 | part = np.zeros_like(pred_test) 90 | part[connected_label == sorted_Props[i]['label']] = 1 91 | if (part * sdf_mask_single).sum() > 0: 92 | mask_single[connected_label == sorted_Props[i]['label']] = 1 93 | 94 | mask_whole[mask_single > 0] = 1 95 | result = mask_whole * pred 96 | return result 97 | 98 | 99 | def maximum_connected_region_post_processor(pred, region_th = 100000): 100 | """ 101 | pred: multi-label 102 | return: multi-label 103 | """ 104 | pred_test = pred.copy() 105 | mask_whole = np.zeros_like(pred) 106 | 107 | for i in range(1, pred.max()+1): 108 | pred_single = np.zeros_like(pred_test) 109 | pred_single[pred_test==i] = 1 110 | 111 | connected_label = label(pred_single, connectivity=pred_single.ndim) 112 | props = regionprops(connected_label) 113 | sorted_Props = sorted(props, key=lambda e: e.__getitem__('area'), reverse=True) 114 | 115 | mask_single = np.zeros_like(pred) 116 | for i in range(len(sorted_Props)): 117 | # print(sorted_Props[i]['area'],sorted_Props[i]['label']) 118 | if sorted_Props[i]['area']>region_th or i==0: # i==0, make sure the biggest area can be keeped. 119 | mask_single[connected_label==sorted_Props[i]['label']] = 1 120 | else: 121 | pass 122 | 123 | mask_whole[mask_single>0]=1 124 | 125 | return mask_whole*pred 126 | 127 | if __name__ == '__main__': 128 | """ 129 | SDF post processor 130 | 131 | """ 132 | def func(name, path, savepath, post='sdf'): 133 | _, image, meta = _sitk_Image_reader(path+'/'+name) 134 | print(name, image.shape) 135 | if post=='sdf': 136 | post_image = newsdf_post_processor(pred=image, region_th=2000, sdf_th=35) 137 | elif post=='mcr': 138 | post_image = maximum_connected_region_post_processor(image, region_th=100000) 139 | else: 140 | raise NotImplementedError 141 | 142 | _sitk_image_writer(post_image, meta, savepath+'/'+name) 143 | 144 | base_dir = os.environ['HOME'] 145 | pred_path = base_dir + "/all_data/nnUNet/rawdata/ipcai2021_ALL_Test/SDF_show" \ 146 | "/Task22_ipcai2021_T__nnUNet_without_mirror_IPCAI2021_deeps_exclusion__nnUNet_without_mirror_IPCAI2021_deeps_exclusion__fold0_3dcascadefullres_pred" 147 | save_path = pred_path+'___newSDFpost' 148 | if not os.path.exists(save_path): 149 | os.makedirs(save_path) 150 | print(pred_path) 151 | 152 | files = gatherfiles(pred_path, postfix='.nii.gz') 153 | 154 | print(len(files)) 155 | files = list(set(files)) 156 | print(len(files)) 157 | files = sorted(files) 158 | 159 | pool = Pool(16) 160 | fu = partial(func, path=pred_path, savepath=save_path) 161 | _=pool.map(fu, files) 162 | pool.close() 163 | pool.join() 164 | 165 | -------------------------------------------------------------------------------- /nnunet/training/loss_functions/ND_Crossentropy.py: -------------------------------------------------------------------------------- 1 | import torch.nn 2 | import SimpleITK as sitk 3 | import numpy as np 4 | import torch.nn.functional as F 5 | from .LovaszSoftmax import lovasz_softmax 6 | 7 | class CrossentropyND(torch.nn.CrossEntropyLoss): 8 | """ 9 | Network has to have NO NONLINEARITY! 10 | """ 11 | def __init__(self): 12 | self.reduction = 'none' 13 | super(CrossentropyND, self).__init__(reduction=self.reduction) 14 | 15 | def forward(self, inp, target, heat_map=None): 16 | ### lovaszSoftmax loss ### 17 | # inp = F.softmax(inp, dim=1) 18 | # return lovasz_softmax(inp, target, ignore=255) 19 | 20 | ## raw implementation ### 21 | 22 | target = target.long() 23 | num_classes = inp.size()[1] 24 | 25 | i0 = 1 26 | i1 = 2 27 | 28 | while i1 < len(inp.shape): # this is ugly but torch only allows to transpose two axes at once 29 | inp = inp.transpose(i0, i1) 30 | i0 += 1 31 | i1 += 1 32 | 33 | inp = inp.contiguous() 34 | inp = inp.view(-1, num_classes) 35 | 36 | target = target.view(-1,) 37 | if heat_map is None: 38 | return (super(CrossentropyND, self).forward(inp, target)).mean() 39 | else: 40 | heat_map = heat_map.view(-1,) 41 | 42 | return (super(CrossentropyND, self).forward(inp, target)*heat_map).mean() 43 | 44 | ### sdf heatmap ### 45 | loss_all = [] 46 | for i in range(num_classes): 47 | target_index = torch.zeros_like(target).float() 48 | target_index[target==i] = 1 49 | inp_index = inp[:,i] 50 | loss_index = F.binary_cross_entropy_with_logits(inp_index, target_index, reduction=self.reduction).unsqueeze(1) 51 | loss_all.append(loss_index) 52 | 53 | loss_tmp = torch.cat(loss_all, 1) 54 | ### heatmap ### 55 | i0 = 1 56 | i1 = 2 57 | while i1 < len(heatmap.shape): # this is ugly but torch only allows to transpose two axes at once 58 | heatmap = heatmap.transpose(i0, i1) 59 | i0 += 1 60 | i1 += 1 61 | 62 | heatmap = heatmap.contiguous() 63 | heatmap = heatmap.view(-1, num_classes) 64 | 65 | assert heatmap.shape == loss_tmp.shape, print(heatmap.shape, loss_tmp.shape) 66 | 67 | return (loss_tmp*heatmap.cuda()).mean() 68 | 69 | 70 | class CrossentropyND_DeepS(torch.nn.CrossEntropyLoss): 71 | """ 72 | Network has to have NO NONLINEARITY! 73 | """ 74 | def __init__(self, weights): 75 | self.DSweights = weights 76 | self.reduction = 'none' 77 | super(CrossentropyND_DeepS, self).__init__(reduction=self.reduction) 78 | 79 | def forward(self, inps, targets, heat_map=None): 80 | inps = inps[:len(targets)] # resolution sorted from high to low 81 | # print('inps shape: ',[inp.shape for inp in inps]) 82 | # print('tars shape: ',[tar.shape for tar in targets]) 83 | targets = [target.long() for target in targets] 84 | num_classes = inps[0].size()[1] 85 | 86 | i0 = 1 87 | i1 = 2 88 | while i1 < len(inps[0].shape): # this is ugly but torch only allows to transpose two axes at once 89 | inps = [inp.transpose(i0, i1) for inp in inps] 90 | i0 += 1 91 | i1 += 1 92 | 93 | inps = [inp.contiguous() for inp in inps] 94 | inps = [inp.view(-1, num_classes) for inp in inps] 95 | 96 | targets = [target.view(-1,) for target in targets] 97 | 98 | losses = [super(CrossentropyND_DeepS, self).forward(inp, target) for inp, target in zip(inps, targets)] 99 | if not heat_map is None: 100 | heat_map = heat_map.view(-1,) 101 | losses[0] = losses[0]*heat_map 102 | 103 | losses = [loss.mean()*weight for loss, weight in zip(losses, self.DSweights)] 104 | return sum(losses) 105 | 106 | # ------------------------------ sdf ----------------------------------------- 107 | def sdf_func(segImg): 108 | """ 109 | segImg is a sitk Image 110 | """ 111 | Sdf = sitk.SignedMaurerDistanceMap(segImg, squaredDistance=False) 112 | Sdf = sitk.Sigmoid(-Sdf, 50, 0, 1, 0) # alpha, beta, max, min 113 | seg = sitk.GetArrayFromImage(Sdf) 114 | seg[seg > 0.4999] = 1 # convert sdf back to numpy array, and clip 0.5 above to 1 (inside) 115 | heat_map = seg + 0.5 # putGaussianMapF(mask, sigma=50.0) 116 | return heat_map 117 | 118 | def sdf_func_convert(segImg): 119 | """ 120 | segImg is a sitk Image 121 | """ 122 | Sdf = sitk.SignedMaurerDistanceMap(segImg, squaredDistance=False) 123 | Sdf = sitk.Sigmoid(-Sdf, 50, 0, 1, 0) # alpha, beta, max, min 124 | seg = sitk.GetArrayFromImage(Sdf) 125 | seg[seg > 0.4999] = 0.7 # convert sdf back to numpy array, and clip 0.5 above to 1 (inside) 126 | seg = 0.7-seg 127 | heat_map = seg + 1 # putGaussianMapF(mask, sigma=50.0) 128 | return heat_map 129 | 130 | def get_converted_sdf_from_target(target): 131 | """ 132 | 133 | :param target: (batch, c, h, w):(uint16) 134 | :return: (batch, c, h, w):(float32) 135 | """ 136 | assert target.shape[1] == 1 137 | batch, _, h, w = target.shape 138 | heatmap = np.zeros((batch, 5, h, w), dtype='float32') ## float32 is important 139 | for i in range(5): 140 | sep_target = np.zeros_like(target) # b,c,h,w 141 | sep_target[target==i] = 1 # with one kind of target 142 | for slice in range(batch): 143 | slice_sep_target = sep_target[slice,:,:,:] # c,h,w 144 | Slice_Sep_Target = sitk.GetImageFromArray(slice_sep_target) 145 | heatmap[slice, i, :, :] = sdf_func_convert(Slice_Sep_Target) 146 | assert heatmap.max()>1.1 147 | 148 | return heatmap 149 | def get_sdf_from_target(target): 150 | """ 151 | 152 | :param target: (batch, c, h, w):(uint16) 153 | :return: (batch, c, h, w):(float32) 154 | """ 155 | assert target.shape[1] == 1 156 | batch, _, h, w = target.shape 157 | heatmap = np.zeros((batch, 5, h, w), dtype='float32') ## float32 is important 158 | for i in range(5): 159 | sep_target = np.zeros_like(target) # b,c,h,w 160 | sep_target[target==i] = 1 # with one kind of target 161 | for slice in range(batch): 162 | slice_sep_target = sep_target[slice,:,:,:] # c,h,w 163 | Slice_Sep_Target = sitk.GetImageFromArray(slice_sep_target) 164 | heatmap[slice, i, :, :] = sdf_func(Slice_Sep_Target) 165 | assert heatmap.max()>1.1 166 | 167 | return heatmap 168 | 169 | if __name__ == '__main__': 170 | path = '/Volumes/Pengbo_Kani/ICTDATA/kits19/nnunet_results/Spa_2_OK/new_spc_2/get_rid_of_pixel/new_case_00202.nii.gz' 171 | im = sitk.GetArrayFromImage(sitk.ReadImage(path)) 172 | im = im[:,None] 173 | heat_map = get_converted_sdf_from_target(im) 174 | Im = sitk.GetImageFromArray(heat_map) 175 | sitk.WriteImage(Im, '/Users/pengbo/Desktop/aa.nii.gz') -------------------------------------------------------------------------------- /nnunet/training/model_restore.py: -------------------------------------------------------------------------------- 1 | import nnunet 2 | import torch 3 | from batchgenerators.utilities.file_and_folder_operations import * 4 | import importlib 5 | import pkgutil 6 | from nnunet.training.network_training.nnUNetTrainer import nnUNetTrainer 7 | 8 | 9 | def recursive_find_trainer(folder, trainer_name, current_module): 10 | tr = None 11 | for importer, modname, ispkg in pkgutil.iter_modules(folder): 12 | # print(modname, ispkg) 13 | if not ispkg: 14 | m = importlib.import_module(current_module + "." + modname) 15 | if hasattr(m, trainer_name): 16 | tr = getattr(m, trainer_name) 17 | break 18 | 19 | if tr is None: 20 | for importer, modname, ispkg in pkgutil.iter_modules(folder): 21 | if ispkg: 22 | next_current_module = current_module + "." + modname 23 | tr = recursive_find_trainer([join(folder[0], modname)], trainer_name, current_module=next_current_module) 24 | if tr is not None: 25 | break 26 | 27 | return tr 28 | 29 | 30 | def restore_model(pkl_file, checkpoint=None, train=False): 31 | """ 32 | This is a utility function to load any nnUNet trainer from a pkl. It will recursively search 33 | nnunet.trainig.network_training for the file that contains the trainer and instantiate it with the arguments saved in the pkl file. If checkpoint 34 | is specified, it will furthermore load the checkpoint file in train/test mode (as specified by train). 35 | The pkl file required here is the one that will be saved automatically when calling nnUNetTrainer.save_checkpoint. 36 | :param pkl_file: 37 | :param checkpoint: 38 | :param train: 39 | :return: 40 | """ 41 | info = load_pickle(pkl_file) 42 | init = info['init'] 43 | name = info['name'] 44 | search_in = join(nnunet.__path__[0], "training", "network_training") 45 | tr = recursive_find_trainer([search_in], name, current_module="nnunet.training.network_training") 46 | 47 | if tr is None: 48 | """ 49 | Fabian only. This will trigger searching for trainer classes in other repositories as well 50 | """ 51 | try: 52 | import meddec 53 | search_in = join(meddec.__path__[0], "model_training") 54 | tr = recursive_find_trainer([search_in], name, current_module="meddec.model_training") 55 | except ImportError: 56 | pass 57 | 58 | if tr is None: 59 | raise RuntimeError("Could not find the model trainer specified in checkpoint in nnunet.trainig.network_training. If it " 60 | "is not located there, please move it or change the code of restore_model. Your model " 61 | "trainer can be located in any directory within nnunet.trainig.network_training (search is recursive)." 62 | "\nDebug info: \ncheckpoint file: %s\nName of trainer: %s " % (checkpoint, name)) 63 | assert issubclass(tr, nnUNetTrainer), "The network trainer was found but is not a subclass of nnUNetTrainer. " \ 64 | "Please make it so!" 65 | 66 | print(len(init),'!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!') 67 | 68 | if init[0].endswith('2D.pkl'): 69 | initnew = list(init)+['2d'] 70 | del init 71 | init = tuple(initnew) 72 | 73 | print(len(init)) 74 | print(*init) 75 | print(type(init)) 76 | if len(init) == 7: 77 | print("warning: this model seems to have been saved with a previous version of nnUNet. Attempting to load it " 78 | "anyways. Expect the unexpected.") 79 | print("manually editing init args...") 80 | init = [init[i] for i in range(len(init)) if i != 2] 81 | 82 | # init[0] is the plans file. This argument needs to be replaced because the original plans file may not exist 83 | # anymore. 84 | trainer = tr(*init) 85 | trainer.process_plans(info['plans']) 86 | if checkpoint is not None: 87 | trainer.load_checkpoint(checkpoint, train) 88 | return trainer 89 | 90 | 91 | def load_best_model_for_inference(folder): 92 | checkpoint = join(folder, "model_best.model") 93 | pkl_file = checkpoint + ".pkl" 94 | return restore_model(pkl_file, checkpoint, False) 95 | 96 | 97 | def load_model_and_checkpoint_files(folder, folds=None): 98 | """ 99 | used for if you need to ensemble the five models of a cross-validation. This will restore the model from the 100 | checkpoint in fold 0, load all parameters of the five folds in ram and return both. This will allow for fast 101 | switching between parameters (as opposed to loading them form disk each time). 102 | 103 | This is best used for inference and test prediction 104 | :param folder: 105 | :return: 106 | """ 107 | if isinstance(folds, str): 108 | folds = [join(folder, "all")] 109 | assert isdir(folds[0]), "no output folder for fold %s found" % folds 110 | elif isinstance(folds, (list, tuple)): 111 | if len(folds) == 1 and folds[0] == "all": 112 | folds = [join(folder, "all")] 113 | else: 114 | folds = [join(folder, "fold_%d" % i) for i in folds] 115 | assert all([isdir(i) for i in folds]), "list of folds specified but not all output folders are present" 116 | elif isinstance(folds, int): 117 | print('folds: ', folds) 118 | folds = [join(folder, "fold_%d" % folds)] 119 | assert all([isdir(i) for i in folds]), "output folder missing for fold %d" % folds 120 | elif folds is None: 121 | print("folds is None so we will automatically look for output folders (not using \'all\'!)") 122 | folds = subfolders(folder, prefix="fold") 123 | print("found the following folds: ", folds) 124 | else: 125 | raise ValueError("Unknown value for folds. Type: %s. Expected: list of int, int, str or None", str(type(folds))) 126 | 127 | print('we will really load fold[0]: ',folds[0]) 128 | trainer = restore_model(join(folds[0], "model_best.model.pkl")) 129 | trainer.output_folder = folder 130 | trainer.output_folder_base = folder 131 | trainer.update_fold(0) 132 | # trainer.update_fold(None) 133 | # raise NotImplementedError("I think you should check here... liupengbo-20200621") 134 | print('prediction trainer.output_folder: ',trainer.output_folder) 135 | trainer.initialize(False) 136 | # raise NotImplementedError("I think you should check here... liupengbo-20200621") 137 | all_best_model_files = [join(i, "model_best.model") for i in folds] 138 | print("!!using the following model files: ", all_best_model_files) 139 | all_params = [torch.load(i, map_location=torch.device('cuda', torch.cuda.current_device())) for i in all_best_model_files] 140 | return trainer, all_params 141 | 142 | 143 | if __name__ == "__main__": 144 | pkl = "/home/fabian/PhD/results/nnUNetV2/nnUNetV2_3D_fullres/Task04_Hippocampus/fold0/model_best.model.pkl" 145 | checkpoint = pkl[:-4] 146 | train = False 147 | trainer = restore_model(pkl, checkpoint, train) 148 | -------------------------------------------------------------------------------- /nnunet/experiment_planning/common_utils.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | from copy import deepcopy 5 | from nnunet.network_architecture.generic_UNet import Generic_UNet 6 | from nnunet.experiment_planning.configuration import FEATUREMAP_MIN_EDGE_LENGTH_BOTTLENECK 7 | import SimpleITK as sitk 8 | import shutil 9 | from batchgenerators.utilities.file_and_folder_operations import join 10 | 11 | 12 | def split_4d_nifti(filename, output_folder): 13 | img_itk = sitk.ReadImage(filename) 14 | dim = img_itk.GetDimension() 15 | file_base = filename.split("/")[-1] 16 | if dim == 3: 17 | shutil.copy(filename, join(output_folder, file_base[:-7] + "_0000.nii.gz")) 18 | return 19 | elif dim != 4: 20 | raise RuntimeError("Unexpected dimensionality: %d of file %s, cannot split" % (dim, filename)) 21 | else: 22 | img_npy = sitk.GetArrayFromImage(img_itk) 23 | spacing = img_itk.GetSpacing() 24 | origin = img_itk.GetOrigin() 25 | direction = np.array(img_itk.GetDirection()).reshape(4,4) 26 | # now modify these to remove the fourth dimension 27 | spacing = tuple(list(spacing[:-1])) 28 | origin = tuple(list(origin[:-1])) 29 | direction = tuple(direction[:-1, :-1].reshape(-1)) 30 | for i, t in enumerate(range(img_npy.shape[0])): 31 | img = img_npy[t] 32 | img_itk_new = sitk.GetImageFromArray(img) 33 | img_itk_new.SetSpacing(spacing) 34 | img_itk_new.SetOrigin(origin) 35 | img_itk_new.SetDirection(direction) 36 | sitk.WriteImage(img_itk_new, join(output_folder, file_base[:-7] + "_%04.0d.nii.gz" % i)) 37 | 38 | 39 | def get_pool_and_conv_props_poolLateV2(patch_size, min_feature_map_size, max_numpool, spacing): 40 | """ 41 | 42 | :param spacing: 43 | :param patch_size: 44 | :param min_feature_map_size: min edge length of feature maps in bottleneck 45 | :return: 46 | """ 47 | initial_spacing = deepcopy(spacing) 48 | reach = max(initial_spacing) 49 | dim = len(patch_size) 50 | 51 | num_pool_per_axis = get_network_numpool(patch_size, max_numpool, min_feature_map_size) # properity of network 52 | 53 | net_num_pool_op_kernel_sizes = [] 54 | net_conv_kernel_sizes = [] 55 | net_numpool = max(num_pool_per_axis) 56 | 57 | current_spacing = spacing 58 | for p in range(net_numpool): 59 | reached = [current_spacing[i] / reach > 0.5 for i in range(dim)] 60 | pool = [2 if num_pool_per_axis[i] + p >= net_numpool else 1 for i in range(dim)] 61 | if all(reached): 62 | conv = [3] * dim 63 | else: 64 | conv = [3 if not reached[i] else 1 for i in range(dim)] 65 | net_num_pool_op_kernel_sizes.append(pool) 66 | net_conv_kernel_sizes.append(conv) 67 | current_spacing = [i * j for i, j in zip(current_spacing, pool)] 68 | 69 | net_conv_kernel_sizes.append([3] * dim) 70 | 71 | must_be_divisible_by = get_shape_must_be_divisible_by(num_pool_per_axis) 72 | patch_size = pad_shape(patch_size, must_be_divisible_by) 73 | 74 | # we need to add one more conv_kernel_size for the bottleneck. We always use 3x3(x3) conv here 75 | return num_pool_per_axis, net_num_pool_op_kernel_sizes, net_conv_kernel_sizes, patch_size, must_be_divisible_by 76 | 77 | 78 | def get_pool_and_conv_props(spacing, patch_size, min_feature_map_size, max_numpool): 79 | """ 80 | 81 | :param spacing: 82 | :param patch_size: 83 | :param min_feature_map_size: min edge length of feature maps in bottleneck 84 | :return: 85 | """ 86 | dim = len(spacing) 87 | 88 | current_spacing = deepcopy(list(spacing)) 89 | current_size = deepcopy(list(patch_size)) 90 | 91 | pool_op_kernel_sizes = [] 92 | conv_kernel_sizes = [] 93 | 94 | num_pool_per_axis = [0] * dim 95 | 96 | while True: 97 | # find axes that are within factor 2 of min axis spacing 98 | min_spacing = min(current_spacing) 99 | valid_axes_for_pool = [i for i in range(dim) if current_spacing[i] / min_spacing < 2] 100 | axes = [] 101 | for a in range(dim): 102 | my_spacing = current_spacing[a] 103 | partners = [i for i in range(dim) if current_spacing[i] / my_spacing < 2 and my_spacing / current_spacing[i] < 2] 104 | if len(partners) > len(axes): 105 | axes = partners 106 | conv_kernel_size = [3 if i in axes else 1 for i in range(dim)] 107 | 108 | # exclude axes that we cannot pool further because of min_feature_map_size constraint 109 | #before = len(valid_axes_for_pool) 110 | valid_axes_for_pool = [i for i in valid_axes_for_pool if current_size[i] >= 2*min_feature_map_size] 111 | #after = len(valid_axes_for_pool) 112 | #if after == 1 and before > 1: 113 | # break 114 | 115 | valid_axes_for_pool = [i for i in valid_axes_for_pool if num_pool_per_axis[i] < max_numpool] 116 | 117 | if len(valid_axes_for_pool) == 0: 118 | break 119 | 120 | #print(current_spacing, current_size) 121 | 122 | other_axes = [i for i in range(dim) if i not in valid_axes_for_pool] 123 | 124 | pool_kernel_sizes = [0] * dim 125 | for v in valid_axes_for_pool: 126 | pool_kernel_sizes[v] = 2 127 | num_pool_per_axis[v] += 1 128 | current_spacing[v] *= 2 129 | current_size[v] /= 2 130 | for nv in other_axes: 131 | pool_kernel_sizes[nv] = 1 132 | 133 | pool_op_kernel_sizes.append(pool_kernel_sizes) 134 | conv_kernel_sizes.append(conv_kernel_size) 135 | #print(conv_kernel_sizes) 136 | 137 | must_be_divisible_by = get_shape_must_be_divisible_by(num_pool_per_axis) 138 | patch_size = pad_shape(patch_size, must_be_divisible_by) 139 | 140 | # we need to add one more conv_kernel_size for the bottleneck. We always use 3x3(x3) conv here 141 | conv_kernel_sizes.append([3]*dim) 142 | return num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, patch_size, must_be_divisible_by 143 | 144 | 145 | def get_shape_must_be_divisible_by(net_numpool_per_axis): 146 | return 2 ** np.array(net_numpool_per_axis) 147 | 148 | 149 | def pad_shape(shape, must_be_divisible_by): 150 | """ 151 | pads shape so that it is divisibly by must_be_divisible_by 152 | :param shape: 153 | :param must_be_divisible_by: 154 | :return: 155 | """ 156 | if not isinstance(must_be_divisible_by, (tuple, list, np.ndarray)): 157 | must_be_divisible_by = [must_be_divisible_by] * len(shape) 158 | else: 159 | assert len(must_be_divisible_by) == len(shape) 160 | 161 | new_shp = [shape[i] + must_be_divisible_by[i] - shape[i] % must_be_divisible_by[i] for i in range(len(shape))] 162 | 163 | for i in range(len(shape)): 164 | if shape[i] % must_be_divisible_by[i] == 0: 165 | new_shp[i] -= must_be_divisible_by[i] 166 | new_shp = np.array(new_shp).astype(int) 167 | return new_shp 168 | 169 | 170 | def get_network_numpool(patch_size, maxpool_cap=Generic_UNet.MAX_NUMPOOL_3D, 171 | min_feature_map_size=FEATUREMAP_MIN_EDGE_LENGTH_BOTTLENECK): 172 | network_numpool_per_axis = np.floor([np.log(i / min_feature_map_size) / np.log(2) for i in patch_size]).astype(int) 173 | network_numpool_per_axis = [min(i, maxpool_cap) for i in network_numpool_per_axis] 174 | return network_numpool_per_axis 175 | -------------------------------------------------------------------------------- /nnunet/training/data_augmentation/pyramid_augmentations.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | from skimage.morphology import label, ball 3 | from skimage.morphology.binary import binary_erosion, binary_dilation, binary_closing, binary_opening 4 | import numpy as np 5 | from batchgenerators.transforms import AbstractTransform 6 | 7 | 8 | class RemoveRandomConnectedComponentFromOneHotEncodingTransform(AbstractTransform): 9 | def __init__(self, channel_idx, key="data", p_per_sample=0.2, fill_with_other_class_p=0.25, 10 | dont_do_if_covers_more_than_X_percent=0.25): 11 | """ 12 | :param dont_do_if_covers_more_than_X_percent: dont_do_if_covers_more_than_X_percent=0.25 is 25\%! 13 | :param channel_idx: can be list or int 14 | :param key: 15 | """ 16 | self.dont_do_if_covers_more_than_X_percent = dont_do_if_covers_more_than_X_percent 17 | self.fill_with_other_class_p = fill_with_other_class_p 18 | self.p_per_sample = p_per_sample 19 | self.key = key 20 | if not isinstance(channel_idx, (list, tuple)): 21 | channel_idx = [channel_idx] 22 | self.channel_idx = channel_idx 23 | 24 | def __call__(self, **data_dict): 25 | data = data_dict.get(self.key) 26 | for b in range(data.shape[0]): 27 | if np.random.uniform() < self.p_per_sample: 28 | for c in self.channel_idx: 29 | workon = np.copy(data[b, c]) 30 | num_voxels = np.prod(workon.shape) 31 | lab, num_comp = label(workon, return_num=True) 32 | if num_comp > 0: 33 | component_ids = [] 34 | component_sizes = [] 35 | for i in range(1, num_comp + 1): 36 | component_ids.append(i) 37 | component_sizes.append(np.sum(lab == i)) 38 | component_ids = [i for i, j in zip(component_ids, component_sizes) if j < num_voxels*self.dont_do_if_covers_more_than_X_percent] 39 | #_ = component_ids.pop(np.argmax(component_sizes)) 40 | #else: 41 | # component_ids = list(range(1, num_comp + 1)) 42 | if len(component_ids) > 0: 43 | random_component = np.random.choice(component_ids) 44 | data[b, c][lab == random_component] = 0 45 | if np.random.uniform() < self.fill_with_other_class_p: 46 | other_ch = [i for i in self.channel_idx if i != c] 47 | if len(other_ch) > 0: 48 | other_class = np.random.choice(other_ch) 49 | data[b, other_class][lab == random_component] = 1 50 | data_dict[self.key] = data 51 | return data_dict 52 | 53 | 54 | class MoveSegAsOneHotToData(AbstractTransform): 55 | def __init__(self, channel_id, all_seg_labels, key_origin="seg", key_target="data", remove_from_origin=True): 56 | self.remove_from_origin = remove_from_origin 57 | self.all_seg_labels = all_seg_labels 58 | self.key_target = key_target 59 | self.key_origin = key_origin 60 | self.channel_id = channel_id 61 | 62 | def __call__(self, **data_dict): 63 | origin = data_dict.get(self.key_origin) 64 | target = data_dict.get(self.key_target) 65 | seg = origin[:, self.channel_id:self.channel_id+1] 66 | seg_onehot = np.zeros((seg.shape[0], len(self.all_seg_labels), *seg.shape[2:]), dtype=seg.dtype) 67 | for i, l in enumerate(self.all_seg_labels): 68 | seg_onehot[:, i][seg[:, 0] == l] = 1 69 | target = np.concatenate((target, seg_onehot), 1) 70 | data_dict[self.key_target] = target 71 | 72 | if self.remove_from_origin: 73 | remaining_channels = [i for i in range(origin.shape[1]) if i != self.channel_id] 74 | origin = origin[:, remaining_channels] 75 | data_dict[self.key_origin] = origin 76 | 77 | return data_dict 78 | 79 | 80 | class ApplyRandomBinaryOperatorTransform(AbstractTransform): 81 | def __init__(self, channel_idx, p_per_sample=0.3, any_of_these=(binary_dilation, binary_erosion, binary_closing, binary_opening), 82 | key="data", strel_size=(1, 10)): 83 | """ 84 | 85 | :param channel_idx: can be list or int 86 | :param p_per_sample: 87 | :param any_of_these: 88 | :param fill_diff_with_other_class: 89 | :param key: 90 | :param strel_size: 91 | """ 92 | self.strel_size = strel_size 93 | self.key = key 94 | self.any_of_these = any_of_these 95 | self.p_per_sample = p_per_sample 96 | 97 | assert not isinstance(channel_idx, tuple), "bäh" 98 | 99 | if not isinstance(channel_idx, list): 100 | channel_idx = [channel_idx] 101 | self.channel_idx = channel_idx 102 | 103 | def __call__(self, **data_dict): 104 | data = data_dict.get(self.key) 105 | for b in range(data.shape[0]): 106 | if np.random.uniform() < self.p_per_sample: 107 | ch = deepcopy(self.channel_idx) 108 | np.random.shuffle(ch) 109 | for c in ch: 110 | operation = np.random.choice(self.any_of_these) 111 | selem = ball(np.random.uniform(*self.strel_size)) 112 | workon = np.copy(data[b, c]).astype(int) 113 | res = operation(workon, selem).astype(workon.dtype) 114 | data[b, c] = res 115 | 116 | # if class was added, we need to remove it in ALL other channels to keep one hot encoding 117 | # properties 118 | # we modify data 119 | other_ch = [i for i in ch if i != c] 120 | if len(other_ch) > 0: 121 | was_added_mask = (res - workon) > 0 122 | for oc in other_ch: 123 | data[b, oc][was_added_mask] = 0 124 | # if class was removed, leave it at backgound 125 | data_dict[self.key] = data 126 | return data_dict 127 | 128 | class MoveLastFewDataToSeg_pbl(AbstractTransform): 129 | """ 130 | used when there introduce other channel data. like contour, sdf... 131 | """ 132 | def __init__(self, channel_ids, key_origin="data", key_target="seg", remove_from_origin=True): 133 | self.remove_from_origin = remove_from_origin 134 | self.key_target = key_target 135 | self.key_origin = key_origin 136 | self.channel_ids = channel_ids 137 | 138 | assert isinstance(self.channel_ids, list), self.channel_ids 139 | def __call__(self, **data_dict): 140 | origin = data_dict.get(self.key_origin) 141 | target = data_dict.get(self.key_target) 142 | 143 | if np.all(np.array(self.channel_ids)<0): 144 | self.channel_ids = [i+origin.shape[1] for i in self.channel_ids] 145 | assert not (np.any(np.array(self.channel_ids)>=0) and np.any(np.array(self.channel_ids)<0)) 146 | 147 | data_o = origin[:, self.channel_ids] 148 | target = np.concatenate((target, data_o), 1) 149 | data_dict[self.key_target] = target 150 | 151 | if self.remove_from_origin: 152 | remaining_channels = [i for i in range(origin.shape[1]) if not i in self.channel_ids] 153 | origin = origin[:, remaining_channels] 154 | data_dict[self.key_origin] = origin 155 | 156 | return data_dict 157 | 158 | 159 | 160 | -------------------------------------------------------------------------------- /nnunet/preprocessing/cropping.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import SimpleITK as sitk 4 | import numpy as np 5 | import shutil 6 | from batchgenerators.utilities.file_and_folder_operations import * 7 | from multiprocessing import Pool 8 | from collections import OrderedDict 9 | 10 | 11 | def create_nonzero_mask(data): 12 | from scipy.ndimage import binary_fill_holes 13 | assert len(data.shape) == 4 or len(data.shape) == 3, "data must have shape (C, X, Y, Z) or shape (C, X, Y)" # ......??? 3 ???? 14 | nonzero_mask = np.zeros(data.shape[1:], dtype=bool) 15 | for c in range(data.shape[0]): 16 | this_mask = data[c] != 0 17 | nonzero_mask = nonzero_mask | this_mask # | at every voxel. Array should be bool dtype 18 | nonzero_mask = binary_fill_holes(nonzero_mask) 19 | return nonzero_mask 20 | 21 | 22 | def get_bbox_from_mask(mask, outside_value=0): 23 | mask_voxel_coords = np.where(mask != outside_value) 24 | minzidx = int(np.min(mask_voxel_coords[0])) 25 | maxzidx = int(np.max(mask_voxel_coords[0])) + 1 26 | minxidx = int(np.min(mask_voxel_coords[1])) 27 | maxxidx = int(np.max(mask_voxel_coords[1])) + 1 28 | minyidx = int(np.min(mask_voxel_coords[2])) 29 | maxyidx = int(np.max(mask_voxel_coords[2])) + 1 30 | return [[minzidx, maxzidx], [minxidx, maxxidx], [minyidx, maxyidx]] 31 | 32 | 33 | def crop_to_bbox(image, bbox): 34 | assert len(image.shape) == 3, "only supports 3d images" 35 | resizer = (slice(bbox[0][0], bbox[0][1]), slice(bbox[1][0], bbox[1][1]), slice(bbox[2][0], bbox[2][1])) # ???? used for what ??? slice 36 | return image[resizer] 37 | 38 | 39 | def get_case_identifier(case): 40 | case_identifier = case[0].split("/")[-1].split(".nii.gz")[0][:-5] 41 | return case_identifier 42 | 43 | 44 | def get_case_identifier_from_npz(case): 45 | case_identifier = case.split("/")[-1][:-4] 46 | return case_identifier 47 | 48 | 49 | def load_case_from_list_of_files(data_files, seg_file=None): 50 | assert isinstance(data_files, list) or isinstance(data_files, tuple), "case must be either a list or a tuple" 51 | properties = OrderedDict() 52 | data_itk = [sitk.ReadImage(f) for f in data_files] 53 | 54 | properties["original_size_of_raw_data"] = np.array(data_itk[0].GetSize())[[2, 1, 0]] 55 | properties["original_spacing"] = np.array(data_itk[0].GetSpacing())[[2, 1, 0]] 56 | properties["list_of_data_files"] = data_files 57 | properties["seg_file"] = seg_file 58 | 59 | properties["itk_origin"] = data_itk[0].GetOrigin() 60 | properties["itk_spacing"] = data_itk[0].GetSpacing() 61 | properties["itk_direction"] = data_itk[0].GetDirection() 62 | 63 | data_npy = np.vstack([sitk.GetArrayFromImage(d)[None] for d in data_itk]) # shape (d, h, w) will change to (1, d, h, w) with [None]. Then come to (modalities, d, h, w) 64 | if seg_file is not None: 65 | seg_itk = sitk.ReadImage(seg_file) 66 | seg_npy = sitk.GetArrayFromImage(seg_itk)[None].astype(np.float32) 67 | else: 68 | seg_npy = None 69 | return data_npy.astype(np.float32), seg_npy, properties 70 | 71 | 72 | def crop_to_nonzero(data, seg=None, nonzero_label=-1): 73 | """ 74 | 75 | :param data: 76 | :param seg: 77 | :param nonzero_label: this will be written into the segmentation map 78 | :return: 79 | """ 80 | nonzero_mask = create_nonzero_mask(data) 81 | bbox = get_bbox_from_mask(nonzero_mask, 0) 82 | 83 | cropped_data = [] 84 | for c in range(data.shape[0]): 85 | cropped = crop_to_bbox(data[c], bbox) 86 | cropped_data.append(cropped[None]) 87 | data = np.vstack(cropped_data) 88 | 89 | if seg is not None: 90 | cropped_seg = [] 91 | for c in range(seg.shape[0]): 92 | cropped = crop_to_bbox(seg[c], bbox) 93 | cropped_seg.append(cropped[None]) 94 | seg = np.vstack(cropped_seg) 95 | 96 | nonzero_mask = crop_to_bbox(nonzero_mask, bbox)[None] 97 | if seg is not None: 98 | seg[(seg == 0) & (nonzero_mask == 0)] = nonzero_label # ??????? why ?????? 99 | else: 100 | nonzero_mask = nonzero_mask.astype(int) 101 | nonzero_mask[nonzero_mask == 0] = nonzero_label 102 | nonzero_mask[nonzero_mask > 0] = 0 103 | seg = nonzero_mask 104 | return data, seg, bbox 105 | 106 | 107 | def get_patient_identifiers_from_cropped_files(folder): 108 | return [i.split("/")[-1][:-4] for i in subfiles(folder, join=True, suffix=".npz")] ### ...... feizhejingansha.... 109 | 110 | 111 | class ImageCropper(object): 112 | def __init__(self, num_threads, output_folder=None): 113 | """ 114 | This one finds a mask of nonzero elements (must be nonzero in all modalities) and crops the image to that mask. 115 | In the case of BRaTS and ISLES data this results in a significant reduction in image size 116 | :param num_threads: 117 | :param output_folder: whete to store the cropped data 118 | :param list_of_files: 119 | """ 120 | self.output_folder = output_folder 121 | self.num_threads = num_threads 122 | 123 | if self.output_folder is not None: 124 | maybe_mkdir_p(self.output_folder) 125 | 126 | @staticmethod 127 | def crop(data, properties, seg=None): 128 | shape_before = data.shape 129 | data, seg, bbox = crop_to_nonzero(data, seg, nonzero_label=-1) 130 | shape_after = data.shape 131 | print("before crop:", shape_before, 132 | "after crop:", shape_after, 133 | "spacing:", np.array(properties["original_spacing"]), "\n") 134 | 135 | properties["crop_bbox"] = bbox 136 | properties['classes'] = np.unique(seg) # ???? with -1 ?????? 137 | seg[seg < -1] = 0 # ???? why there could be voxel value smaller than -1 138 | properties["size_after_cropping"] = data[0].shape 139 | return data, seg, properties 140 | 141 | @staticmethod 142 | def crop_from_list_of_files(data_files, seg_file=None): 143 | data, seg, properties = load_case_from_list_of_files(data_files, seg_file) # (m, d, h, w), (1, d, h, w), dict of image information 144 | return ImageCropper.crop(data, properties, seg) 145 | 146 | def load_crop_save(self, case, case_identifier, overwrite_existing=False): 147 | print(case_identifier) 148 | if overwrite_existing \ 149 | or (not os.path.isfile(os.path.join(self.output_folder, "%s.npz" % case_identifier)) 150 | or not os.path.isfile(os.path.join(self.output_folder, "%s.pkl" % case_identifier))): 151 | 152 | data, seg, properties = self.crop_from_list_of_files(case[:-1], case[-1]) 153 | 154 | all_data = np.vstack((data, seg)) 155 | np.savez_compressed(os.path.join(self.output_folder, "%s.npz" % case_identifier), data=all_data) # (modality+1, d h, w) array 156 | with open(os.path.join(self.output_folder, "%s.pkl" % case_identifier), 'wb') as f: 157 | pickle.dump(properties, f) 158 | 159 | def _load_crop_save_star(self, args): 160 | return self.load_crop_save(*args) 161 | 162 | def get_list_of_cropped_files(self): 163 | return subfiles(self.output_folder, join=True, suffix=".npz") 164 | 165 | def get_patient_identifiers_from_cropped_files(self): 166 | return [i.split("/")[-1][:-4] for i in self.get_list_of_cropped_files()] 167 | 168 | def run_cropping(self, list_of_files, overwrite_existing=False, output_folder=None): 169 | """ 170 | also copied ground truth nifti segmentation into the preprocessed folder so that we can use them for evaluation 171 | on the cluster 172 | :param list_of_files: list of list of files [[PATIENTID_TIMESTEP_0000.nii.gz], [PATIENTID_TIMESTEP_0000.nii.gz]] 173 | :param overwrite_existing: 174 | :param output_folder: 175 | :return: 176 | """ 177 | if output_folder is not None: 178 | self.output_folder = output_folder 179 | 180 | output_folder_gt = os.path.join(self.output_folder, "gt_segmentations") 181 | maybe_mkdir_p(output_folder_gt) 182 | for j, case in enumerate(list_of_files): 183 | if case[-1] is not None: 184 | shutil.copy(case[-1], output_folder_gt) 185 | 186 | list_of_args = [] 187 | for j, case in enumerate(list_of_files): 188 | case_identifier = get_case_identifier(case) # patient name 189 | list_of_args.append((case, case_identifier, overwrite_existing)) # (path s, name, bool) 190 | 191 | p = Pool(self.num_threads) 192 | p.map(self._load_crop_save_star, list_of_args) 193 | p.close() 194 | p.join() 195 | 196 | def load_properties(self, case_identifier): 197 | with open(os.path.join(self.output_folder, "%s.pkl" % case_identifier), 'rb') as f: 198 | properties = pickle.load(f) 199 | return properties 200 | 201 | def save_properties(self, case_identifier, properties): 202 | with open(os.path.join(self.output_folder, "%s.pkl" % case_identifier), 'wb') as f: 203 | pickle.dump(properties, f) 204 | -------------------------------------------------------------------------------- /nnunet/inference/predict_simple.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import sys 4 | import os 5 | sys.path.append(os.path.dirname(os.getcwd())) 6 | 7 | import argparse 8 | from nnunet.inference.predict import predict_from_folder 9 | from nnunet.paths import default_plans_identifier, network_training_output_dir 10 | from batchgenerators.utilities.file_and_folder_operations import join, isdir 11 | 12 | if __name__ == "__main__": 13 | parser = argparse.ArgumentParser() 14 | parser.add_argument("-i", '--input_folder', help="Must contain all modalities for each patient in the correct" 15 | " order (same as training). Files must be named " 16 | "CASENAME_XXXX.nii.gz where XXXX is the modality " 17 | "identifier (0000, 0001, etc)", required=True) 18 | parser.add_argument('-o', "--output_folder", required=True, help="folder for saving predictions") 19 | parser.add_argument('-t', '--task_name', help='task name, required.', 20 | default=default_plans_identifier, required=True) 21 | 22 | 23 | parser.add_argument('-tr', '--nnunet_trainer', help='nnUNet trainer class. Default: nnUNetTrainer', required=False, 24 | default='nnUNetTrainer') 25 | parser.add_argument('-m', '--model', help="2d, 3d_lowres, 3d_fullres or 3d_cascade_fullres. Default: 3d_fullres", 26 | default="3d_fullres", required=False) 27 | parser.add_argument('-p', '--plans_identifier', help='do not touch this unless you know what you are doing', 28 | default=default_plans_identifier, required=False) 29 | 30 | parser.add_argument('-f', '--folds', nargs='+', default='None', help="folds to use for prediction. Default is None " 31 | "which means that folds will be detected " 32 | "automatically in the model output folder") 33 | parser.add_argument('-z', '--save_npz', required=False, action='store_true', help="use this if you want to ensemble" 34 | " these predictions with those of" 35 | " other models. Softmax " 36 | "probabilities will be saved as " 37 | "compresed numpy arrays in " 38 | "output_folder and can be merged " 39 | "between output_folders with " 40 | "merge_predictions.py") 41 | parser.add_argument('-l', '--lowres_segmentations', required=False, default='None', help="if model is the highres " 42 | "stage of the cascade then you need to use -l to specify where the segmentations of the " 43 | "corresponding lowres unet are. Here they are required to do a prediction") 44 | parser.add_argument("--part_id", type=int, required=False, default=0, help="Used to parallelize the prediction of " 45 | "the folder over several GPUs. If you " 46 | "want to use n GPUs to predict this " 47 | "folder you need to run this command " 48 | "n times with --part_id=0, ... n-1 and " 49 | "--num_parts=n (each with a different " 50 | "GPU (for example via " 51 | "CUDA_VISIBLE_DEVICES=X)") 52 | parser.add_argument("--num_parts", type=int, required=False, default=1, help="Used to parallelize the prediction of " 53 | "the folder over several GPUs. If you " 54 | "want to use n GPUs to predict this " 55 | "folder you need to run this command " 56 | "n times with --part_id=0, ... n-1 and " 57 | "--num_parts=n (each with a different " 58 | "GPU (via " 59 | "CUDA_VISIBLE_DEVICES=X)") 60 | parser.add_argument("--num_threads_preprocessing", required=False, default=6, type=int, help= 61 | "Determines many background processes will be used for data preprocessing. Reduce this if you " 62 | "run into out of memory (RAM) problems. Default: 6") 63 | parser.add_argument("--num_threads_nifti_save", required=False, default=2, type=int, help= 64 | "Determines many background processes will be used for segmentation export. Reduce this if you " 65 | "run into out of memory (RAM) problems. Default: 2") 66 | parser.add_argument("--tta", required=False, type=int, default=0, help="Set to 0 to disable test time data " 67 | "augmentation (speedup of factor " 68 | "4(2D)/8(3D)), " 69 | "lower quality segmentations") 70 | parser.add_argument("--overwrite_existing", required=False, type=int, default=1, help="Set this to 0 if you need " 71 | "to resume a previous " 72 | "prediction. Default: 1 " 73 | "(=existing segmentations " 74 | "in output_folder will be " 75 | "overwritten)") 76 | parser.add_argument('--gpu', default='0', type=str) 77 | 78 | args = parser.parse_args() 79 | import os 80 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu 81 | 82 | input_folder = args.input_folder 83 | output_folder = args.output_folder 84 | 85 | part_id = args.part_id 86 | num_parts = args.num_parts 87 | folds = args.folds 88 | save_npz = args.save_npz 89 | lowres_segmentations = args.lowres_segmentations 90 | num_threads_preprocessing = args.num_threads_preprocessing 91 | num_threads_nifti_save = args.num_threads_nifti_save 92 | tta = args.tta # default = 0 93 | overwrite = args.overwrite_existing 94 | 95 | output_folder_name = join(network_training_output_dir, # result_path/nnUNet_without_mirror 96 | args.model, # 3d_lowres 97 | args.task_name, # Task12_ 98 | args.nnunet_trainer + "__" + args.plans_identifier) 99 | print('\n'*2, "using model stored in: ", output_folder_name) 100 | assert isdir(output_folder_name), "model output folder not found: %s" % output_folder_name 101 | 102 | if lowres_segmentations == "None": 103 | lowres_segmentations = None 104 | 105 | print("folds: ", folds) 106 | if isinstance(folds, list): 107 | if folds[0] == 'all' and len(folds) == 1: 108 | pass 109 | else: 110 | folds = [int(i) for i in folds] 111 | elif folds == "None": 112 | folds = None 113 | else: 114 | raise ValueError("Unexpected value for argument folds") 115 | 116 | if tta == 0: 117 | tta = False 118 | elif tta == 1: 119 | tta = True 120 | else: 121 | raise ValueError("Unexpected value for tta, Use 1 or 0") 122 | 123 | if overwrite == 0: 124 | overwrite = False 125 | elif overwrite == 1: 126 | overwrite = True 127 | else: 128 | raise ValueError("Unexpected value for overwrite, Use 1 or 0") 129 | 130 | predict_from_folder(output_folder_name, 131 | input_folder, 132 | output_folder, 133 | folds, save_npz, 134 | num_threads_preprocessing, 135 | num_threads_nifti_save, 136 | lowres_segmentations, 137 | part_id, 138 | num_parts, 139 | tta, 140 | overwrite_existing=overwrite) 141 | 142 | -------------------------------------------------------------------------------- /nnunet/training/loss_functions/LovaszSoftmax.py: -------------------------------------------------------------------------------- 1 | """ 2 | Lovasz-Softmax and Jaccard hinge loss in PyTorch 3 | Maxim Berman 2018 ESAT-PSI KU Leuven (MIT License) 4 | """ 5 | 6 | from __future__ import print_function, division 7 | 8 | import torch 9 | from torch.autograd import Variable 10 | import torch.nn.functional as F 11 | import numpy as np 12 | 13 | try: 14 | from itertools import ifilterfalse 15 | except ImportError: # py3k 16 | from itertools import filterfalse as ifilterfalse 17 | 18 | 19 | def lovasz_grad(gt_sorted): 20 | """ 21 | Computes gradient of the Lovasz extension w.r.t sorted errors 22 | See Alg. 1 in paper 23 | """ 24 | p = len(gt_sorted) 25 | gts = gt_sorted.sum() 26 | intersection = gts - gt_sorted.float().cumsum(0) 27 | union = gts + (1 - gt_sorted).float().cumsum(0) 28 | jaccard = 1. - intersection / union 29 | if p > 1: # cover 1-pixel case 30 | jaccard[1:p] = jaccard[1:p] - jaccard[0:-1] 31 | return jaccard 32 | 33 | 34 | def iou_binary(preds, labels, EMPTY=1., ignore=None, per_image=True): 35 | """ 36 | IoU for foreground class 37 | binary: 1 foreground, 0 background 38 | """ 39 | if not per_image: 40 | preds, labels = (preds,), (labels,) 41 | ious = [] 42 | for pred, label in zip(preds, labels): 43 | intersection = ((label == 1) & (pred == 1)).sum() 44 | union = ((label == 1) | ((pred == 1) & (label != ignore))).sum() 45 | if not union: 46 | iou = EMPTY 47 | else: 48 | iou = float(intersection) / float(union) 49 | ious.append(iou) 50 | iou = mean(ious) # mean accross images if per_image 51 | return 100 * iou 52 | 53 | 54 | def iou(preds, labels, C, EMPTY=1., ignore=None, per_image=False): 55 | """ 56 | Array of IoU for each (non ignored) class 57 | """ 58 | if not per_image: 59 | preds, labels = (preds,), (labels,) 60 | ious = [] 61 | for pred, label in zip(preds, labels): 62 | iou = [] 63 | for i in range(C): 64 | if i != ignore: # The ignored label is sometimes among predicted classes (ENet - CityScapes) 65 | intersection = ((label == i) & (pred == i)).sum() 66 | union = ((label == i) | ((pred == i) & (label != ignore))).sum() 67 | if not union: 68 | iou.append(EMPTY) 69 | else: 70 | iou.append(float(intersection) / float(union)) 71 | ious.append(iou) 72 | ious = [mean(iou) for iou in zip(*ious)] # mean accross images if per_image 73 | return 100 * np.array(ious) 74 | 75 | 76 | # --------------------------- BINARY LOSSES --------------------------- 77 | 78 | 79 | def lovasz_hinge(logits, labels, per_image=True, ignore=None): 80 | """ 81 | Binary Lovasz hinge loss 82 | logits: [B, H, W] Variable, logits at each pixel (between -\infty and +\infty) 83 | labels: [B, H, W] Tensor, binary ground truth masks (0 or 1) 84 | per_image: compute the loss per image instead of per batch 85 | ignore: void class id 86 | """ 87 | if per_image: 88 | loss = mean(lovasz_hinge_flat(*flatten_binary_scores(log.unsqueeze(0), lab.unsqueeze(0), ignore)) 89 | for log, lab in zip(logits, labels)) 90 | else: 91 | loss = lovasz_hinge_flat(*flatten_binary_scores(logits, labels, ignore)) 92 | return loss 93 | 94 | 95 | def lovasz_hinge_flat(logits, labels): 96 | """ 97 | Binary Lovasz hinge loss 98 | logits: [P] Variable, logits at each prediction (between -\infty and +\infty) 99 | labels: [P] Tensor, binary ground truth labels (0 or 1) 100 | ignore: label to ignore 101 | """ 102 | if len(labels) == 0: 103 | # only void pixels, the gradients should be 0 104 | return logits.sum() * 0. 105 | signs = 2. * labels.float() - 1. 106 | errors = (1. - logits * Variable(signs)) 107 | errors_sorted, perm = torch.sort(errors, dim=0, descending=True) 108 | perm = perm.data 109 | gt_sorted = labels[perm] 110 | grad = lovasz_grad(gt_sorted) 111 | loss = torch.dot(F.relu(errors_sorted), Variable(grad)) 112 | return loss 113 | 114 | 115 | def flatten_binary_scores(scores, labels, ignore=None): 116 | """ 117 | Flattens predictions in the batch (binary case) 118 | Remove labels equal to 'ignore' 119 | """ 120 | scores = scores.view(-1) 121 | labels = labels.view(-1) 122 | if ignore is None: 123 | return scores, labels 124 | valid = (labels != ignore) 125 | vscores = scores[valid] 126 | vlabels = labels[valid] 127 | return vscores, vlabels 128 | 129 | 130 | class StableBCELoss(torch.nn.modules.Module): 131 | def __init__(self): 132 | super(StableBCELoss, self).__init__() 133 | 134 | def forward(self, input, target): 135 | neg_abs = - input.abs() 136 | loss = input.clamp(min=0) - input * target + (1 + neg_abs.exp()).log() 137 | return loss.mean() 138 | 139 | 140 | def binary_xloss(logits, labels, ignore=None): 141 | """ 142 | Binary Cross entropy loss 143 | logits: [B, H, W] Variable, logits at each pixel (between -\infty and +\infty) 144 | labels: [B, H, W] Tensor, binary ground truth masks (0 or 1) 145 | ignore: void class id 146 | """ 147 | logits, labels = flatten_binary_scores(logits, labels, ignore) 148 | loss = StableBCELoss()(logits, Variable(labels.float())) 149 | return loss 150 | 151 | 152 | # --------------------------- MULTICLASS LOSSES --------------------------- 153 | 154 | 155 | def lovasz_softmax(probas, labels, classes='present', per_image=False, ignore=None): 156 | """ 157 | Multi-class Lovasz-Softmax loss 158 | probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1). 159 | Interpreted as binary (sigmoid) output with outputs of size [B, H, W]. 160 | labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1) 161 | classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average. 162 | per_image: compute the loss per image instead of per batch 163 | ignore: void class labels 164 | """ 165 | if per_image: 166 | loss = mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), classes=classes) 167 | for prob, lab in zip(probas, labels)) 168 | else: 169 | loss = lovasz_softmax_flat(*flatten_probas(probas, labels, ignore), classes=classes) 170 | return loss 171 | 172 | 173 | def lovasz_softmax_flat(probas, labels, classes='present'): 174 | """ 175 | Multi-class Lovasz-Softmax loss 176 | probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1) 177 | labels: [P] Tensor, ground truth labels (between 0 and C - 1) 178 | classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average. 179 | """ 180 | if probas.numel() == 0: 181 | # only void pixels, the gradients should be 0 182 | return probas * 0. 183 | C = probas.size(1) 184 | losses = [] 185 | class_to_sum = list(range(C)) if classes in ['all', 'present'] else classes 186 | for c in class_to_sum: 187 | fg = (labels == c).float() # foreground for class c 188 | if (classes is 'present' and fg.sum() == 0): 189 | continue 190 | if C == 1: 191 | if len(classes) > 1: 192 | raise ValueError('Sigmoid output possible only with 1 class') 193 | class_pred = probas[:, 0] 194 | else: 195 | class_pred = probas[:, c] 196 | errors = (Variable(fg) - class_pred).abs() 197 | errors_sorted, perm = torch.sort(errors, 0, descending=True) 198 | perm = perm.data 199 | fg_sorted = fg[perm] 200 | losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted)))) 201 | return mean(losses) 202 | 203 | 204 | def flatten_probas(probas, labels, ignore=None): 205 | """ 206 | Flattens predictions in the batch 207 | """ 208 | if probas.dim() == 3: 209 | # assumes output of a sigmoid layer 210 | B, H, W = probas.size() 211 | probas = probas.view(B, 1, H, W) 212 | B, C, H, W = probas.size() 213 | probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C) # B * H * W, C = P, C 214 | labels = labels.view(-1) 215 | if ignore is None: 216 | return probas, labels 217 | valid = (labels != ignore) 218 | vprobas = probas[valid.nonzero().squeeze()] 219 | vlabels = labels[valid] 220 | return vprobas, vlabels 221 | 222 | 223 | def xloss(logits, labels, ignore=None): 224 | """ 225 | Cross entropy loss 226 | """ 227 | return F.cross_entropy(logits, Variable(labels), ignore_index=255) 228 | 229 | 230 | # --------------------------- HELPER FUNCTIONS --------------------------- 231 | def isnan(x): 232 | return x != x 233 | 234 | 235 | def mean(l, ignore_nan=False, empty=0): 236 | """ 237 | nanmean compatible with generators. 238 | """ 239 | l = iter(l) 240 | if ignore_nan: 241 | l = ifilterfalse(isnan, l) 242 | try: 243 | n = 1 244 | acc = next(l) 245 | except StopIteration: 246 | if empty == 'raise': 247 | raise ValueError('Empty mean') 248 | return empty 249 | for n, v in enumerate(l, 2): 250 | acc += v 251 | if n == 1: 252 | return acc 253 | return acc / n -------------------------------------------------------------------------------- /evaluation.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import SimpleITK as sitk 3 | import time 4 | import os 5 | from utils import _sitk_Image_reader, _sitk_image_writer, save_pkl 6 | from functools import partial 7 | from multiprocessing import Pool 8 | from collections import OrderedDict 9 | from postprocessing import newsdf_post_processor, maximum_connected_region_post_processor 10 | 11 | def write2singlefile(content, savepath): 12 | with open(savepath,'a+') as f: 13 | f.write(content) 14 | 15 | def computeQualityMeasures(lP, lT): 16 | """ 17 | Binary [0,1] 18 | """ 19 | quality = dict() 20 | try: 21 | labelPred = sitk.GetImageFromArray(lP, isVector=False) 22 | labelTrue = sitk.GetImageFromArray(lT, isVector=False) 23 | hausdorffcomputer = sitk.HausdorffDistanceImageFilter() 24 | hausdorffcomputer.Execute(labelTrue > 0.5, labelPred > 0.5) 25 | quality["avgHausdorff"] = hausdorffcomputer.GetAverageHausdorffDistance() 26 | quality["Hausdorff"] = hausdorffcomputer.GetHausdorffDistance() 27 | 28 | dicecomputer = sitk.LabelOverlapMeasuresImageFilter() 29 | dicecomputer.Execute(labelTrue > 0.5, labelPred > 0.5) 30 | quality["dice"] = dicecomputer.GetDiceCoefficient() 31 | 32 | quality["pred_pixel_num"] = len(lP[lP==1]) 33 | quality["target_pixel_num"] = len(lT[lT==1]) 34 | return quality 35 | except Exception as e: 36 | quality = dict() 37 | print('Exception: ',e) 38 | quality["avgHausdorff"]=0 39 | quality["Hausdorff"]=0 40 | quality["dice"] = 1 41 | quality["pred_pixel_num"] = 0 42 | quality["target_pixel_num"] = 0 43 | return quality 44 | 45 | def computeQualityMeasures_oneCases(name, pred_path, target_path_file, postprocessor, region_th=2000, sdf_th = 0.4): 46 | """ 47 | need modified to suited format 48 | """ 49 | _, pred, _ = _sitk_Image_reader(os.path.join(pred_path, name+'.nii.gz')) 50 | _, target, _ = _sitk_Image_reader(os.path.join(target_path_file, name+'_mask_4label.nii.gz')) 51 | print("computing {} ...".format(name), np.unique(target)) 52 | # write2singlefile("computing {} ...".format(name)+ str(np.unique(target))+'\n', LOG_save_path) 53 | 54 | """ 55 | Come on!!!!!! 56 | 57 | """ 58 | if postprocessor == 'sdf': 59 | pred = newsdf_post_processor(pred, sdf_th=sdf_th, region_th = region_th) 60 | elif postprocessor == 'mcr': 61 | pred = maximum_connected_region_post_processor(pred, region_th=region_th) 62 | elif postprocessor is None: 63 | pass 64 | else: 65 | raise NotImplementedError 66 | 67 | one_case_qualities = OrderedDict() 68 | 69 | range_right = target.max()+1 70 | for i in range(1, range_right): 71 | class_pred = np.zeros_like(pred) 72 | class_target = np.zeros_like(target) 73 | class_pred[pred == i] = 1 74 | class_target[target == i] = 1 75 | 76 | class_quality = computeQualityMeasures(class_pred, class_target) 77 | one_case_qualities[i] = class_quality 78 | 79 | 80 | pred[pred>1] = 1 81 | target[target>1 ]= 1 82 | assert len(np.unique(pred)) == 2 83 | assert len(np.unique(target)) == 2 84 | one_case_qualities['whole'] = computeQualityMeasures(pred, target) 85 | del pred, target 86 | 87 | 88 | hausdorffs = [one_case_qualities[i]["Hausdorff"] for i in range(1,range_right)] 89 | dices = [one_case_qualities[i]["dice"] for i in range(1,range_right)] 90 | pixel_nums = [one_case_qualities[i]["target_pixel_num"] for i in range(1,range_right)] 91 | hausdorffs = np.array(hausdorffs) 92 | dices = np.array(dices) 93 | pixel_nums = np.array(pixel_nums) 94 | 95 | mean_hausdorff = hausdorffs.mean() 96 | weighted_mean_hausdorff = (pixel_nums*hausdorffs).sum()/pixel_nums.sum() 97 | mean_dice = dices.mean() 98 | weighted_mean_dice = (pixel_nums*dices).sum()/pixel_nums.sum() 99 | 100 | one_case_qualities["mean_hausdorff"] = mean_hausdorff 101 | one_case_qualities["mean_dice"] = mean_dice 102 | one_case_qualities["weighted_mean_hausdorff"] = weighted_mean_hausdorff 103 | one_case_qualities["weighted_mean_dice"] = weighted_mean_dice 104 | 105 | for i in range(range_right, 5): 106 | print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! {}: {} tianjia !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!".format(name, i)) 107 | one_case_qualities[i]=dict() 108 | one_case_qualities[i]["Hausdorff"]=0 109 | one_case_qualities[i]["dice"] = 1 110 | print(name,'\n', 111 | '1: ',one_case_qualities[1],'\n', 112 | '2: ',one_case_qualities[2],'\n', 113 | '3: ',one_case_qualities[3],'\n', 114 | '4: ',one_case_qualities[4],'\n', 115 | 'mean_hausdorff: ',one_case_qualities['mean_hausdorff'],'\n', 116 | 'mean_dice: ',one_case_qualities['mean_dice'],'\n', 117 | '-'*33,'\n') 118 | 119 | return (name, one_case_qualities) 120 | 121 | 122 | def computeQualityMeasures_oneModel(pred_path, target_path_file, subdataset, postprocessor, thread, region_th,sdf_th): 123 | """ 124 | pred_path: prediction path 125 | target_path_file: ground truth 126 | subdataset: dataset1, dataset2... 127 | postprocessor: sdf, mcr 128 | 129 | results structure: 130 | dict{patient:{1:quality, 131 | 2:quality, 132 | ...}} 133 | """ 134 | print('pred_path: ', pred_path) 135 | files = os.listdir(pred_path) 136 | 137 | ''' 138 | need modify to suited format 139 | ''' 140 | names = [i.replace('.nii.gz','') for i in files if i.endswith('.nii.gz') and not i.endswith('_metal.nii.gz')] 141 | 142 | if subdataset=='all': 143 | pass 144 | elif subdataset.startswith('dataset'): 145 | names = [i for i in names if i.startswith(subdataset)] 146 | else: 147 | raise EOFError 148 | 149 | print(names) 150 | print(pred_path) 151 | print('names: ',len(names)) 152 | print('post: ',postprocessor) 153 | print('thread: ', thread) 154 | print('thresh:',region_th,sdf_th) 155 | 156 | """ 157 | Come on!!!!!! 158 | """ 159 | pool = Pool(thread) 160 | func = partial(computeQualityMeasures_oneCases, pred_path = pred_path, target_path_file = target_path_file, postprocessor = postprocessor,region_th=region_th,sdf_th=sdf_th) 161 | results = pool.map(func, names) 162 | pool.close() 163 | pool.join() 164 | 165 | di = OrderedDict() 166 | idx = 0 167 | mean_1_dice = 0 168 | mean_1_haud = 0 169 | mean_2_dice = 0 170 | mean_2_haud = 0 171 | mean_3_dice = 0 172 | mean_3_haud = 0 173 | mean_4_dice = 0 174 | mean_4_haud = 0 175 | mwhole_dice = 0 176 | mwhole_haud = 0 177 | mean_dice = 0 178 | mean_haud = 0 179 | w_mean_dice = 0 180 | w_mean_haud = 0 181 | for name, quailty in results: 182 | di[name] = quailty 183 | idx += 1 184 | mean_1_dice += quailty[1]["dice"] 185 | mean_2_dice += quailty[2]["dice"] 186 | mean_3_dice += quailty[3]["dice"] 187 | mean_4_dice += quailty[4]["dice"] 188 | mwhole_dice += quailty["whole"]["dice"] 189 | mean_1_haud += quailty[1]["Hausdorff"] 190 | mean_2_haud += quailty[2]["Hausdorff"] 191 | mean_3_haud += quailty[3]["Hausdorff"] 192 | mean_4_haud += quailty[4]["Hausdorff"] 193 | mwhole_haud += quailty["whole"]["Hausdorff"] 194 | 195 | mean_dice += quailty["mean_dice"] 196 | mean_haud += quailty["mean_hausdorff"] 197 | w_mean_dice += quailty["weighted_mean_dice"] 198 | w_mean_haud += quailty["weighted_mean_hausdorff"] 199 | 200 | print("mean_1_dice: ", mean_1_dice / idx) 201 | print("mean_1_huad: ", mean_1_haud / idx) 202 | print("mean_2_dice: ", mean_2_dice / idx) 203 | print("mean_2_huad: ", mean_2_haud / idx) 204 | print("mean_3_dice: ", mean_3_dice / idx) 205 | print("mean_3_huad: ", mean_3_haud / idx) 206 | print("mean_4_dice: ", mean_4_dice / idx) 207 | print("mean_4_huad: ", mean_4_haud / idx) 208 | print("mwhole_dice: ", mwhole_dice / idx) 209 | print("mwhole_haud: ", mwhole_haud / idx) 210 | print("mean_dice: ", mean_dice / idx) 211 | print("mean_haud: ", mean_haud / idx) 212 | print("w_mean_dice: ", w_mean_dice / idx) 213 | print("w_mean_haud: ", w_mean_haud / idx) 214 | print(pred_path) 215 | print('post: ',postprocessor) 216 | 217 | if postprocessor.endswith('sdf'): 218 | pklsave = os.path.join(pred_path, "evaluation_{}_{}__{}.pkl".format(postprocessor,sdf_th,region_th)) 219 | else: 220 | pklsave = os.path.join(pred_path, "evaluation_{}__{}.pkl".format(postprocessor, region_th)) 221 | save_pkl(di, pklsave) 222 | 223 | print(pklsave,'saved...') 224 | 225 | 226 | if __name__ == '__main__': 227 | t_begin = time.time() 228 | predbasePath = os.path.join(os.environ['HOME'],'all_data/nnUNet/rawdata/ipcai2021_ALL_Test/') 229 | tarPath = os.path.join(os.environ['HOME'],'all_data/nnUNet/rawdata/ipcai2021/') 230 | 231 | print(predbasePath) 232 | for fo in [0]: 233 | computeQualityMeasures_oneModel( 234 | pred_path=predbasePath+ 235 | f'Task22_ipcai2021_T__nnUNet_without_mirror_IPCAI2021_deeps_exclusion__nnUNet_without_mirror_IPCAI2021_deeps_exclusion__fold{fo}_3dcascadefullres_pred', 236 | target_path_file=tarPath, 237 | subdataset='all', 238 | postprocessor='sdf', 239 | thread=64, 240 | region_th=2000, 241 | sdf_th=35) 242 | 243 | t_end = time.time() 244 | print(f'time consuming {t_end-t_begin} s ...') 245 | --------------------------------------------------------------------------------