1919import shlex
2020import shutil
2121import subprocess
22- from tempfile import NamedTemporaryFile , TemporaryDirectory
2322import warnings
2423from concurrent .futures import ThreadPoolExecutor
24+ from tempfile import NamedTemporaryFile , TemporaryDirectory
2525from typing import Any
2626
2727import monai
3030 analyze_data ,
3131 create_new_data_copy ,
3232 create_new_dataset_json ,
33- glob_to_datalist ,
34- check_existing_data_indices ,
35- get_next_available_index ,
3633 get_info_from_dataset_json ,
37- move_predictions
34+ get_next_available_index ,
35+ glob_to_datalist ,
36+ move_predictions ,
3837)
3938from monai .bundle import ConfigParser
4039from monai .utils import ensure_tuple , optional_import
@@ -212,8 +211,8 @@ def __init__(
212211 os .environ ["OMP_NUM_THREADS" ] = str (1 )
213212
214213 # dataset_name_or_id has to be a string
215- if ' dataset_name_or_id' in self .input_info :
216- self .dataset_name_or_id = str (self .input_info [' dataset_name_or_id' ])
214+ if " dataset_name_or_id" in self .input_info :
215+ self .dataset_name_or_id = str (self .input_info [" dataset_name_or_id" ])
217216 else :
218217 # we get the next available index
219218 self .dataset_name_or_id = str (get_next_available_index (self .nnunet_raw ))
@@ -285,7 +284,9 @@ def convert_dataset(self, testing=False):
285284 if isinstance (datalist_json [test_key ][0 ], dict ) and "label" in datalist_json [test_key ][0 ]:
286285 os .makedirs (os .path .join (raw_data_foldername , "labelsTs" ))
287286
288- num_input_channels , num_foreground_classes = self .input_info .get ('num_input_channels' ), self .input_info .get ('num_foreground_classes' )
287+ num_input_channels , num_foreground_classes = self .input_info .get ("num_input_channels" ), self .input_info .get (
288+ "num_foreground_classes"
289+ )
289290
290291 if num_input_channels is None or num_foreground_classes is None :
291292 # can't get num_foreground classes from the data, so should be inserted by user
@@ -296,7 +297,7 @@ def convert_dataset(self, testing=False):
296297 modality = [modality ]
297298
298299 create_new_dataset_json (
299- # pyrefly: ignore [bad-argument-type]
300+ # pyrefly: ignore [bad-argument-type]
300301 modality = modality ,
301302 num_foreground_classes = num_foreground_classes ,
302303 num_input_channels = num_input_channels ,
@@ -439,13 +440,7 @@ def preprocess(
439440 logger .info ("Preprocessing..." )
440441
441442 plans_name = overwrite_plans_name if overwrite_plans_name is not None else self .plans_identifier
442- preprocess (
443- [int (self .dataset_name_or_id )],
444- plans_name ,
445- configurations = c ,
446- num_processes = n_proc ,
447- verbose = verbose ,
448- )
443+ preprocess ([int (self .dataset_name_or_id )], plans_name , configurations = c , num_processes = n_proc , verbose = verbose )
449444 self .plans_identifier = plans_name
450445
451446 def plan_and_process (
@@ -1058,7 +1053,7 @@ def predict_datalist(
10581053 modality : str = "CT" ,
10591054 num_foreground_classes : int | None = None ,
10601055 num_input_channels : int | None = None ,
1061- work_dir : str = ' work_dir' ,
1056+ work_dir : str = " work_dir" ,
10621057 ):
10631058 """Method to run inference on a datalist using a model trained by this runner.
10641059 Handles all nnUNet boilerplate, instantiation of the runner, etc.
@@ -1112,17 +1107,13 @@ def predict_datalist(
11121107 # these things are hardcoded upstream
11131108 raw_data_foldername_prefix = str (int (runner .dataset_name_or_id ) + 1000 )
11141109 raw_data_foldername_prefix = "Dataset" + raw_data_foldername_prefix [- 3 :]
1115- raw_data_foldername = raw_data_foldername_prefix + "_" + input_config [' dataroot' ].split (os .sep )[- 1 ]
1116- raw_data_foldername = os .path .join (input_config [' nnunet_raw' ], raw_data_foldername )
1110+ raw_data_foldername = raw_data_foldername_prefix + "_" + input_config [" dataroot" ].split (os .sep )[- 1 ]
1111+ raw_data_foldername = os .path .join (input_config [" nnunet_raw" ], raw_data_foldername )
11171112
11181113 with TemporaryDirectory () as pred_work_folder :
11191114 test_images_dir = os .path .join (raw_data_foldername , "imagesTs" ) # Also hardcoded upstream
11201115
1121- runner .predict (
1122- test_images_dir ,
1123- output_folder = pred_work_folder ,
1124- model_training_output_dir = model_dir ,
1125- )
1116+ runner .predict (test_images_dir , output_folder = pred_work_folder , model_training_output_dir = model_dir )
11261117 move_predictions (raw_data_foldername , pred_work_folder , output_dir )
11271118
11281119 # now we can delete the raw data folder too
@@ -1141,18 +1132,18 @@ def predict_files_glob(
11411132 modality : str = "CT" ,
11421133 ):
11431134 """Method to run inference on a glob of files using a model trained by this runner.
1144-
1135+
11451136 Creates a temporary datalist json file from the glob of files, and then calls predict_datalist.
1146-
1137+
11471138 Args:
11481139 input_files_glob: glob pattern to match input files (e.g., "/path/to/images/*.nii.gz")
11491140 input_files_root: root directory for the input files (e.g., "/path/to/images")
11501141 model_dir: path to the folder containing the trained model (full path inside the work_dir, e.g., work_dir/nnUNet_trained_models/Dataset001_data/nnUNetTrainer__nnUNetPlans__3d_fullres)
11511142 output_dir: path to the output directory, predictions will be saved here under their original names.
1152- work_dir: path to the work_dir created by the runner during training.
1143+ work_dir: path to the work_dir created by the runner during training.
11531144 modality: modality of the input data (default: "CT")
11541145 """
1155- with NamedTemporaryFile (mode = 'w+' , delete = False , suffix = ' .json' ) as temp_json_file :
1146+ with NamedTemporaryFile (mode = "w+" , delete = False , suffix = " .json" ) as temp_json_file :
11561147 temp_json_path = temp_json_file .name
11571148 glob_to_datalist (input_files_glob , output_json = temp_json_path , key = "testing" , dataroot = input_files_root )
11581149
@@ -1162,23 +1153,25 @@ def predict_files_glob(
11621153 model_dir = model_dir ,
11631154 work_dir = work_dir ,
11641155 output_dir = output_dir ,
1165- modality = modality
1156+ modality = modality ,
11661157 )
11671158
11681159 def _determine_configs (self ):
11691160 from nnunetv2 .paths import nnUNet_preprocessed
11701161 from nnunetv2 .utilities .dataset_name_id_conversion import maybe_convert_to_dataset_name
11711162
1172- preprocessed_dataset_folder_base = join (nnUNet_preprocessed , maybe_convert_to_dataset_name (self .dataset_name_or_id ))
1173- plans_file = join (preprocessed_dataset_folder_base , self .plans_identifier + '.json' )
1163+ preprocessed_dataset_folder_base = join (
1164+ nnUNet_preprocessed , maybe_convert_to_dataset_name (self .dataset_name_or_id )
1165+ )
1166+ plans_file = join (preprocessed_dataset_folder_base , self .plans_identifier + ".json" )
11741167
1175- with open (plans_file , 'r' ) as f :
1168+ with open (plans_file ) as f :
11761169 plans = json .load (f )
11771170
1178- configurations = plans .get (' configurations' , [])
1171+ configurations = plans .get (" configurations" , [])
11791172 if not configurations :
11801173 raise ValueError (f"No configurations found in plans file: { plans_file } " )
1181-
1174+
11821175 config_names = list (configurations .keys ())
11831176 return config_names
11841177
0 commit comments