from pathlib import Path import matplotlib.pyplot as plt import numpy import torch import cv2 from auto_detect_breast_mri.data.transforms import * from auto_detect_breast_mri.models.unet_crop import UNet_Crop from auto_detect_breast_mri.data.loaders import get_multiple_subjects_dataloader from auto_detect_breast_mri.config import resolve_path from auto_detect_breast_mri.data.metadata import get_uka_metatensor import unittest class TestAugmentation(unittest.TestCase): def setUp(self) -> None: # Parameters that can be adapted: self.batch_size = 1 # data locations come from the site config; point $AIMRI_CONFIG at the synthetic dataset # (tools/make_synthetic_dataset.py) to run this without any patient data feature_path = resolve_path(None, "metadata_file", "metadata export") path_base = resolve_path(None, "data_root", "root folder of the NIfTI data") split_files_folder = resolve_path(None, "split_root", "folder holding the split files") pretrained_model = str(Path(resolve_path(None, "model_root", "checkpoint folder")) / "unet_best.pth") protocol = ['pre'] pre_image_shape = (32, 512, 512) # transformations self.norm = ZNormalization() self.crop = UNet_Crop(input_channel=len(protocol), path_to_unet=pretrained_model) self.mirror = tio.RandomFlip((0,1,2), flip_probability=0.5) self.scale = tio.RandomAffine(scales=0.2, degrees=0, center='image', default_pad_value='minimum', check_shape=True) self.rotate = tio.RandomAffine(scales=0, degrees=25, center='image', default_pad_value='minimum', check_shape=True) transform = tio.Compose([ self.scale, self.mirror, self.crop, self.norm, ImageToTensor() ]) # Dataset and DataLoader feature_dataframe = get_uka_metatensor(0, feature_path) self.train_loader, self.eval_loader, self.test_loader = get_multiple_subjects_dataloader(path_base, feature_dataframe, pre_image_shape, protocol=protocol, set_folder=split_files_folder, transform=transform, batch_size=self.batch_size, segmentation_task=True, stratified=True) self.dataset = self.train_loader.dataset self.circle_image = numpy.zeros(pre_image_shape) self.circle_center = (112, 357) # col, row self.radius = 84 self.mask_width = 256 for d in range(pre_image_shape[0]): half = pre_image_shape[0] / 2 radius = self.radius * d / pre_image_shape[0] if d < half else self.radius * (pre_image_shape[0] - d) / \ pre_image_shape[0] radius = int(radius) cv2.circle(self.circle_image[d], self.circle_center, radius, (1, 1, 1), -1) self.probability_tensor = torch.ones((2,) + pre_image_shape) self.probability_tensor[1][self.circle_image == 1] = 2.2 self.probability_tensor[0][self.circle_image == 1] = 0.1 self.probability_tensor[1][self.circle_image != 1] = 0.3 self.probability_tensor[0][self.circle_image != 1] = 0.7 plt.imshow(self.circle_image[16], cmap='gray', interpolation='none') plt.show() def test_mask_extraction(self): height_position = self.circle_center[1] # radius is defined as x, y (x is horizontal) height_min = height_position - int(self.mask_width / 2) height_max = height_position + int(self.mask_width / 2) if height_min < 0: # cropping area would move out of image to the left height_min = 0 height_max = self.mask_width elif height_max > self.circle_image.shape[-1]: # cropping area would move out of image to the right height_min = self.mask_width height_max = self.circle_image.shape[-1] expected_output = torch.Tensor(self.circle_image)[:, height_min:height_max, :] plt.imshow(expected_output[16], cmap='gray', interpolation='none') plt.show() cropped_mask, (comp_min, comp_max) = self.crop.extract_breast_mask(self.probability_tensor, self.mask_width) calculated_output = cropped_mask[0] plt.imshow(cropped_mask[0, 16], cmap='gray', interpolation='none') plt.show() self.assertTrue(torch.all(calculated_output.eq(expected_output))) def test_augmentation_steps_visually(self): for batch_id, batch in enumerate(self.train_loader): data = batch['image']['data'] mask = batch['label']['data'] # image (one half) plt.imshow(data[0][0][21], cmap='gray', interpolation='none') plt.show() # mask plt.imshow(mask[0][1][21], cmap='gray', interpolation='none') plt.show() # Preprocessing: # 0. norm images_full_data to be between 0 and 1 preprocessed_input = self.norm.__call__(data[0]) # image (one half) plt.imshow(preprocessed_input[0][21], cmap='gray', interpolation='none') plt.show() # 1. crop images_full_data with unet cropped_image, (w_min, w_max), _ = self.crop.__call__(preprocessed_input) cropped_mask = mask[0, :, :, :, w_min:w_max] plt.imshow(torch.cat([cropped_image[0][21], cropped_mask[1][21]]), cmap='gray', interpolation='none') plt.show() transformation_tensor = torch.cat([cropped_image, cropped_mask], dim=0) # 2. Mirror randomly flipped_image = self.mirror.__call__(transformation_tensor) plt.imshow(torch.cat([flipped_image[0][21], flipped_image[4][21]]), cmap='gray', interpolation='none') plt.show() # 3. Random rotation between [-25, 25] rotated_image = self.rotate.__call__(flipped_image) plt.imshow(torch.cat([rotated_image[0][21], rotated_image[4][21]]), cmap='gray', interpolation='none') plt.show() # 4. Rescale between [0.9, 1.1] rescaled_image = self.scale.__call__(rotated_image) plt.imshow(torch.cat([rescaled_image[0][21], rescaled_image[4][21]]), cmap='gray', interpolation='none') plt.show() self.assertEqual(batch_id, len(self.train_loader)/self.batch_size) if __name__ == "__main__": unittest.main()