def plot_planes(save_dir, suffix, name, img_siz, trans_vec_final, quat_final,
mesh_final, trans_vec_gt, quat_gt, mesh_gt):
"""Plot GT and predicted planes
"""
img_c = (img_siz - 1) / 2
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
surf_gt = ax.plot_surface(mesh_gt[1], mesh_gt[0], mesh_gt[2])
surf_gt.set_facecolor([1, 0, 0, 0.3])
mat_gt = geometry.quaternion_matrix(quat_gt)
ax.quiver(trans_vec_gt[1] + img_c[1],
trans_vec_gt[0] + img_c[0],
trans_vec_gt[2] + img_c[2],
mat_gt[:, 1],
mat_gt[:, 0],
mat_gt[:, 2],
length=30,
color='r')
surf_final = ax.plot_surface(mesh_final[1], mesh_final[0], mesh_final[2])
surf_final.set_facecolor([0, 1, 0, 0.3])
mat_final = geometry.quaternion_matrix(quat_final)
ax.quiver(trans_vec_final[1] + img_c[1],
trans_vec_final[0] + img_c[0],
trans_vec_final[2] + img_c[2],
mat_final[:, 1],
mat_final[:, 0],
mat_final[:, 2],
length=30,
color='g')
plt.axis('equal')
ax.set_title('{0}'.format(name))
ax.set_xlim(0, img_siz[1])
ax.set_ylim(0, img_siz[0])
ax.set_zlim(0, img_siz[2])
ax.set_xlabel('y')
ax.set_ylabel('x')
ax.set_zlabel('z')
ax.view_init(elev=8., azim=63)
ax.invert_xaxis()
save_dir = os.path.join(save_dir, suffix)
if not os.path.isdir(save_dir):
os.makedirs(save_dir)
fig.savefig(os.path.join(save_dir, name + '.png'), bbox_inches='tight')
plt.close(fig)
def extract_plane_from_pose(image, t, q, plane_siz, order):
"""Extract a 2D plane image from the 3D volume given the pose wrt the identity plane.
Args:
image: 3D volume. [x,y,z]
t: translation of the pose [3]
q: rotation of the pose in quaternions [4]
plane_siz: size of plane [2]
order: interpolation order (0-5)
Returns:
slice: 2D plane image [plane_siz[0], plane_siz[1]]
mesh: mesh coordinates of the plane. Origin at volume corner. numpy array of size [3, plane_siz[0], plane_siz[1]]
"""
# Initialise identity plane
xyz_coords = init_mesh(plane_siz)
# Rotate and translate plane
mat = geometry.quaternion_matrix(q)
mat[:3, 3] = t
xyz_coords = np.dot(mat, xyz_coords)
# Extract image plane
slice, xyz_coords_new = extract_plane_from_mesh(image, xyz_coords, plane_siz, order)
return slice, xyz_coords_new
def save_planes_tform(save_dir, suffix, names, trans_vecs, quats):
"""Save the plane as a 4x4 transformation matrix in a txt file.
Args:
save_dir: Directory storing the results.
suffix: 'train' or 'test'
names: list of names of the patients.
trans_vecs: translation vectors. [img_count, 3]
quats: quaternions. [img_count, 4]
"""
save_dir = os.path.join(save_dir, suffix)
if not os.path.isdir(save_dir):
os.makedirs(save_dir)
for i in xrange(len(names)):
mat = geometry.quaternion_matrix(quats[i, :])
mat[:3, 3] = trans_vecs[i, :]
np.savetxt(os.path.join(save_dir, names[i] + '_mat.txt'), mat)
def update_plane(n):
ax1.cla()
fig1.set_tight_layout(True)
# GT plane
surf_gt = ax1.plot_surface(mesh_gt[1], mesh_gt[0], mesh_gt[2])
surf_gt.set_facecolor([1, 0, 0, 0.3])
mat_gt = geometry.quaternion_matrix(quat_gt)
quiv_gt = ax1.quiver(trans_vec_gt[1] + img_c[1],
trans_vec_gt[0] + img_c[0],
trans_vec_gt[2] + img_c[2],
mat_gt[:, 1],
mat_gt[:, 0],
mat_gt[:, 2],
length=30,
color='r')
# Predicted planes
surf = ax1.plot_surface(meshes[0, n, 0, :, :, 1],
meshes[0, n, 0, :, :, 0], meshes[0, n, 0, :, :,
2])
surf.set_facecolor([0, 1, 0, 0.3])
quiv = ax1.quiver(matrices[0, n, 1, 3] + img_c[1],
matrices[0, n, 0, 3] + img_c[0],
matrices[0, n, 2, 3] + img_c[2],
matrices[0, n, :, 1],
matrices[0, n, :, 0],
matrices[0, n, :, 2],
length=30,
color='g')
plt.axis('equal')
ax1.set_title('iteration {}'.format(n))
ax1.set_xlim(0, img_siz[1])
ax1.set_ylim(0, img_siz[0])
ax1.set_zlim(0, img_siz[2])
ax1.set_xlabel('y')
ax1.set_ylabel('x')
ax1.set_zlabel('z')
ax1.view_init(elev=8., azim=63)
ax1.invert_xaxis()
return surf_gt, quiv_gt, surf, quiv
def get_train_pairs(config, data):
"""Prepare training data.
Args:
batch_size: mini batch size
images: list of img_count images. Each image is [width, height, depth, channel], [x,y,z,channel]
trans_gt: 3D centre point of the ground truth plane wrt the volume centre as origin. [2, img_count, 3]. first dimension is the tv(0) or tc(1) plane
rots_gt: Quaternions that rotate xy-plane to the GT plane. [2, img_count, 4]. first dimension is the tv(0) or tc(1) plane
trans_frac: Percentage of middle volume to sample translation vector from. (0-1)
max_euler: Maximum range of Euler angles to sample from. (+/- max_euler). [3]
box_size: size of 2D plane. [x,y].
input_plane: number of input planes (1 or 3)
plane: TV(0) or TC(1)
Returns:
slices: 2D plane images. [batch_size, box_size[0], box_size[1], input_plane]
actions_tran: [batch_size, 6] the GT classification probability for translation. Hard label, one-hot vector. Gives the axis about which to translate, ie. axis with biggest distance to GT
trans_diff: [batch_size, 3]. 3D centre point of the ground truth plane wrt the centre of the randomly sampled plane as origin.
actions_rot:[batch_size, 6] the GT classification probability for rotation. Hard label, one-hot vector. Gives the axis about which to rotate, ie. rotation axis with biggest rotation angle.
rots_diff: [batch_size, 4]. Rotation that maps the randomly sampled plane to the GT plane.
"""
images = data.images
trans_gt = data.trans_vecs
rots_gt = data.quats
batch_size = config.batch_size
box_size = config.box_size
input_plane = config.input_plane
trans_frac = config.trans_frac
max_euler = config.max_euler
img_count = len(images)
slices = np.zeros((batch_size, box_size[0], box_size[1], input_plane),
np.float32)
trans_diff = np.zeros((batch_size, 3))
trans = np.zeros((batch_size, 3))
rots_diff = np.zeros((batch_size, 4))
rots = np.zeros((batch_size, 4))
euler = np.zeros(
(batch_size, 6, 3)
) # 6 Euler angle conventions. 'sxyz', 'sxzy', 'syxz', 'syzx', 'szxy', 'szyx'
actions_tran = np.zeros((batch_size, 6), np.float32)
actions_rot = np.zeros((batch_size, 6), np.float32)
# get image indices randomly for a mini-batch
ind = np.random.randint(img_count, size=batch_size)
# Random uniform sampling of Euler angles with restricted range
euler_angles = geometry.sample_euler_angles_fix_range(
batch_size, max_euler[0], max_euler[1], max_euler[2])
for i in xrange(batch_size):
image = np.squeeze(images[ind[i]])
img_siz = np.array(image.shape)
# GT translation and quaternions
tran_gt = trans_gt[ind[i], :]
rot_gt = rots_gt[ind[i], :]
# Randomly sample translation (plane centre) and quaternions
tran = (np.random.rand(3) * (img_siz * trans_frac) + img_siz *
(1 - trans_frac) / 2.0) - ((img_siz - 1) / 2.0)
trans[i, :] = tran
rot = geometry.quaternion_from_euler(euler_angles[i, 0],
euler_angles[i, 1],
euler_angles[i, 2], 'rxyz')
rots[i, :] = rot
##### Extract plane image #####
# Initialise identity plane and get orthogonal planes
if input_plane == 1:
xyz_coords = plane.init_mesh_by_plane(box_size, 'z')
elif input_plane == 3:
xyz_coords = plane.init_mesh_ortho(box_size)
# Rotate and translate plane
mat = geometry.quaternion_matrix(rot)
mat[:3, 3] = tran
xyz_coords = np.matmul(mat, xyz_coords)
# Extract image plane
if input_plane == 1:
slices[i, :, :, 0], _ = plane.extract_plane_from_mesh(
image, xyz_coords, box_size, 1)
elif input_plane == 3:
slice_single, _ = plane.extract_plane_from_mesh_batch(
image, xyz_coords, box_size, 1)
slices[i] = np.transpose(slice_single, (1, 2, 0))
##### Compute GT labels #####
# Translation and rotation regression outputs. Compute difference in tran and quat between sampled plane and GT plane (convert to rotation matrices first)
mat_inv = geometry.inv_mat(mat)
mat_gt = geometry.quaternion_matrix(rot_gt)
mat_gt[:3, 3] = tran_gt
mat_diff = np.matmul(mat_inv, mat_gt)
trans_diff[i, :] = mat_diff[:3, 3]
rots_diff[i, :] = geometry.quaternion_from_matrix(mat_diff,
isprecise=True)
# Rotation classification output. Compute Euler angles for the six different conventions
euler[i, 0, :] = np.array(
geometry.euler_from_matrix(mat_diff, axes='sxyz'))
euler[i, 1, :] = np.array(
geometry.euler_from_matrix(mat_diff, axes='sxzy'))
euler[i, 2, :] = np.array(
geometry.euler_from_matrix(mat_diff, axes='syxz'))
euler[i, 3, :] = np.array(
geometry.euler_from_matrix(mat_diff, axes='syzx'))
euler[i, 4, :] = np.array(
geometry.euler_from_matrix(mat_diff, axes='szxy'))
euler[i, 5, :] = np.array(
geometry.euler_from_matrix(mat_diff, axes='szyx'))
# Rotation classification output.
max_ind_rot = np.argmax(np.abs(euler[:, :, 0]), axis=1)
rot_x_max = np.logical_or(max_ind_rot == 0, max_ind_rot == 1)
rot_y_max = np.logical_or(max_ind_rot == 2, max_ind_rot == 3)
rot_z_max = np.logical_or(max_ind_rot == 4, max_ind_rot == 5)
actions_ind_rot = np.zeros((batch_size), dtype=np.uint16)
actions_ind_rot[rot_x_max] = 0
actions_ind_rot[rot_y_max] = 1
actions_ind_rot[rot_z_max] = 2
max_euler = euler[np.arange(batch_size), max_ind_rot,
np.zeros(batch_size, dtype=np.uint16)] # [batch_size]
is_positive = (max_euler > 0)
actions_ind_rot[is_positive] = actions_ind_rot[is_positive] * 2
actions_ind_rot[np.logical_not(
is_positive)] = actions_ind_rot[np.logical_not(is_positive)] * 2 + 1
actions_rot[np.arange(batch_size), actions_ind_rot] = 1
# Translation classification output
max_ind_tran = np.argmax(np.abs(trans_diff), axis=1) # [batch_size]
max_trans_diff = trans_diff[np.arange(batch_size),
max_ind_tran] # [batch_size]
is_positive = (max_trans_diff > 0)
actions_ind_tran = np.zeros((batch_size), dtype=np.uint16)
actions_ind_tran[is_positive] = max_ind_tran[is_positive] * 2
actions_ind_tran[np.logical_not(
is_positive)] = max_ind_tran[np.logical_not(is_positive)] * 2 + 1
actions_tran[np.arange(batch_size), actions_ind_tran] = 1
return slices, actions_tran, trans_diff, actions_rot, rots_diff
def predict_mat_diff(ytc=None,
ytr=None,
yrc=None,
yrr=None,
weight_tran=False,
weight_rot=False):
"""Form predicted transformation matrix from translation vector and quaternions.
Use classification probabilities as weighting if weight_tran, weight_rot set to True
Args:
ytc: predicted translation probabilities [num_examples, 6]
ytr: predicted translation vector [num_examples, 3].
yrc: predicted rotation probabilities [num_examples, 6]
yrr: predicted rotation quaternions [num_examples, 4].
weight_tran: False: Regression only.
True: Soft classification. Multiply classification probabilities with regressed distances.
weight_rot: False: Regression only.
True: Soft classification. Multiply classification probabilities with regressed rotation.
Returns:
mat_diff: Predicted transformation matrix [num_examples, 4, 4]
"""
num_examples = yrr.shape[0]
mat_diff = np.zeros((num_examples, 4, 4))
# Rotation prediction
if (not weight_rot) or (yrc is None):
# Regression only.
for j in xrange(num_examples):
mat_diff[j] = geometry.quaternion_matrix(yrr[j])
else:
# Multiply classification probabilities with regressed rotation.
yrc_max = np.amax(yrc, axis=1) # yrc_max=[num_examples]
rot_axis = (np.argmax(yrc, axis=1) / 2).astype(
int) # rot_axis=[num_examples]
# Convert predicted quaternions to euler angles using the most probable convention
# Then convert to rotation matrix
for j in xrange(num_examples):
if rot_axis[j] == 0:
euler = np.array(
geometry.euler_from_quaternion(yrr[j, :], axes='sxyz'))
euler[1:3] = 0
euler[0] = euler[0] * yrc_max[j]
mat_diff[j] = geometry.euler_matrix(euler[0],
euler[1],
euler[2],
axes='sxyz')
elif rot_axis[j] == 1:
euler = np.array(
geometry.euler_from_quaternion(yrr[j, :], axes='syxz'))
euler[1:3] = 0
euler[0] = euler[0] * yrc_max[j]
mat_diff[j] = geometry.euler_matrix(euler[0],
euler[1],
euler[2],
axes='syxz')
elif rot_axis[j] == 2:
euler = np.array(
geometry.euler_from_quaternion(yrr[j, :], axes='szxy'))
euler[1:3] = 0
euler[0] = euler[0] * yrc_max[j]
mat_diff[j] = geometry.euler_matrix(euler[0],
euler[1],
euler[2],
axes='szxy')
# Translation prediction
if (not weight_tran) or (ytc is None):
# Regression only.
mat_diff[:, :3, 3] = ytr
else:
# Soft classification. Multiply classification probabilities with regressed distances.
mat_diff[:, :3, 3] = ytr * np.amax(np.reshape(ytc,
(ytc.shape[0], 3, 2)),
axis=2)
return mat_diff