|
|
import matplotlib.pyplot as plt |
|
|
from mpl_toolkits.mplot3d import Axes3D |
|
|
from matplotlib.lines import Line2D |
|
|
import numpy as np |
|
|
from ddmr.utils.visualization import add_axes_arrows_3d, remove_tick_labels, set_axes_size |
|
|
import os |
|
|
|
|
|
|
|
|
def _plot_graph(graph, ax, nodes_colour='C3', edges_colour='C1', plot_nodes=True, plot_edges=True, add_axes=True): |
|
|
if plot_edges: |
|
|
for (start_node, end_node) in graph.edges(): |
|
|
edge_pts = graph[start_node][end_node]['pts'] |
|
|
edge_pts = np.vstack([graph.nodes[start_node]['o'], edge_pts]) |
|
|
edge_pts = np.vstack([edge_pts, graph.nodes[end_node]['o']]) |
|
|
ax.plot(edge_pts[:, 0], edge_pts[:, 1], edge_pts[:, 2], edges_colour) |
|
|
|
|
|
if plot_nodes: |
|
|
nodes = graph.nodes() |
|
|
ps = np.array([nodes[i]['o'] for i in nodes]) |
|
|
if len(ps.shape) > 1: |
|
|
ax.scatter(ps[:, 0], ps[:, 1], ps[:, 2], nodes_colour) |
|
|
else: |
|
|
ax.scatter(ps[0], ps[1], ps[2], nodes_colour) |
|
|
ax.set_xlim(0, 63) |
|
|
ax.set_ylim(0, 63) |
|
|
ax.set_zlim(0, 63) |
|
|
remove_tick_labels(ax, True) |
|
|
if add_axes: |
|
|
add_axes_arrows_3d(ax, x_color='r', y_color='g', z_color='b') |
|
|
ax.view_init(None, 45) |
|
|
|
|
|
return ax |
|
|
|
|
|
|
|
|
def plot_skeleton(img, skeleton, graph, filename='skeleton', extension=['.png']): |
|
|
if not isinstance(extension, list): |
|
|
extension = [extension] |
|
|
|
|
|
f = plt.figure(figsize=(5, 5)) |
|
|
ax = f.add_subplot(111, projection='3d') |
|
|
|
|
|
coords = np.argwhere(skeleton) |
|
|
i = coords[:, 0] |
|
|
j = coords[:, 1] |
|
|
k = coords[:, 2] |
|
|
|
|
|
seg = ax.voxels(img, facecolors=(0., 0., 1., 0.3), label='image') |
|
|
ske = ax.scatter(i, j, k, color='C1', label='skeleton', s=1) |
|
|
ax.set_xlim(0, 63) |
|
|
ax.set_ylim(0, 63) |
|
|
ax.set_zlim(0, 63) |
|
|
remove_tick_labels(ax, True) |
|
|
add_axes_arrows_3d(ax, x_color='r', y_color='g', z_color='b') |
|
|
ax.view_init(None, 45) |
|
|
for ex in extension: |
|
|
f.savefig(filename + '_segmentation_skeleton' + ex) |
|
|
|
|
|
|
|
|
ax = _plot_graph(graph, ax, 'r', 'r') |
|
|
|
|
|
for ex in extension: |
|
|
f.savefig(filename + '_combined' + ex) |
|
|
plt.close() |
|
|
|
|
|
|
|
|
f = plt.figure(figsize=(5, 5)) |
|
|
ax = f.add_subplot(111, projection='3d') |
|
|
|
|
|
ax = _plot_graph(graph, ax) |
|
|
|
|
|
for ex in extension: |
|
|
f.savefig(filename + '_graph' + ex) |
|
|
plt.close() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def compare_graphs(graph_0, graph_1, graph_names=None, filename='compare_graphs'): |
|
|
f = plt.figure(figsize=(12, 5)) |
|
|
if graph_names is None: |
|
|
graph_names =['graph_0', 'graph_1'] |
|
|
else: |
|
|
assert len(graph_names) == 2, 'A different name is expected for each graph' |
|
|
ax = f.add_subplot(131, projection='3d') |
|
|
ax = _plot_graph(graph_0, ax) |
|
|
ax.set_title(graph_names[0], y=-0.01) |
|
|
|
|
|
ax = f.add_subplot(132, projection='3d') |
|
|
ax = _plot_graph(graph_1, ax) |
|
|
ax.set_title(graph_names[1]) |
|
|
|
|
|
ax = f.add_subplot(133, projection='3d') |
|
|
ax = _plot_graph(graph_0, ax, 'C2', 'C2', plot_nodes=False) |
|
|
ax = _plot_graph(graph_1, ax, 'C4', 'C4', plot_nodes=False) |
|
|
legend_elements = [Line2D([0], [0], color='C2', lw=2, label=graph_names[0]), |
|
|
Line2D([0], [0], color='C4', lw=2, label=graph_names[1])] |
|
|
ax.legend(handles=legend_elements) |
|
|
|
|
|
f.savefig(filename + '_compare_graphs.png') |
|
|
plt.close() |
|
|
|
|
|
|
|
|
def plot_cpd_registration_step(iteration, error, X, Y, out_folder, add_axes=True, pdf=True): |
|
|
fig = plt.figure(figsize=(8, 8)) |
|
|
ax = fig.add_axes([0, 0, .9, .9], projection='3d') |
|
|
ax.scatter(X[:, 0], X[:, 1], X[:, 2], color='C1', label='Fixed') |
|
|
ax.scatter(Y[:, 0], Y[:, 1], Y[:, 2], color='C2', label='Moving') |
|
|
|
|
|
ax.text2D(0.95, 0.98, 'Iteration: {:d}'.format( |
|
|
iteration), horizontalalignment='right', verticalalignment='center', transform=ax.transAxes, fontsize='x-large') |
|
|
|
|
|
|
|
|
ax.legend(loc='upper left', fontsize='x-large') |
|
|
|
|
|
if add_axes: |
|
|
x_range = [np.min(np.hstack([X[:, 0], Y[:, 0]])), np.max(np.hstack([X[:, 0], Y[:, 0]]))] |
|
|
y_range = [np.min(np.hstack([X[:, 1], Y[:, 1]])), np.max(np.hstack([X[:, 1], Y[:, 1]]))] |
|
|
z_range = [np.min(np.hstack([X[:, 2], Y[:, 2]])), np.max(np.hstack([X[:, 2], Y[:, 2]]))] |
|
|
ax.set_xlim(x_range[0], x_range[1]) |
|
|
ax.set_ylim(y_range[0], y_range[1]) |
|
|
ax.set_zlim(z_range[0], z_range[1]) |
|
|
|
|
|
remove_tick_labels(ax, True) |
|
|
add_axes_arrows_3d(ax, x_color='r', y_color='g', z_color='b', arrow_length=25, dist_arrow_text=3) |
|
|
ax.view_init(None, 45) |
|
|
|
|
|
os.makedirs(out_folder, exist_ok=True) |
|
|
fig.savefig(os.path.join(out_folder, '{:04d}.png'.format(iteration))) |
|
|
if pdf: |
|
|
fig.savefig(os.path.join(out_folder, '{:04d}.pdf'.format(iteration))) |
|
|
plt.close() |
|
|
|
|
|
|
|
|
def plot_cpd(fix_pts, mov_pts, fix_centroid, mov_centroid, file_name): |
|
|
fig = plt.figure(figsize=(8, 8)) |
|
|
ax = fig.add_axes([0, 0, .9, .9], projection='3d') |
|
|
ax.scatter(fix_pts[:, 0], fix_pts[:, 1], fix_pts[:, 2], color='C1', label='Fixed') |
|
|
ax.scatter(mov_pts[:, 0], mov_pts[:, 1], mov_pts[:, 2], color='C2', label='Moving') |
|
|
ax.scatter(fix_centroid[0], fix_centroid[1], fix_centroid[2], color='none', s=100, edgecolor='b', label='Centroid') |
|
|
ax.scatter(mov_centroid[0], mov_centroid[1], mov_centroid[2], color='none', s=100, edgecolor='b') |
|
|
ax.scatter(fix_centroid[0], fix_centroid[1], fix_centroid[2], color='C1') |
|
|
ax.scatter(mov_centroid[0], mov_centroid[1], mov_centroid[2], color='C2') |
|
|
|
|
|
x_range = [np.min(np.hstack([fix_pts[:, 0], mov_pts[:, 0], fix_centroid[0], mov_centroid[0]])), |
|
|
np.max(np.hstack([fix_pts[:, 0], mov_pts[:, 0], fix_centroid[0], mov_centroid[0]]))] |
|
|
y_range = [np.min(np.hstack([fix_pts[:, 1], mov_pts[:, 1], fix_centroid[1], mov_centroid[1]])), |
|
|
np.max(np.hstack([fix_pts[:, 1], mov_pts[:, 1], fix_centroid[1], mov_centroid[1]]))] |
|
|
z_range = [np.min(np.hstack([fix_pts[:, 2], mov_pts[:, 2], fix_centroid[2], mov_centroid[2]])), |
|
|
np.max(np.hstack([fix_pts[:, 2], mov_pts[:, 2], fix_centroid[2], mov_centroid[2]]))] |
|
|
ax.set_xlim(x_range[0], x_range[1]) |
|
|
ax.set_ylim(y_range[0], y_range[1]) |
|
|
ax.set_zlim(z_range[0], z_range[1]) |
|
|
|
|
|
remove_tick_labels(ax, True) |
|
|
add_axes_arrows_3d(ax, x_color='r', y_color='g', z_color='b', arrow_length=25, dist_arrow_text=3) |
|
|
ax.view_init(None, 45) |
|
|
ax.legend(fontsize='xx-large') |
|
|
fig.savefig(file_name + '.png') |
|
|
fig.savefig(file_name + '.pdf') |
|
|
plt.close() |
|
|
|
|
|
|
