def apply_rotation(self, mgrid):
if self.rot_mat is not None:
shape = mgrid[0].shape
rotated = self.rot_mat.dot(mgrid_to_points(mgrid).T).T
return points_to_mgrid(rotated, shape)
else:
return mgrid
def get_voxel_grid_real_space(images, append_ones=False):
# Get shape excluding channels
shape = images.shape[:-1]
# Get affine transforming voxel positions to real space positions
vox_to_real_affine = images.affine[:-1, :-1]
# Transform axes from voxel space to real space
grid_vox_space = np.mgrid[0:shape[0]:1,
0:shape[1]:1,
0:shape[2]:1]
# Move grid to real space
grid_points_real_space = vox_to_real_affine.dot(
mgrid_to_points(grid_vox_space).T).T
# Center
centered_grid_points_real_space = grid_points_real_space - \
np.mean(grid_points_real_space, axis=0)
# Append column of ones?
if append_ones:
centered_grid_points_real_space = np.column_stack(
(grid_points_real_space,
np.ones(len(grid_points_real_space))))
# Return real space grid as mgrid
points = points_to_mgrid(centered_grid_points_real_space, shape)
return points
def sample_plane_at(norm_vector, sample_dim, real_space_span,
offset_from_center, noise_sd, test_mode=False):
# Prepare normal vector to the plane
n_hat = np.array(norm_vector, np.float32)
n_hat /= np.linalg.norm(n_hat)
# Add noise?
if type(noise_sd) is not np.ndarray:
noise_sd = np.random.normal(scale=noise_sd, size=3)
n_hat += noise_sd
n_hat /= np.linalg.norm(n_hat)
if np.all(n_hat[:-1] < 0.2):
# Vector pointing primarily up, noise will have large effect on image
# orientation. We force the first two components to go into the
# positive direction to control variability of sampling
n_hat[:-1] = np.abs(n_hat[:-1])
if np.all(np.isclose(n_hat[:-1], 0)):
u = np.array([1, 0, 0])
v = np.array([0, 1, 0])
else:
# Find vector in same vertical plane as nhat
nhat_vs = n_hat.copy()
nhat_vs[-1] = nhat_vs[-1] + 1
nhat_vs /= np.linalg.norm(nhat_vs)
# Get two orthogonal vectors in plane, u pointing down in z-direction
u = get_rotation_matrix(np.cross(n_hat, nhat_vs), -90).dot(n_hat)
v = np.cross(n_hat, u)
# Define basis matrix + displacement to center (affine transformation)
basis = np.column_stack((u, v, n_hat))
# Define regular grid (centered at origin)
hd = real_space_span // 2
g = np.linspace(-hd, hd, sample_dim)
j = complex(sample_dim)
grid = np.mgrid[-hd:hd:j,
-hd:hd:j,
offset_from_center:offset_from_center:1j]
# Calculate voxel coordinates on the real space grid
points = mgrid_to_points(grid)
real_points = basis.dot(points.T).T
real_grid = points_to_mgrid(real_points, grid.shape[1:])
if test_mode:
return real_grid, g, np.linalg.inv(basis)
else:
return real_grid
def map_real_space_pred(pred, grid, inv_basis, voxel_grid_real_space, method="nearest"):
print("Mapping to real coordinate space...")
# Prepare fill value vector, we set this to 1.0 background
fill = np.zeros(shape=pred.shape[-1], dtype=np.float32)
fill[0] = 1.0
# Initialize interpolator object
intrp = RegularGridInterpolator(grid, pred, fill_value=fill,
bounds_error=False, method=method)
points = inv_basis.dot(mgrid_to_points(voxel_grid_real_space).T).T
transformed_grid = points_to_mgrid(points, voxel_grid_real_space[0].shape)
# Prepare mapped pred volume
mapped = np.empty(transformed_grid[0].shape + (pred.shape[-1],),
dtype=pred.dtype)
# Prepare interpolation function
def _do(xs, ys, zs, index):
return intrp((xs, ys, zs)), index
# Prepare thread pool of 10 workers
from concurrent.futures import ThreadPoolExecutor
from multiprocessing import cpu_count
pool = ThreadPoolExecutor(max_workers=max(7, cpu_count()))
# Perform interpolation async.
inds = np.arange(transformed_grid.shape[1])
result = pool.map(_do, transformed_grid[0], transformed_grid[1],
transformed_grid[2], inds)
i = 1
for map, ind in result:
# Print status
print(" %i/%i" % (i, inds[-1]+1), end="\r", flush=True)
i += 1
# Map the interpolation results into the volume
mapped[ind] = map
# Interpolate
# mapped = intrp(tuple(transformed_grid))
print("")
pool.shutdown()
return mapped
def sample_box_at(real_placement, sample_dim, real_box_dim,
noise_sd, test_mode):
j = complex(sample_dim)
a, b, c = real_placement
grid = np.mgrid[a:a + real_box_dim:j,
b:b + real_box_dim:j,
c:c + real_box_dim:j]
rot_mat = np.eye(3)
rot_grid = grid
if noise_sd:
# Get random rotation vector
rot_axis = get_random_views(N=1, dim=3, pos_z=True)
rot_angle = False
while not rot_angle:
angle = np.abs(np.random.normal(scale=noise_sd, size=1)[0])
if angle < 2*np.pi:
rot_angle = angle
rot_mat = get_rotation_matrix(rot_axis, angle_rad=rot_angle)
# Center --> apply rotation --> revert centering --> mgrid
points = mgrid_to_points(grid)
center = np.mean(points, axis=0)
points -= center
points = rot_mat.dot(points.T).T + center
rot_grid = points_to_mgrid(points, grid.shape[1:])
if test_mode:
axes = (np.linspace(a, a+real_box_dim, sample_dim),
np.linspace(b, b+real_box_dim, sample_dim),
np.linspace(c, c+real_box_dim, sample_dim))
return rot_grid, axes, np.linalg.inv(rot_mat)
else:
return rot_grid