|
import argparse |
|
import os |
|
import shutil |
|
import random |
|
|
|
import uuid |
|
|
|
from nnunet.inference.predict import predict_from_folder |
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
parser = argparse.ArgumentParser(description='Inference using nnU-Net predict_from_folder Python API') |
|
parser.add_argument('-i', '--input_list', help='Input image file_list.txt') |
|
parser.add_argument('-t', '--tmp_folder', help='Temporary folder', required=True) |
|
parser.add_argument('-o', '--output_folder', help='Output Segmentation folder', required=True) |
|
parser.add_argument('-m', '--model', help='Trained Model', required=True) |
|
parser.add_argument('-v', '--verbose', help='Verbose Output', action='store_true', default=False) |
|
args = vars(parser.parse_args()) |
|
|
|
|
|
args['tmp_folder'] += f'_{str(uuid.uuid4().hex)}' |
|
|
|
|
|
os.mkdir(args['tmp_folder']) |
|
|
|
|
|
with open(args['input_list']) as f: |
|
image_list = f.read().splitlines() |
|
|
|
|
|
image_list_link = [os.path.join(args['tmp_folder'], os.path.basename(x).replace('.nii.gz', '_0000.nii.gz')) |
|
for x in image_list] |
|
|
|
|
|
for src, dst in zip(image_list, image_list_link): |
|
try: |
|
os.link(src, dst) |
|
except: |
|
shutil.copyfile(src, dst) |
|
|
|
|
|
|
|
predict_from_folder(args['model'], args['tmp_folder'], args['output_folder'], folds=None, save_npz=False, |
|
num_threads_preprocessing=6, num_threads_nifti_save=2, |
|
lowres_segmentations=None, part_id=0, num_parts=1, tta=False, |
|
overwrite_existing=False, mode="fastest", overwrite_all_in_gpu=None, |
|
mixed_precision=True, step_size=0.5, checkpoint_name="model_final_checkpoint") |
|
|
|
|
|
|
|
|
|
shutil.rmtree(args['tmp_folder']) |
|
|