import nibabel as nib
import numpy as np
from unittest import TestCase
from fetal_net.utils.patches import compute_patch_indices, get_patch_from_3d_data, reconstruct_from_patches
class TestPrediction(TestCase):
def setUp(self):
image_shape = (120, 144, 90)
data = np.arange(0, image_shape[0]*image_shape[1]*image_shape[2]).reshape(image_shape)
affine = np.diag(np.ones(4))
self.image = nib.Nifti1Image(data, affine)
def test_reconstruct_from_patches(self):
patch_shape = (32, 32, 32)
patch_overlap = 0
patch_indices = compute_patch_indices(self.image.shape, patch_shape, patch_overlap)
patches = [get_patch_from_3d_data(self.image.get_data(), patch_shape, index) for index in patch_indices]
reconstruced_data = reconstruct_from_patches(patches, patch_indices, self.image.shape)
# noinspection PyTypeChecker
self.assertTrue(np.all(self.image.get_data() == reconstruced_data))
def test_reconstruct_with_overlapping_patches(self):
patch_overlap = 0
patch_shape = (32, 32, 32)
patch_indices = compute_patch_indices(self.image.shape, patch_shape, patch_overlap)
patches = [get_patch_from_3d_data(self.image.get_data(), patch_shape, index) for index in patch_indices]
# extend patches with modified patches that are 2 lower than the original patches
patches.extend([patch - 2 for patch in patches])
patch_indices = np.concatenate([patch_indices, patch_indices], axis=0)
reconstruced_data = reconstruct_from_patches(patches, patch_indices, self.image.shape)
# The reconstructed data should be 1 lower than the original data as 2 was subtracted from half the patches.
# The resulting reconstruction should be the average.
# noinspection PyTypeChecker
self.assertTrue(np.all((self.image.get_data() - 1) == reconstruced_data))
def test_reconstruct_with_overlapping_patches2(self):
image_shape = (144, 144, 144)
data = np.arange(0, image_shape[0]*image_shape[1]*image_shape[2]).reshape(image_shape)
patch_overlap = 16
patch_shape = (64, 64, 64)
patch_indices = compute_patch_indices(data.shape, patch_shape, patch_overlap)
patches = [get_patch_from_3d_data(data, patch_shape, index) for index in patch_indices]
no_overlap_indices = compute_patch_indices(data.shape, patch_shape, 32)
patch_indices = np.concatenate([patch_indices, no_overlap_indices])
patches.extend([get_patch_from_3d_data(data, patch_shape, index) for index in no_overlap_indices])
reconstruced_data = reconstruct_from_patches(patches, patch_indices, data.shape)
# noinspection PyTypeChecker
self.assertTrue(np.all(data == reconstruced_data))
def test_reconstruct_with_multiple_channels(self):
image_shape = (144, 144, 144)
n_channels = 4
data = np.arange(0, image_shape[0]*image_shape[1]*image_shape[2]*n_channels).reshape(
[n_channels] + list(image_shape))
patch_overlap = 16
patch_shape = (64, 64, 64)
patch_indices = compute_patch_indices(image_shape, patch_shape, patch_overlap)
patches = [get_patch_from_3d_data(data, patch_shape, index) for index in patch_indices]
self.assertEqual(patches[0].shape, tuple([4] + list(patch_shape)))
reconstruced_data = reconstruct_from_patches(patches, patch_indices, data.shape)
# noinspection PyTypeChecker
self.assertTrue(np.all(data == reconstruced_data))
