| 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]) |
|
|
| |
| 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): |
| |
| 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 |
|
|
| sub_edge = [edge[0]] |
| include_last = bool((edge_len - 2) % num_samples) |
| 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 |
|
|
|
|
