File size: 3,683 Bytes
b10768a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import networkx as nx
import numpy as np
from scipy.interpolate import RegularGridInterpolator, LinearNDInterpolator
def graph_to_ndarray(graph):
out_nodes = np.empty((1, 3))
out_edges = np.empty((1, 3))
visited_nodes = list()
visited_node_pairs = list()
for (start_node, end_node) in graph.edges():
if (not (start_node, end_node) in visited_node_pairs) and (not (end_node, start_node) in visited_node_pairs):
edge = graph[start_node][end_node]['pts']
out_edges = np.vstack([out_edges, edge])
# Avoid duplicates
if not (start_node in visited_nodes):
out_nodes = np.vstack([out_nodes, graph.nodes[start_node]['o']])
visited_nodes.append(start_node)
if not (end_node in visited_nodes):
out_nodes = np.vstack([out_nodes, graph.nodes[end_node]['o']])
visited_nodes.append(end_node)
visited_node_pairs.append((start_node, end_node))
return np.vstack([out_edges, out_nodes]), out_nodes, out_edges
def get_bifurcation_nodes(graph: nx.Graph):
# Vertex degree relates to the number of branches connected to a given node
out_nodes = np.empty((1, 3))
bif_nodes_id = list()
for node_num, deg in graph.degree:
if deg > 1:
bif_nodes_id.append(node_num)
out_nodes = np.vstack([out_nodes, graph.nodes[node_num]['o']])
return out_nodes, bif_nodes_id
def apply_displacement(pts_list: np.ndarray, interpolator: [RegularGridInterpolator, LinearNDInterpolator]):
pts_list = pts_list.astype(np.float)
ret_val = pts_list + interpolator(pts_list).squeeze()
return ret_val
def deform_graph(graph, dm_interpolator: [RegularGridInterpolator, LinearNDInterpolator]):
def_graph = nx.Graph()
for (start_node, end_node) in graph.edges():
edge = graph[start_node][end_node]['pts']
def_edge = apply_displacement(edge, dm_interpolator)
def_start_node_pts = apply_displacement(graph.nodes[start_node]['pts'], dm_interpolator)
def_end_node_pts = apply_displacement(graph.nodes[end_node]['pts'], dm_interpolator)
def_start_node_o = apply_displacement(graph.nodes[start_node]['o'], dm_interpolator)
def_end_node_o = apply_displacement(graph.nodes[end_node]['o'], dm_interpolator)
def_graph.add_node(start_node, pts=def_start_node_pts, o=def_start_node_o)
def_graph.add_node(end_node, pts=def_end_node_pts, o=def_end_node_o)
def_graph.add_edge(start_node, end_node, pts=def_edge, weight=len(def_edge))
return def_graph
def subsample_graph(graph: nx.Graph, num_samples=3):
sub_graph = nx.Graph()
for (start_node, end_node) in graph.edges():
edge = graph[start_node][end_node]['pts']
edge_len = edge.shape[0]
sub_edge_len = (edge_len - 2) // num_samples # Do not count the pts corresponding to the nodes (-2)
sub_edge = [edge[0]]
include_last = bool((edge_len - 2) % num_samples) # Skip the last point, as this is too close to the node
if sub_edge_len:
idxs = np.arange(0, edge_len, num_samples)[1:] if include_last else np.arange(0, edge_len, num_samples)[1:-1]
for i in idxs:
sub_edge.append(edge[i])
sub_edge.append(edge[-1])
sub_edge = np.asarray(sub_edge)
sub_graph.add_node(start_node, pts=graph.nodes[start_node]['pts'], o=graph.nodes[start_node]['o'])
sub_graph.add_node(end_node, pts=graph.nodes[end_node]['pts'], o=graph.nodes[end_node]['o'])
sub_graph.add_edge(start_node, end_node, pts=sub_edge, weight=len(sub_edge))
return sub_graph
|