AI_MRI / tests /test_augmentation.py

init

a23d562 2 months ago

6.97 kB

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()

	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()