def test_show_non_square_images(test_output_dirs: OutputFolderForTests) -> None: input_file = full_ml_test_data_path("patch_sampling") / "scan_small.nii.gz" input = load_nifti_image(input_file) image = input.image shape = image.shape mask = np.zeros_like(image) mask[shape[0] // 2, shape[1] // 2, shape[2] // 2] = 1 for dim in range(3): scan_with_transparent_overlay(image, mask, dim, shape[dim] // 2, spacing=input.header.spacing) actual_file = Path(test_output_dirs.root_dir) / f"dim_{dim}.png" resize_and_save(5, 5, actual_file) expected = full_ml_test_data_path("patch_sampling") / f"overlay_with_aspect_dim{dim}.png" # To update the stored results, uncomment this line: # expected.write_bytes(actual_file.read_bytes()) assert_binary_files_match(actual_file, expected)
def test_plot_overlay(test_output_dirs: OutputFolderForTests, dimension: int) -> None: set_random_seed(0) shape = (10, 30, 30) image = np.random.rand(*shape).astype(np.float32) * 1000 mask = np.zeros(shape).flatten() for i in range(len(mask)): mask[i] = i mask = mask.reshape(shape) plt.figure() scan_with_transparent_overlay(image, mask, dimension, shape[dimension] // 2, spacing=(1.0, 1.0, 1.0)) file = Path(test_output_dirs.root_dir) / "plot.png" resize_and_save(5, 5, file) assert file.exists() expected = full_ml_test_data_path("patch_sampling") / f"overlay_{dimension}.png" # To update the stored results, uncomment this line: # expected.write_bytes(file.read_bytes()) assert_binary_files_match(file, expected)
def visualize_random_crops(sample: Sample, config: SegmentationModelBase, output_folder: Path) -> np.ndarray: """ Simulate the effect of sampling random crops (as is done for trainig segmentation models), and store the results as a Nifti heatmap and as 3 axial/sagittal/coronal slices. The heatmap and the slices are stored in the given output folder, with filenames that contain the patient ID as the prefix. :param sample: The patient information from the dataset, with scans and ground truth labels. :param config: The model configuration. :param output_folder: The folder into which the heatmap and thumbnails should be written. :return: A numpy array that has the same size as the image, containing how often each voxel was contained in """ output_folder.mkdir(exist_ok=True, parents=True) sample = CroppingDataset.create_possibly_padded_sample_for_cropping( sample=sample, crop_size=config.crop_size, padding_mode=config.padding_mode) logging.info(f"Processing sample: {sample.patient_id}") # Exhaustively sample with random crop function image_channel0 = sample.image[0] heatmap = np.zeros(image_channel0.shape, dtype=np.uint16) # Number of repeats should fit into the range of UInt16, because we will later save the heatmap as an integer # Nifti file of that datatype. repeats = 200 for _ in range(repeats): slicers, _ = augmentation.slicers_for_random_crop( sample=sample, crop_size=config.crop_size, class_weights=config.class_weights) heatmap[slicers[0], slicers[1], slicers[2]] += 1 is_3dim = heatmap.shape[0] > 1 header = sample.metadata.image_header if not header: logging.warning( f"No image header found for patient {sample.patient_id}. Using default header." ) header = get_unit_image_header() if is_3dim: ct_output_name = str(output_folder / f"{sample.patient_id}_ct.nii.gz") heatmap_output_name = str( output_folder / f"{sample.patient_id}_sampled_patches.nii.gz") io_util.store_as_nifti(image=heatmap, header=header, file_name=heatmap_output_name, image_type=heatmap.dtype, scale=False) io_util.store_as_nifti(image=image_channel0, header=header, file_name=ct_output_name, image_type=sample.image.dtype, scale=False) heatmap_scaled = heatmap.astype(dtype=np.float) / heatmap.max() # If the incoming image is effectively a 2D image with degenerate Z dimension, then only plot a single # axial thumbnail. Otherwise, plot thumbnails for all 3 dimensions. dimensions = list(range(3)) if is_3dim else [0] # Center the 3 thumbnails at one of the points where the heatmap attains a maximum. This should ensure that # the thumbnails are in an area where many of the organs of interest are located. max_heatmap_index = np.unravel_index( heatmap.argmax(), heatmap.shape) if is_3dim else (0, 0, 0) for dimension in dimensions: plt.clf() scan_with_transparent_overlay( scan=image_channel0, overlay=heatmap_scaled, dimension=dimension, position=max_heatmap_index[dimension] if is_3dim else 0, spacing=header.spacing) # Construct a filename that has a dimension suffix if we are generating 3 of them. For 2dim images, skip # the suffix. thumbnail = f"{sample.patient_id}_sampled_patches" if is_3dim: thumbnail += f"_dim{dimension}" thumbnail += ".png" resize_and_save(width_inch=5, height_inch=5, filename=output_folder / thumbnail) return heatmap