From c014a27a7c14408b4996a2916296bb95fd3ab0ab Mon Sep 17 00:00:00 2001 From: Anna Grim <108307071+anna-grim@users.noreply.github.com> Date: Sun, 23 Jun 2024 13:09:29 -0700 Subject: [PATCH] feat: long range proposals, profiles (#170) Co-authored-by: anna-grim --- src/deep_neurographs/generate_proposals.py | 54 ++++++++++++++++--- src/deep_neurographs/geometry.py | 30 ++++++----- .../machine_learning/feature_generation.py | 28 +++++++--- .../machine_learning/graph_trainer.py | 3 +- src/deep_neurographs/neurograph.py | 6 ++- 5 files changed, 89 insertions(+), 32 deletions(-) diff --git a/src/deep_neurographs/generate_proposals.py b/src/deep_neurographs/generate_proposals.py index 1fd4c86..d66dbff 100644 --- a/src/deep_neurographs/generate_proposals.py +++ b/src/deep_neurographs/generate_proposals.py @@ -15,7 +15,9 @@ ENDPOINT_DIST = 10 -def run(neurograph, search_radius, complex_bool=True): +def run( + neurograph, search_radius, complex_bool=True, long_range_proposals=False +): """ Generates proposals emanating from "leaf". @@ -28,6 +30,9 @@ def run(neurograph, search_radius, complex_bool=True): complex_bool : bool, optional Indication of whether to generate complex proposals. The default is True. + long_range_proposals : bool + Indication of whether to generate simple proposals within distance of + 2 * search_radius of leaf. Returns ------- @@ -38,13 +43,24 @@ def run(neurograph, search_radius, complex_bool=True): connections = dict() for leaf in neurograph.leafs: neurograph, connections = run_on_leaf( - neurograph, connections, leaf, search_radius, complex_bool + neurograph, + connections, + leaf, + search_radius, + complex_bool, + long_range_proposals, ) - # neurograph.filter_nodes() return neurograph -def run_on_leaf(neurograph, connections, leaf, search_radius, complex_bool): +def run_on_leaf( + neurograph, + connections, + leaf, + search_radius, + complex_bool, + long_range_proposals, +): """ Generates proposals emanating from "leaf". @@ -62,6 +78,9 @@ def run_on_leaf(neurograph, connections, leaf, search_radius, complex_bool): Maximum Euclidean distance between endpoints of proposal. complex_bool : bool Indication of whether to generate complex proposals. + long_range_proposals : bool + Indication of whether to generate simple proposals within distance of + 2 * search_radius of leaf. Returns ------- @@ -72,10 +91,17 @@ def run_on_leaf(neurograph, connections, leaf, search_radius, complex_bool): were added to "neurograph". """ + # Get candidates leaf_swc_id = neurograph.nodes[leaf]["swc_id"] - for xyz in get_candidates(neurograph, leaf, search_radius): + candidates = get_candidates(neurograph, leaf, search_radius) + if len(candidates) == 0 and long_range_proposals: + candidates = get_candidates(neurograph, leaf, 2 * search_radius) + candidates = parse_long_range(neurograph, candidates, leaf) + + # Parse candidates + for xyz in candidates: # Get connection - neurograph, node = get_conection(neurograph, leaf, xyz, search_radius) + neurograph, node = get_conection(neurograph, leaf, xyz) if not complex_bool and neurograph.degree[node] > 1: continue @@ -100,6 +126,19 @@ def run_on_leaf(neurograph, connections, leaf, search_radius, complex_bool): return neurograph, connections +def parse_long_range(neurograph, candidates, leaf): + hit_swc_ids = set() + filtered_candidates = [] + for xyz in candidates: + neurograph, i = get_conection(neurograph, leaf, xyz) + if neurograph.degree[i] > 1: + continue + else: + filtered_candidates.append(xyz) + hit_swc_ids.add(neurograph.nodes[i]["swc_id"]) + return filtered_candidates if len(hit_swc_ids) == 1 else [] + + def get_candidates(neurograph, leaf, search_radius): """ Generates proposals for node "leaf" in "neurograph" by finding candidate @@ -150,11 +189,10 @@ def get_best_candidates(neurograph, candidates, dists): return list(candidates.values()) -def get_conection(neurograph, leaf, xyz, search_radius): +def get_conection(neurograph, leaf, xyz): edge = neurograph.xyz_to_edge[xyz] node, d = get_closer_endpoint(neurograph, edge, xyz) if d > ENDPOINT_DIST: - # or neurograph.dist(leaf, node) > search_radius: attrs = neurograph.get_edge_data(*edge) idx = np.where(np.all(attrs["xyz"] == xyz, axis=1))[0][0] node = neurograph.split_edge(edge, attrs, idx) diff --git a/src/deep_neurographs/geometry.py b/src/deep_neurographs/geometry.py index 47924bc..1b00c65 100644 --- a/src/deep_neurographs/geometry.py +++ b/src/deep_neurographs/geometry.py @@ -16,7 +16,7 @@ # Directional Vectors -def get_directional(neurograph, i, origin, window_size): +def get_directional(neurograph, i, origin, depth): """ Computes the directional vector of a branch or bifurcation in a neurograph relative to a specified origin. @@ -30,9 +30,9 @@ def get_directional(neurograph, i, origin, window_size): origin : numpy.ndarray The origin point xyz relative to which the directional vector is computed. - window_size : numpy.ndarry - The size of the window around the branch or bifurcation to consider - for computing the directional vector. + depth : numpy.ndarry + The size of the window in microns around the branch or bifurcation to + consider for computing the directional vector. Returns ------- @@ -44,17 +44,17 @@ def get_directional(neurograph, i, origin, window_size): branches = neurograph.get_branches(i, ignore_reducibles=True) branches = shift_branches(branches, origin) if len(branches) == 1: - return compute_tangent(get_subarray(branches[0], window_size)) + return compute_tangent(get_subarray(branches[0], depth)) elif len(branches) == 2: - branch_1 = get_subarray(branches[0], window_size) - branch_2 = get_subarray(branches[1], window_size) + branch_1 = get_subarray(branches[0], depth) + branch_2 = get_subarray(branches[1], depth) branch = np.concatenate((branch_1, branch_2)) return compute_tangent(branch) else: return np.array([0, 0, 0]) -def get_subarray(arr, window_size): +def get_subarray(arr, depth): """ Extracts a sub-array of a specified window size from a given input array. @@ -62,8 +62,8 @@ def get_subarray(arr, window_size): ---------- branch : numpy.ndarray Array from which the sub-branch will be extracted. - window_size : int - Size of the window to extract from "arr". + depth : int + Size of the window in microns to extract from "arr". Returns ------- @@ -72,10 +72,12 @@ def get_subarray(arr, window_size): smaller than the window size, the entire branch array is returned. """ - if arr.shape[0] < window_size: - return arr - else: - return arr[0:window_size, :] + length = 0 + for i in range(1, arr.shape[0]): + length += dist(arr[i - 1], arr[i]) + if length > depth: + return arr[0:i, :] + return arr def compute_svd(xyz): diff --git a/src/deep_neurographs/machine_learning/feature_generation.py b/src/deep_neurographs/machine_learning/feature_generation.py index a0838d8..cca5807 100644 --- a/src/deep_neurographs/machine_learning/feature_generation.py +++ b/src/deep_neurographs/machine_learning/feature_generation.py @@ -273,10 +273,6 @@ def proposal_profiles(neurograph, proposals, img): return profiles -def generate_edge_profiles(neurograph, img): - pass - - def get_profile(img, coords, thread_id): """ Gets the image intensity profile for a given proposal. @@ -301,12 +297,24 @@ def get_profile(img, coords, thread_id): coords["bbox"]["max"] = [coords["bbox"]["max"][i] + 1 for i in range(3)] chunk = img_utils.read_tensorstore_with_bbox(img, coords["bbox"]) chunk = img_utils.normalize(chunk) - profile = [chunk[tuple(xyz)] for xyz in coords["profile_path"]] + profile = read_intensities(chunk, coords) avg, std = utils.get_avg_std(profile) profile.extend([avg, std]) return thread_id, profile +def read_intensities(img, coords): + profile = [] + for xyz in coords["profile_path"]: + start = xyz - 1 + end = xyz + 2 + val = np.max( + img[start[0]: end[0], start[1]: end[1], start[2]: end[2]] + ) + profile.append(val) + return profile + + def get_proposal_profile_coords(neurograph, proposal): """ Gets coordinates needed to compute an image intensity profile. @@ -330,10 +338,14 @@ def get_proposal_profile_coords(neurograph, proposal): coord_0 = utils.to_voxels(xyz_0) coord_1 = utils.to_voxels(xyz_1) - # Store coordinates + # Store local coordinates bbox = utils.get_minimal_bbox(coord_0, coord_1) - start = [coord_0[i] - bbox["min"][i] for i in range(3)] - end = [coord_1[i] - bbox["min"][i] for i in range(3)] + start = [coord_0[i] - bbox["min"][i] + 1 for i in range(3)] + end = [coord_1[i] - bbox["min"][i] + 1 for i in range(3)] + + # Shift bbox + bbox["min"] = [bbox["min"][i] - 1 for i in range(3)] + bbox["max"] = [bbox["max"][i] + 2 for i in range(3)] coords = { "bbox": bbox, "profile_path": geometry.make_line(start, end, N_PROFILE_PTS), diff --git a/src/deep_neurographs/machine_learning/graph_trainer.py b/src/deep_neurographs/machine_learning/graph_trainer.py index d6ab01b..dd968ea 100644 --- a/src/deep_neurographs/machine_learning/graph_trainer.py +++ b/src/deep_neurographs/machine_learning/graph_trainer.py @@ -27,6 +27,7 @@ # Training LR = 1e-3 +MODEL_TYPE = "GraphNeuralNet" N_EPOCHS = 200 SCHEDULER_GAMMA = 0.5 SCHEDULER_STEP_SIZE = 1000 @@ -224,7 +225,7 @@ def forward(self, data): """ self.optimizer.zero_grad() - x, edge_index = toGPU(data) + x, edge_index = gnn_utils.get_inputs(data, MODEL_TYPE) hat_y = self.model(x, edge_index) y = data.y # .to("cuda:0", dtype=torch.float32) return y, truncate(hat_y, y) diff --git a/src/deep_neurographs/neurograph.py b/src/deep_neurographs/neurograph.py index 7d0626e..a8d7008 100644 --- a/src/deep_neurographs/neurograph.py +++ b/src/deep_neurographs/neurograph.py @@ -275,6 +275,7 @@ def generate_proposals( self, search_radius, complex_bool=True, + long_range_proposals=False, proposals_per_leaf=3, optimize=False, optimization_depth=10, @@ -294,7 +295,10 @@ def generate_proposals( # Main self = generate_proposals.run( - self, search_radius, complex_bool=complex_bool + self, + search_radius, + complex_bool=complex_bool, + long_range_proposals=long_range_proposals, ) # Finish