def store_and_upload_example(
dataset_example: DatasetExample,
segmentation_config: SegmentationModelBase) -> None:
"""
Stores an example input and output of the network to Nifti files.
:param dataset_example: The dataset example, with image, label and prediction, that should be written.
:param segmentation_config: configuration information to be used for normalization and example_images_folder
"""
folder = segmentation_config.example_images_folder
os.makedirs(folder, exist_ok=True)
def create_file_name(suffix: str) -> str:
fn = "p" + str(dataset_example.patient_id) + "_e_" + str(
dataset_example.epoch) + "_" + suffix + ".nii.gz"
fn = os.path.join(folder, fn)
return fn
io_util.store_image_as_short_nifti(
image=dataset_example.image,
header=dataset_example.header,
file_name=create_file_name(suffix="image"),
args=segmentation_config)
# merge multiple binary masks (one per class) into a single multi-label map image
labels = image_util.merge_masks(dataset_example.labels)
io_util.store_as_ubyte_nifti(image=labels,
header=dataset_example.header,
file_name=create_file_name(suffix="label"))
io_util.store_as_ubyte_nifti(
image=dataset_example.prediction,
header=dataset_example.header,
file_name=create_file_name(suffix="prediction"))
def test_store_as_ubyte_nifti(test_output_dirs: TestOutputDirectories) -> None:
image = np.random.random_sample((dim_z, dim_y, dim_x))
# get values in [0, 255] range
image = np.array((image + 1) * 255).astype(int)
header = ImageHeader(origin=(1, 1, 1), direction=(1, 0, 0, 0, 1, 0, 0, 0, 1), spacing=(1, 2, 4))
io_util.store_as_ubyte_nifti(image, header, test_output_dirs.create_file_or_folder_path(default_image_name))
t = np.unique(image).astype(np.ubyte)
assert_nifti_content(test_output_dirs.create_file_or_folder_path(default_image_name),
image.shape, header, list(t), np.ubyte)
def score_image(args: ScorePipelineConfig) -> Path:
"""
Perform model inference on a single image. By doing the following:
1) Copy the provided data root directory to the root (this contains the model checkpoints and image to infer)
2) Instantiate an inference pipeline based on the provided model_inference.json in the snapshot
3) Store the segmentation file in the current directory
4) Upload the segmentation to AML
:param args:
:return:
"""
logging.getLogger().setLevel(logging.INFO)
score_py_folder = Path(__file__).parent
model_folder = Path(args.model_folder or str(score_py_folder))
run_context = Run.get_context()
logging.info(f"Run context={run_context.id}")
if args.use_dicom:
# Only a single zip file is supported.
if len(args.image_files) > 1:
raise ValueError("Supply exactly one zip file in args.images.")
input_zip_file = check_input_file(args.data_folder,
args.image_files[0])
reference_series_folder = model_folder / "temp_extraction"
nifti_filename = model_folder / "temp_nifti.nii.gz"
convert_zipped_dicom_to_nifti(input_zip_file, reference_series_folder,
nifti_filename)
test_images = [nifti_filename]
else:
test_images = [
check_input_file(args.data_folder, file)
for file in args.image_files
]
images = [load_nifti_image(file) for file in test_images]
inference_pipeline, config = init_from_model_inference_json(
model_folder, args.use_gpu)
segmentation = run_inference(images, inference_pipeline, config)
segmentation_file_name = model_folder / args.result_image_name
result_dst = store_as_ubyte_nifti(segmentation, images[0].header,
segmentation_file_name)
if args.use_dicom:
result_dst = convert_nifti_to_zipped_dicom_rt(
result_dst, reference_series_folder, model_folder, config,
args.result_zip_dicom_name, args.model_id)
if not is_offline_run_context(run_context):
upload_file_name = args.result_zip_dicom_name if args.use_dicom else args.result_image_name
run_context.upload_file(upload_file_name, str(result_dst))
logging.info(f"Segmentation completed: {result_dst}")
return result_dst
def store_and_upload_example(dataset_example: DatasetExample,
args: Optional[SegmentationModelBase],
images_folder: Optional[Path] = None) -> None:
"""
Stores an example input and output of the network to Nifti files.
:param dataset_example: The dataset example, with image, label and prediction, that should be written.
:param args: configuration information to be used for normalization. TODO: This should not be optional why is this
assigning to example_images_folder
:param images_folder: The folder to which the result Nifti files should be written. If args is not None,
the args.example_images_folder is used instead.
"""
folder = Path("") if images_folder is None else images_folder
if args is not None:
folder = args.example_images_folder
if folder != "" and not os.path.exists(folder):
os.mkdir(folder)
def create_file_name(suffix: str) -> str:
fn = "p" + str(dataset_example.patient_id) + "_e_" + str(
dataset_example.epoch) + "_" + suffix + ".nii.gz"
fn = os.path.join(folder, fn)
return fn
io_util.store_image_as_short_nifti(
image=dataset_example.image,
header=dataset_example.header,
file_name=create_file_name(suffix="image"),
args=args)
# merge multiple binary masks (one per class) into a single multi-label map image
labels = image_util.merge_masks(dataset_example.labels)
io_util.store_as_ubyte_nifti(image=labels,
header=dataset_example.header,
file_name=create_file_name(suffix="label"))
io_util.store_as_ubyte_nifti(
image=dataset_example.prediction,
header=dataset_example.header,
file_name=create_file_name(suffix="prediction"))
def score_image(args: ScorePipelineConfig) -> Path:
"""
Perform model inference on a single image. By doing the following:
1) Copy the provided data root directory to the root (this contains the model checkpoints and image to infer)
2) Instantiate an inference pipeline based on the provided model_inference.json in the snapshot
3) Store the segmentation file in the current directory
4) Upload the segmentation to AML
:param args:
:return:
"""
logging.getLogger().setLevel(logging.INFO)
score_py_folder = Path(__file__).parent
model_folder = Path(args.model_folder or str(score_py_folder))
run_context = Run.get_context()
logging.info(f"Run context={run_context.id}")
test_images = []
data_folder = args.data_folder
for file in args.image_files:
full_file_path = data_folder / file
if not full_file_path.exists():
message = \
str(data_folder) if data_folder.is_absolute() else f"{data_folder}, absolute: {data_folder.absolute()}"
raise ValueError(
f"File {file} does not exist in data folder {message}")
test_images.append(full_file_path)
images = [load_nifti_image(file) for file in test_images]
inference_pipeline, config = init_from_model_inference_json(
model_folder, args.use_gpu)
segmentation = run_inference(images, inference_pipeline, config)
segmentation_file_name = str(model_folder / args.result_image_name)
result_dst = store_as_ubyte_nifti(segmentation, images[0].header,
segmentation_file_name)
if not is_offline_run_context(run_context):
run_context.upload_file(args.result_image_name, segmentation_file_name)
logging.info(f"Segmentation completed: {result_dst}")
return result_dst
def score_image(args: ScorePipelineConfig) -> Path:
"""
Perform model inference on a single image. By doing the following:
1) Copy the provided data root directory to the root (this contains the model checkpoints and image to infer)
2) Instantiate an inference pipeline based on the provided model_inference.json in the snapshot
3) Store the segmentation file in the current directory
4) Upload the segmentation to AML
:param args:
:return:
"""
logging.getLogger().setLevel(logging.INFO)
project_root = Path(args.project_root)
# copy the model to the current directory
copy_tree(args.data_root, str(project_root))
logging.info(
f'Copied contents of data_root: {args.data_root} to {project_root}')
run_context = Run.get_context()
logging.info(f"Run context={run_context.id}")
images = [
load_nifti_image(project_root / DEFAULT_DATA_FOLDER / x)
for x in args.test_image_channels
]
inference_pipeline, config = init_from_model_inference_json(
project_root, args.use_gpu)
segmentation = run_inference(images, inference_pipeline, config)
segmentation_file_name = str(project_root / args.result_image_name)
result_dst = store_as_ubyte_nifti(segmentation, images[0].header,
segmentation_file_name)
if not is_offline_run_context(run_context):
run_context.upload_file(args.result_image_name, segmentation_file_name)
logging.info(f"Segmentation completed: {result_dst}")
return Path(result_dst)
def store_inference_results(inference_result: InferencePipeline.Result,
config: SegmentationModelBase,
results_folder: Path,
image_header: ImageHeader) -> List[str]:
"""
Store the segmentation, posteriors, and binary predictions into Nifti files.
:param inference_result: The inference result for a given patient_id and epoch. Posteriors must be in
(Classes x Z x Y x X) shape, segmentation in (Z, Y, X)
:param config: The test configurations.
:param results_folder: The folder where the prediction should be stored.
:param image_header: The image header that was used in the input image.
"""
def create_file_path(_results_folder: Path, _file_name: str) -> Path:
"""
Create filename with Nifti extension
:param _results_folder: The results folder
:param _file_name: The name of the file
:return: A full path to the results folder for the file
"""
file_path = _file_name + MedicalImageFileType.NIFTI_COMPRESSED_GZ.value
return _results_folder / Path(file_path)
# create the directory for the given patient inside the results dir
patient_results_folder = get_patient_results_folder(
results_folder, inference_result.patient_id)
patient_results_folder.mkdir(exist_ok=True, parents=True)
# write the segmentations to disk
image_paths = [
io_util.store_as_ubyte_nifti(image=inference_result.segmentation,
header=image_header,
file_name=str(
create_file_path(
patient_results_folder,
"segmentation")))
]
class_names_and_indices = config.class_and_index_with_background().items()
binaries = binaries_from_multi_label_array(inference_result.segmentation,
config.number_of_classes)
# rescale posteriors if required and save them
for (class_name, index), binary in zip(class_names_and_indices, binaries):
posterior = inference_result.posteriors[index, ...]
# save the posterior map
file_name = "posterior_{}".format(class_name)
image_path = io_util.store_posteriors_as_nifti(
image=posterior,
header=image_header,
file_name=str(create_file_path(patient_results_folder, file_name)))
image_paths.append(image_path)
# save the binary mask
image_path = io_util.store_binary_mask_as_nifti(
image=binary,
header=image_header,
file_name=str(create_file_path(patient_results_folder,
class_name)))
image_paths.append(image_path)
# rescale and store uncertainty map as nifti
image_path = io_util.store_posteriors_as_nifti(
image=inference_result.uncertainty,
header=image_header,
file_name=str(create_file_path(patient_results_folder, "uncertainty")))
image_paths.append(image_path)
return image_paths
