Skip to content

Commit f600b7d

Browse files
SimoneBendazzoli93SimoneBendazzoli93
authored
Merge branch 'MIC-DKFZ:master' into master
2 parents e64fda4 + ea8baa3 commit f600b7d

File tree

49 files changed

+1330
-782
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

49 files changed

+1330
-782
lines changed

‎documentation/assets/nnUNetMagician.png

160 KB
Loading

‎documentation/competitions/Toothfairy2/inference_script_semseg_only_customInf2.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def predict_single_npy_array(self, input_image: np.ndarray, image_properties: di
3535
if self.verbose:
3636
print('preprocessing')
3737
preprocessor = self.configuration_manager.preprocessor_class(verbose=self.verbose)
38-
data, _ = preprocessor.run_case_npy(input_image, None, image_properties,
38+
data, _, image_properties = preprocessor.run_case_npy(input_image, None, image_properties,
3939
self.plans_manager,
4040
self.configuration_manager,
4141
self.dataset_json)
@@ -73,7 +73,7 @@ def initialize_from_trained_model_folder(self, model_training_output_dir: str,
7373
for i, f in enumerate(use_folds):
7474
f = int(f) if f != 'all' else f
7575
checkpoint = torch.load(join(model_training_output_dir, f'fold_{f}', checkpoint_name),
76-
map_location=torch.device('cpu'))
76+
map_location=torch.device('cpu'), weights_only=False)
7777
if i == 0:
7878
trainer_name = checkpoint['trainer_name']
7979
configuration_name = checkpoint['init_args']['configuration']
@@ -145,7 +145,7 @@ def predict_preprocessed_image(self, image):
145145
# we are loading parameters on demand instead of loading them upfront. This reduces memory footprint a lot.
146146
# each set of parameters is only used once on the test set (one image) so run time wise this is almost the
147147
# same
148-
self.network.load_state_dict(torch.load(p, map_location=compute_device)['network_weights'])
148+
self.network.load_state_dict(torch.load(p, map_location=compute_device)['network_weights'], weights_only=False)
149149
with torch.autocast(self.device.type, enabled=True):
150150
for sl in tqdm(slicers, disable=not self.allow_tqdm):
151151
pred = self._internal_maybe_mirror_and_predict(data[sl][None].to(compute_device))[0].to(

‎nnunetv2/batch_running/collect_results_custom_Decathlon.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -94,19 +94,21 @@ def summarize(input_file, output_file, folds: Tuple[int, ...], configs: Tuple[st
9494

9595
if __name__ == '__main__':
9696
use_these_trainers = {
97-
'nnUNetTrainer': ('nnUNetPlans', 'nnUNetResEncUNetPlans', 'nnUNetResEncUNet2Plans', 'nnUNetResBottleneckEncUNetPlans', 'nnUNetResUNetPlans', 'nnUNetResUNet2Plans', 'nnUNetResUNet3Plans', 'nnUNetDeeperResBottleneckEncUNetPlans'),
98-
}
97+
'nnUNetTrainer': ('nnUNetResEncUNetMPlans', ),
98+
'nnUNetTrainer_probabilisticOversampling_033': ('nnUNetResEncUNetMPlans', ),
99+
'nnUNetTrainer_probabilisticOversampling_010': ('nnUNetResEncUNetMPlans',),
100+
}
99101
all_results_file= join(nnUNet_results, 'customDecResults.csv')
100-
datasets = [2, 3, 4, 17, 24, 27, 38, 55, 137, 217, 220, 221, 223] # amos post challenge, kits2023
102+
datasets = [3, 4, 5, 8, 10, 17, 27, 55, 220, 223]
101103
collect_results(use_these_trainers, datasets, all_results_file)
102104

103105
folds = (0, 1, 2, 3, 4)
104-
configs = ("3d_fullres", )
106+
configs = ("2d", )
105107
output_file = join(nnUNet_results, 'customDecResults_summary5fold.csv')
106108
summarize(all_results_file, output_file, folds, configs, datasets, use_these_trainers)
107109

108110
folds = (0, )
109-
configs = ("3d_fullres", )
111+
configs = ("2d", )
110112
output_file = join(nnUNet_results, 'customDecResults_summaryfold0.csv')
111113
summarize(all_results_file, output_file, folds, configs, datasets, use_these_trainers)
112114

‎nnunetv2/batch_running/generate_lsf_runs_customDecathlon.py

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -55,19 +55,20 @@ def merge(dict1, dict2):
5555
}
5656

5757
num_gpus = 1
58-
exclude_hosts = "-R \"select[hname!='e230-dgx2-2']\" -R \"select[hname!='e230-dgx2-1']\" -R \"select[hname!='lsf22-gpu02']\" -R \"select[hname!='lsf22-gpu06']\" -R \"select[hname!='e230-dgx1-1']\""
58+
exclude_hosts = "-R \"select[hname!='e230-dgx2-2']\" -R \"select[hname!='e230-dgx2-1']\""
5959
resources = ""
60-
gpu_requirements = f"-gpu num={num_gpus}:j_exclusive=yes:gmem=33G"#gmodel=NVIDIAA100_PCIE_40GB"
61-
queue = "-q gpu-lowprio"
62-
preamble = "\"source ~/load_env_torch221.sh && " # -L /bin/bash
60+
gpu_requirements = f"-gpu num={num_gpus}:j_exclusive=yes:gmem=1G"#gmodel=NVIDIAA100_PCIE_40GB"
61+
queue = "-q gpu"
62+
preamble = "\". ~/load_env_torch224.sh && " # -L /bin/bash
6363
train_command = 'nnUNetv2_train'
6464

65+
6566
folds = (0, )
6667
# use_this = configurations_2d_only
67-
use_this = configurations_3d_fr_only
68+
use_this = configurations_2d_only
6869
# use_this = merge(use_this, configurations_3d_c_only)
6970

70-
datasets = [3, 4, 5, 8, 10, 27, 55, 220, 223]
71+
datasets = [3, 4, 5, 8, 10, 17, 27, 55, 220, 223]
7172
use_this = {i: use_this[i] for i in datasets}
7273

7374
use_these_modules = {
@@ -85,7 +86,10 @@ def merge(dict1, dict2):
8586
# 'nnUNetTrainerUMambaBot': ('nnUNetPlans',),
8687
# 'nnUNetTrainerUMambaEnc': ('nnUNetPlans',),
8788
# 'nnUNetTrainer_fasterDA': ('nnUNetPlans', 'nnUNetResEncUNetLPlans'),
88-
'nnUNetTrainer_fasterDA_bg_style_seg_sampling': ('nnUNetResEncUNetLPlans', ),
89+
# 'nnUNetTrainer_noDummy2DDA': ('nnUNetResEncUNetMPlans', ),
90+
'nnUNetTrainer': ('nnUNetResEncUNetMPlans', ),
91+
# 'nnUNetTrainer_probabilisticOversampling_033': ('nnUNetResEncUNetMPlans', ),
92+
# 'nnUNetTrainer_probabilisticOversampling_010': ('nnUNetResEncUNetMPlans',),
8993
# BN
9094
}
9195

‎nnunetv2/batch_running/jobs.sh

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# lsf22-gpu01
2+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=3 nnUNetv2_train 3 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
3+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=4 nnUNetv2_train 4 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
4+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=5 nnUNetv2_train 5 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
5+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=6 nnUNetv2_train 8 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
6+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=7 nnUNetv2_train 10 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
7+
8+
# lsf22-gpu03
9+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=0 nnUNetv2_train 17 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
10+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 27 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
11+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=2 nnUNetv2_train 55 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
12+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=6 nnUNetv2_train 220 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
13+
14+
# lsf22-gpu05
15+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=0 nnUNetv2_train 223 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
16+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 3 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
17+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=2 nnUNetv2_train 4 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
18+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=3 nnUNetv2_train 5 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
19+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=4 nnUNetv2_train 8 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
20+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=6 nnUNetv2_train 10 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
21+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=7 nnUNetv2_train 17 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
22+
23+
# lsf22-gpu06
24+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=0 nnUNetv2_train 27 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
25+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 55 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
26+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=2 nnUNetv2_train 220 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
27+
screen -dm bash -c ". ~/load_env_torch224.sh && CUDA_VISIBLE_DEVICES=3 nnUNetv2_train 223 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
28+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=4 nnUNetv2_train 3 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
29+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=5 nnUNetv2_train 4 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
30+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=6 nnUNetv2_train 5 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
31+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=7 nnUNetv2_train 8 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
32+
33+
# lsf22-gpu07
34+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=0 nnUNetv2_train 10 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
35+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 17 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
36+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=2 nnUNetv2_train 27 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
37+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=3 nnUNetv2_train 55 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
38+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=4 nnUNetv2_train 220 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
39+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=5 nnUNetv2_train 223 3d_fullres 0 -tr nnUNetTrainer_noDummy2DDA -p nnUNetResEncUNetMPlans --disable_checkpointing"
40+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=6 nnUNetv2_train 3 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
41+
screen -dm bash -c ". ~/load_env_torch224_balintsfix.sh && CUDA_VISIBLE_DEVICES=7 nnUNetv2_train 4 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
42+
43+
44+
# launched as jobs
45+
bsub -R "select[hname!='e230-dgx2-2']" -R "select[hname!='e230-dgx2-1']" -q gpu -gpu num=1:j_exclusive=yes:gmem=1G ". ~/load_env_torch224_balintsfix.sh && nnUNetv2_train 5 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
46+
bsub -R "select[hname!='e230-dgx2-2']" -R "select[hname!='e230-dgx2-1']" -q gpu -gpu num=1:j_exclusive=yes:gmem=1G ". ~/load_env_torch224_balintsfix.sh && nnUNetv2_train 8 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
47+
bsub -R "select[hname!='e230-dgx2-2']" -R "select[hname!='e230-dgx2-1']" -q gpu -gpu num=1:j_exclusive=yes:gmem=1G ". ~/load_env_torch224_balintsfix.sh && nnUNetv2_train 10 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
48+
bsub -R "select[hname!='e230-dgx2-2']" -R "select[hname!='e230-dgx2-1']" -q gpu -gpu num=1:j_exclusive=yes:gmem=1G ". ~/load_env_torch224_balintsfix.sh && nnUNetv2_train 17 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
49+
bsub -R "select[hname!='e230-dgx2-2']" -R "select[hname!='e230-dgx2-1']" -q gpu -gpu num=1:j_exclusive=yes:gmem=1G ". ~/load_env_torch224_balintsfix.sh && nnUNetv2_train 27 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
50+
bsub -R "select[hname!='e230-dgx2-2']" -R "select[hname!='e230-dgx2-1']" -q gpu -gpu num=1:j_exclusive=yes:gmem=1G ". ~/load_env_torch224_balintsfix.sh && nnUNetv2_train 55 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
51+
bsub -R "select[hname!='e230-dgx2-2']" -R "select[hname!='e230-dgx2-1']" -q gpu -gpu num=1:j_exclusive=yes:gmem=1G ". ~/load_env_torch224_balintsfix.sh && nnUNetv2_train 220 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
52+
bsub -R "select[hname!='e230-dgx2-2']" -R "select[hname!='e230-dgx2-1']" -q gpu -gpu num=1:j_exclusive=yes:gmem=1G ". ~/load_env_torch224_balintsfix.sh && nnUNetv2_train 223 3d_fullres 0 -tr nnUNetTrainer -p nnUNetResEncUNetMPlans --disable_checkpointing"
53+

‎nnunetv2/dataset_conversion/Dataset015_018_RibFrac_RibSeg.py

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
from copy import deepcopy
2+
3+
import numpy as np
4+
from batchgenerators.utilities.file_and_folder_operations import *
5+
import shutil
6+
from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json
7+
from nnunetv2.paths import nnUNet_raw
8+
import SimpleITK as sitk
9+
10+
if __name__ == '__main__':
11+
"""
12+
Download RibFrac dataset. Links are at https://ribfrac.grand-challenge.org/
13+
Download everything. Part1, 2, validation and test
14+
Extract EVERYTHING into one folder so that all images and labels are in there. Don't worry they all have unique
15+
file names.
16+
17+
For RibSeg also download the dataset from https://github.com/M3DV/RibSeg
18+
(https://drive.google.com/file/d/1ZZGGrhd0y1fLyOZGo_Y-wlVUP4lkHVgm/view?usp=sharing) and extract in to that same
19+
folder (seg only, files end with -rib-seg.nii.gz)
20+
"""
21+
# extracted traiing.zip file is here
22+
base = '/home/isensee/Downloads/RibFrac_all'
23+
24+
files = nifti_files(base, join=False)
25+
identifiers = np.unique([i.split('-')[0] for i in files])
26+
27+
# RibFrac
28+
target_dataset_id = 15
29+
target_dataset_name = f'Dataset{target_dataset_id:03.0f}_RibFrac'
30+
31+
maybe_mkdir_p(join(nnUNet_raw, target_dataset_name))
32+
imagesTr = join(nnUNet_raw, target_dataset_name, 'imagesTr')
33+
imagesTs = join(nnUNet_raw, target_dataset_name, 'imagesTs')
34+
labelsTr = join(nnUNet_raw, target_dataset_name, 'labelsTr')
35+
maybe_mkdir_p(imagesTr)
36+
maybe_mkdir_p(imagesTs)
37+
maybe_mkdir_p(labelsTr)
38+
39+
n_tr = 0
40+
for c in identifiers:
41+
print(c)
42+
img_file = join(base, c + '-image.nii.gz')
43+
seg_file = join(base, c + '-label.nii.gz')
44+
if not isfile(seg_file):
45+
# test case
46+
shutil.copy(img_file, join(imagesTs, c + '_0000.nii.gz'))
47+
continue
48+
n_tr += 1
49+
shutil.copy(img_file, join(imagesTr, c + '_0000.nii.gz'))
50+
51+
# we must open seg and map -1 to 5
52+
seg_itk = sitk.ReadImage(seg_file)
53+
seg_npy = sitk.GetArrayFromImage(seg_itk)
54+
seg_npy[seg_npy == -1] = 5
55+
seg_itk_out = sitk.GetImageFromArray(seg_npy.astype(np.uint8))
56+
seg_itk_out.SetSpacing(seg_itk.GetSpacing())
57+
seg_itk_out.SetDirection(seg_itk.GetDirection())
58+
seg_itk_out.SetOrigin(seg_itk.GetOrigin())
59+
sitk.WriteImage(seg_itk_out, join(labelsTr, c + '.nii.gz'))
60+
61+
# - 0: it is background
62+
# - 1: it is a displaced rib fracture
63+
# - 2: it is a non-displaced rib fracture
64+
# - 3: it is a buckle rib fracture
65+
# - 4: it is a segmental rib fracture
66+
# - -1: it is a rib fracture, but we could not define its type due to
67+
# ambiguity, diagnosis difficulty, etc. Ignore it in the
68+
# classification task.
69+
70+
generate_dataset_json(
71+
join(nnUNet_raw, target_dataset_name),
72+
channel_names={0: 'CT'},
73+
labels = {
74+
'background': 0,
75+
'fracture': (1, 2, 3, 4, 5),
76+
'displaced rib fracture': 1,
77+
'non-displaced rib fracture': 2,
78+
'buckle rib fracture': 3,
79+
'segmental rib fracture': 4,
80+
},
81+
num_training_cases=n_tr,
82+
file_ending='.nii.gz',
83+
regions_class_order=(5, 1, 2, 3, 4),
84+
dataset_name=target_dataset_name,
85+
reference='https://ribfrac.grand-challenge.org/'
86+
)
87+
88+
# RibSeg
89+
# overall I am not happy with the GT quality here. But eh what can I do
90+
91+
target_dataset_name_ribfrac = deepcopy(target_dataset_name)
92+
target_dataset_id = 18
93+
target_dataset_name = f'Dataset{target_dataset_id:03.0f}_RibSeg'
94+
95+
maybe_mkdir_p(join(nnUNet_raw, target_dataset_name))
96+
imagesTr = join(nnUNet_raw, target_dataset_name, 'imagesTr')
97+
labelsTr = join(nnUNet_raw, target_dataset_name, 'labelsTr')
98+
maybe_mkdir_p(imagesTr)
99+
maybe_mkdir_p(labelsTr)
100+
101+
# the authors have a google shet where they highlight problems with their dataset:
102+
# https://docs.google.com/spreadsheets/d/1lz9liWPy8yHybKCdO3BCA9K76QH8a54XduiZS_9fK70/edit?gid=1416415020#gid=1416415020
103+
# we exclude the cases marked in red. They have unannotated ribs
104+
skip_identifiers = [
105+
'RibFrac452',
106+
'RibFrac485',
107+
'RibFrac490',
108+
'RibFrac471',
109+
'RibFrac462',
110+
'RibFrac487',
111+
]
112+
113+
n_tr = 0
114+
dataset = {}
115+
for c in identifiers:
116+
if c in skip_identifiers:
117+
continue
118+
print(c)
119+
tr_file = join('$nnUNet_raw', target_dataset_name_ribfrac, 'imagesTr', c + '_0000.nii.gz')
120+
ts_file = join('$nnUNet_raw', target_dataset_name_ribfrac, 'imagesTs', c + '_0000.nii.gz')
121+
if isfile(os.path.expandvars(tr_file)):
122+
img_file = tr_file
123+
elif isfile(os.path.expandvars(ts_file)):
124+
img_file = ts_file
125+
else:
126+
raise RuntimeError(f'Missing image file for identifier {identifiers}')
127+
seg_file = join(base, c + '-rib-seg.nii.gz')
128+
n_tr += 1
129+
shutil.copy(seg_file, join(labelsTr, c + '.nii.gz'))
130+
dataset[c] = {
131+
'images': [img_file],
132+
'label': join('labelsTr', c + '.nii.gz')
133+
}
134+
135+
generate_dataset_json(
136+
join(nnUNet_raw, target_dataset_name),
137+
channel_names={0: 'CT'},
138+
labels = {
139+
'background': 0,
140+
**{'rib%02.0d' % i: i for i in range(1, 25)}
141+
},
142+
num_training_cases=n_tr,
143+
file_ending='.nii.gz',
144+
dataset_name=target_dataset_name,
145+
reference='https://github.com/M3DV/RibSeg, https://ribfrac.grand-challenge.org/',
146+
dataset=dataset
147+
)

0 commit comments

Comments
 (0)