From 0bf602bddc3866c902ce8053004e8800529aa99a Mon Sep 17 00:00:00 2001 From: anna-grim Date: Thu, 11 Jan 2024 20:59:07 +0000 Subject: [PATCH] major upd : asynchronous intake --- src/deep_neurographs/feature_extraction.py | 36 +--- src/deep_neurographs/geometry_utils.py | 4 +- src/deep_neurographs/graph_utils.py | 23 ++- src/deep_neurographs/intake.py | 220 ++++++++++----------- src/deep_neurographs/neurograph.py | 147 ++++---------- src/deep_neurographs/swc_utils.py | 9 + src/deep_neurographs/utils.py | 35 ++-- 7 files changed, 197 insertions(+), 277 deletions(-) diff --git a/src/deep_neurographs/feature_extraction.py b/src/deep_neurographs/feature_extraction.py index 41e1eb0..65b36f5 100644 --- a/src/deep_neurographs/feature_extraction.py +++ b/src/deep_neurographs/feature_extraction.py @@ -164,42 +164,14 @@ def generate_img_profiles(neurograph, path, anisotropy=[1.0, 1.0, 1.0]): ) img = utils.normalize_img(img) for edge in neurograph.mutable_edges: - if neurograph.optimize_alignment or neurograph.optimize_path: - xyz = to_img_coords(neurograph, edge) - path = geometry_utils.sample_path(xyz, N_PROFILE_POINTS) - else: - i, j = tuple(edge) - xyz_i = utils.world_to_img(neurograph, i) - xyz_j = utils.world_to_img(neurograph, j) - path = geometry_utils.make_line(xyz_i, xyz_j, N_PROFILE_POINTS) + xyz_i, xyz_j = neurograph.get_edge_attr("xyz", edge) + xyz_i = utils.world_to_img(neurograph, xyz_i) + xyz_j = utils.world_to_img(neurograph, xyz_j) + path = geometry_utils.make_line(xyz_i, xyz_j, N_PROFILE_POINTS) features[edge] = geometry_utils.get_profile(img, path, window=WINDOW) return features -def to_img_coords(neurograph, edge): - """ - Converts xyz coordinate of each vertex in "edge" from real world to image - coordinates. - - Parameters - ---------- - neurograph : NeuroGraph - NeuroGraph generated from a directory of swcs generated from a - predicted segmentation. - edge : frozenset - The edge from "neurograph" to perform coordinate transformation. - - Returns - ------- - img_coords : numpy.ndarray - Image coordinates of each vertex from "edge". - - """ - xyz_list = neurograph.edges[edge]["xyz"] - img_coords = [utils.world_to_img(neurograph, xyz) for xyz in xyz_list] - return np.array(img_coords) - - def generate_mutable_skel_features(neurograph): features = dict() for edge in neurograph.mutable_edges: diff --git a/src/deep_neurographs/geometry_utils.py b/src/deep_neurographs/geometry_utils.py index db7ea80..7bcb57b 100644 --- a/src/deep_neurographs/geometry_utils.py +++ b/src/deep_neurographs/geometry_utils.py @@ -9,9 +9,9 @@ # Directional Vectors -def get_directional(neurograph, i, proposal_tangent, window=5): +def get_directional(neurograph, i, proposal_tangent, window=5, n_svd_points=10): directionals = [] - d = neurograph.optimize_depth + d = n_svd_points for branch in neurograph.get_branches(i): if branch.shape[0] >= window + d: xyz = deepcopy(branch[d:, :]) diff --git a/src/deep_neurographs/graph_utils.py b/src/deep_neurographs/graph_utils.py index c94d1e7..79db58e 100644 --- a/src/deep_neurographs/graph_utils.py +++ b/src/deep_neurographs/graph_utils.py @@ -8,6 +8,15 @@ Routines that extract the irreducible components of a graph. --define what an irreducible is + leafs : set + Nodes with degreee 1. + junctions : set + Nodes with degree > 2. + edges : dict + Set of edges connecting nodes in leafs and junctions. The keys are + pairs of nodes connected by an edge and values are a dictionary of + attributes. + --define what a branch is """ @@ -41,14 +50,10 @@ def get_irreducibles(swc_dict, prune=True, depth=16, smooth=True): Returns ------- - leafs : set - Nodes with degreee 1. - junctions : set - Nodes with degree > 2. - edges : dict - Set of edges connecting nodes in leafs and junctions. The keys are - pairs of nodes connected by an edge and values are a dictionary of - attributes. + dict + Irreducibles stored in a dictionary where key-values are type of + irreducible (i.e. leaf, junction, or edge) and corresponding set of + all irreducibles from the graph of that type. """ # Initializations @@ -80,7 +85,7 @@ def get_irreducibles(swc_dict, prune=True, depth=16, smooth=True): nbs = append_value(nbs, root, j) nbs = append_value(nbs, j, root) root = None - return leafs, junctions, edges + return {"leafs": leafs, "junctions": junctions, "edges": edges} def get_irreducible_nodes(graph): diff --git a/src/deep_neurographs/intake.py b/src/deep_neurographs/intake.py index 5f98381..997a18f 100644 --- a/src/deep_neurographs/intake.py +++ b/src/deep_neurographs/intake.py @@ -20,60 +20,65 @@ from deep_neurographs import graph_utils as gutils from deep_neurographs import swc_utils, utils from deep_neurographs.neurograph import NeuroGraph -from deep_neurographs.swc_utils import parse_gcs_zip +from deep_neurographs.swc_utils import parse_gcs_zip, process_local_paths N_PROPOSALS_PER_LEAF = 3 -OPTIMIZE_ALIGNMENT = False -OPTIMIZE_DEPTH = 15 +OPTIMIZE_PROPOSALS = False +OPTIMIZATION_DEPTH = 15 PRUNE = True PRUNE_DEPTH = 16 SEARCH_RADIUS = 0 MIN_SIZE = 30 -SMOOTH = False +SMOOTH = True -# --- Build graph --- +# --- Build graph wrappers --- def build_neurograph_from_local( swc_dir=None, swc_paths=None, - img_patch_shape=None, img_patch_origin=None, + img_patch_shape=None, img_path=None, + min_size=MIN_SIZE, n_proposals_per_leaf=N_PROPOSALS_PER_LEAF, prune=PRUNE, prune_depth=PRUNE_DEPTH, - optimize_alignment=OPTIMIZE_ALIGNMENT, - optimize_depth=OPTIMIZE_DEPTH, + optimize_proposals=OPTIMIZE_PROPOSALS, + optimization_depth=OPTIMIZATION_DEPTH, search_radius=SEARCH_RADIUS, - min_size=MIN_SIZE, smooth=SMOOTH, ): + # Process swc files + t0 = time() assert utils.xor(swc_dir, swc_paths), "Error: provide swc_dir or swc_paths" - neurograph = NeuroGraph( - swc_dir=swc_dir, + bbox = utils.get_bbox(img_patch_origin, img_patch_shape) + paths = get_paths(swc_dir) if swc_dir else swc_paths + swc_dicts = process_local_paths(paths, min_size, bbox=bbox) + print(f"process_local_paths(): {time() - t0} seconds") + + # Build neurograph + t0 = time() + neurograph = build_neurograph( + swc_dicts, + bbox=bbox, img_path=img_path, - optimize_depth=optimize_depth, - optimize_alignment=optimize_alignment, - origin=img_patch_origin, - shape=img_patch_shape, - ) - print("Build Graph...") - neurograph = init_immutables_from_local( - neurograph, - swc_dir=swc_dir, - swc_paths=swc_paths, prune=prune, prune_depth=prune_depth, - min_size=min_size, smooth=smooth, - ) + ) + print(f"build_neurograph(): {time() - t0} seconds") + + # Generate proposals t0 = time() if search_radius > 0: neurograph.generate_proposals( + search_radius, n_proposals_per_leaf=n_proposals_per_leaf, - search_radius=search_radius, + optimize=optimize_proposals, + optimization_depth=optimization_depth, ) - print(f" generate_proposals(): {time() - t0} seconds") + print(f"generate_proposals(): {time() - t0} seconds") + return neurograph @@ -86,8 +91,8 @@ def build_neurograph_from_gcs_zips( search_radius=SEARCH_RADIUS, prune=PRUNE, prune_depth=PRUNE_DEPTH, - optimize_alignment=OPTIMIZE_ALIGNMENT, - optimize_depth=OPTIMIZE_DEPTH, + optimize_proposals=OPTIMIZE_PROPOSALS, + optimization_depth=OPTIMIZATION_DEPTH, smooth=SMOOTH, ): """ @@ -99,27 +104,34 @@ def build_neurograph_from_gcs_zips( Name of GCS bucket where zips are stored. cloud_path : str Path within GCS bucket to directory containing zips. - img_path : str + img_path : str, optional Path to image stored GCS Bucket that swc files were generated from. - min_size : int - Minimum path length of swc files which are stored. - n_proposals_per_leaf : int + The default is None. + min_size : int, optional + Minimum path length of swc files which are stored. The default is the + global variable "MIN_SIZE". + n_proposals_per_leaf : int, optional Number of edge proposals generated from each leaf node in an swc file. - search_radius : float - Maximum Euclidean length of an edge proposal. - prune : bool - Indication of whether to prune short branches. - prune_depth : int + The default is the global variable "N_PROPOSALS_PER_LEAF". + search_radius : float, optional + Maximum Euclidean length of an edge proposal. The default is the + global variable "SEARCH_RADIUS". + prune : bool, optional + Indication of whether to prune short branches. The default is the + global variable "PRUNE". + prune_depth : int, optional Branches less than "prune_depth" microns are pruned if "prune" is - True. - optimize_alignment : bool + True. The default is the global variable "PRUNE_DEPTH". + optimize_proposals : bool, optional Indication of whether to optimize alignment of edge proposals to image - signal. - optimize_depth : int + signal. The default is the global variable "OPTIMIZE_PROPOSALS". + optimization_depth : int, optional Distance from each edge proposal end point that is search during - alignment optimization. - smooth : bool - Indication of whether to smooth branches from swc files. + alignment optimization. The default is the global variable + "OPTIMIZATION_DEPTH". + smooth : bool, optional + Indication of whether to smooth branches from swc files. The default + is the global variable "SMOOTH". Returns ------- @@ -127,55 +139,24 @@ def build_neurograph_from_gcs_zips( Neurograph generated from zips of swc files stored in a GCS bucket. """ - swc_dicts = download_gcs_zips(bucket_name, cloud_path, min_size=min_size) + swc_dicts = download_gcs_zips(bucket_name, cloud_path, min_size) neurograph = build_neurograph( swc_dicts, img_path=img_path, prune=prune, prune_depth=prune_depth, smooth=smooth, - optimize_alignment=OPTIMIZE_ALIGNMENT, - optimize_depth=OPTIMIZE_DEPTH, ) if search_radius > 0: neurograph.generate_proposals( + search_radius, n_proposals_per_leaf=n_proposals_per_leaf, - search_radius=search_radius, ) return neurograph -def init_immutables_from_local( - neurograph, - swc_dir=None, - swc_paths=None, - prune=PRUNE, - prune_depth=PRUNE_DEPTH, - min_size=MIN_SIZE, - smooth=SMOOTH, -): - neurograph.extraction = 0 - neurograph.add_edges_timer = 0 - swc_paths = get_paths(swc_dir) if swc_dir else swc_paths - for path in swc_paths: - neurograph.ingest_swc_from_local( - path, prune=True, prune_depth=16, smooth=smooth - ) - print( - f" extract_irreducible_graph(): {neurograph.extraction} seconds" - ) - print(f" add_edges(): {neurograph.add_edges_timer} seconds") - return neurograph - - -def get_paths(swc_dir): - paths = [] - for f in utils.listdir(swc_dir, ext=".swc"): - paths.append(os.path.join(swc_dir, f)) - return paths - - -def download_gcs_zips(bucket_name, cloud_path, min_size=0): +# -- Read swc files -- +def download_gcs_zips(bucket_name, cloud_path, min_size): """ Downloads swc files from zips stored in a GCS bucket. @@ -262,58 +243,43 @@ def list_gcs_filenames(bucket, cloud_path, extension): return [blob.name for blob in blobs if extension in blob.name] -def report_runtimes( - n_files, n_files_completed, chunk_size, start, start_chunk -): - runtime = time() - start - chunk_runtime = time() - start_chunk - n_files_remaining = n_files - n_files_completed - rate = chunk_runtime / chunk_size - eta = (runtime + n_files_remaining * rate) / 60 - files_processed = f"{n_files_completed - chunk_size}-{n_files_completed}" - print(f"Completed: {round(100 * n_files_completed / n_files, 2)}%") - print( - f"Runtime for Zips {files_processed}: {round(chunk_runtime, 4)} seconds" - ) - print(f"Zip Processing Rate: {file_rate} seconds") - print(f"Approximate Total Runtime: {round(eta, 4)} minutes") - print("") - - +# -- Build neurograph --- def build_neurograph( swc_dicts, + bbox=None, img_path=None, - optimize_alignment=OPTIMIZE_ALIGNMENT, - optimize_depth=OPTIMIZE_DEPTH, prune=PRUNE, prune_depth=PRUNE_DEPTH, smooth=SMOOTH, ): - graph_list = build_graphs(swc_dicts, prune, prune_depth, smooth) - start_ids = get_start_ids(swc_dicts) - print("Total Runtime:", 1600 * t) - stop + # Extract irreducibles + t0 = time() + n_components = len(swc_dicts) + irreducibles = [None] * n_components + for i in range(n_components): + irreducibles[i] = gutils.get_irreducibles( + swc_dicts[i], prune=prune, depth=prune_depth, smooth=smooth + ) + print(f" --> get_irreducibles(): {time() - t0} seconds") + + # Build neurograph + t0 = time() neurograph = NeuroGraph( + bbox=bbox, img_path=img_path, - optimize_alignment=optimize_alignment, - optimize_depth=optimize_depth, ) + for i in range(n_components): + neurograph.add_immutables(swc_dicts[i], irreducibles[i]) + print(f" --> add_irreducibles(): {time() - t0} seconds") + return neurograph -def build_graphs(swc_dicts, prune, prune_depth, smooth): - t0 = time() - graphs = [None] * len(swc_dicts) - for i, swc_dict in enumerate(swc_dicts): - graphs[i] = build_subgraph(swc_dict) - t = time() - t0 - print(f"build_subgraphs(): {t} seconds") - return graphs - - -def build_subgraph(swc_dict): - graph = nx.Graph() - graph.add_edges_from(zip(swc_dict["id"][1:], swc_dict["pid"][1:])) - return graph +# -- Utils -- +def get_paths(swc_dir): + paths = [] + for f in utils.listdir(swc_dir, ext=".swc"): + paths.append(os.path.join(swc_dir, f)) + return paths def get_start_ids(swc_dicts): @@ -327,3 +293,21 @@ def get_start_ids(swc_dicts): node_ids.append(cnt) cnt += len(leafs) + len(junctions) return node_ids + + +def report_runtimes( + n_files, n_files_completed, chunk_size, start, start_chunk +): + runtime = time() - start + chunk_runtime = time() - start_chunk + n_files_remaining = n_files - n_files_completed + rate = chunk_runtime / chunk_size + eta = (runtime + n_files_remaining * rate) / 60 + files_processed = f"{n_files_completed - chunk_size}-{n_files_completed}" + print(f"Completed: {round(100 * n_files_completed / n_files, 2)}%") + print( + f"Runtime for Zips {files_processed}: {round(chunk_runtime, 4)} seconds" + ) + print(f"Zip Processing Rate: {file_rate} seconds") + print(f"Approximate Total Runtime: {round(eta, 4)} minutes") + print("") diff --git a/src/deep_neurographs/neurograph.py b/src/deep_neurographs/neurograph.py index 0cb9b3a..e648c6c 100644 --- a/src/deep_neurographs/neurograph.py +++ b/src/deep_neurographs/neurograph.py @@ -35,54 +35,37 @@ class NeuroGraph(nx.Graph): def __init__( self, + bbox=None, swc_dir=None, img_path=None, label_mask=None, - optimize_depth=10, - optimize_alignment=False, - optimize_path=False, - origin=None, - shape=None, - size_threshold=30, ): - """ - Parameters - ---------- - None. - - Returns - ------- - None. - - """ super(NeuroGraph, self).__init__() - self.path = swc_dir + # Initialize paths + self.img_path = img_path self.label_mask = label_mask + self.swc_paths = swc_dir + + # Initialize node and edge sets self.leafs = set() self.junctions = set() - self.size_threshold = size_threshold - self.immutable_edges = set() self.mutable_edges = set() self.target_edges = set() - self.xyz_to_edge = dict() - self.kdtree = None - self.img_path = img_path - self.optimize_depth = optimize_depth - self.optimize_alignment = optimize_alignment - self.optimize_path = optimize_path - self.simple_proposals = set() + # Initialize data structures for proposals self.complex_proposals = set() + self.simple_proposals = set() + self.xyz_to_edge = dict() + self.kdtree = None - self.bbox = None - self.shape = shape - if origin and shape: - self.bbox = { - "min": np.array(origin), - "max": np.array([origin[i] + shape[i] for i in range(3)]), - } - self.origin = np.array(origin) + # Initialize bounding box (if exists) + self.bbox = bbox + if self.bbox: + self.origin = bbox["min"] + self.shape = (bbox["max"] - bbox["min"]).astype(int) + else: + self.shape = None def init_immutable_graph(self, add_attrs=False): immutable_graph = nx.Graph() @@ -99,57 +82,26 @@ def init_predicted_graph(self): self.predicted_graph = self.init_immutable_graph() def init_densegraph(self): - self.densegraph = DenseGraph(self.path) + self.densegraph = DenseGraph(self.swc_paths) # --- Add nodes or edges --- - def ingest_swc_from_local( - self, path, prune=True, prune_depth=16, smooth=True - ): - # Parse swc - swc_id = utils.get_id(path) - swc_dict = swc_utils.parse_local_swc(path, bbox=self.bbox) - if len(swc_dict["xyz"]) > self.size_threshold: - return None - - # Build neurograph - t0 = time() - leafs, junctions, edges = gutils.get_irreducibles( - swc_dict, prune=prune, depth=prune_depth, smooth=smooth - ) - self.extraction += time() - t0 - self.add_immutables(swc_id, swc_dict, leafs, junctions, edges) - - def add_immutables(self, swc_id, swc_dict, leafs, junctions, edges): - """ - Adds nodes to graph from a dictionary generated from an swc files. - - Parameters - ---------- - node_id : int - Node id. - swc_dict : dict - Dictionary generated from an swc where the keys are swc - attributes. - - Returns - ------- - None. - - """ + def add_immutables(self, swc_dict, irreducibles): # Add nodes - node_id = dict() + leafs = irreducibles["leafs"] + junctions = irreducibles["junctions"] for i in list(leafs) + list(junctions): node_id[i] = len(self.nodes) self.add_node( node_id[i], xyz=np.array(swc_dict["xyz"][i]), radius=swc_dict["radius"][i], - swc_id=swc_id, + swc_id=swc_dict["swc_id"], ) # Add edges t0 = time() + edges = irreducibles["edges"] for i, j in edges.keys(): # Get edge edge = (node_id[i], node_id[j]) @@ -159,7 +111,11 @@ def add_immutables(self, swc_id, swc_dict, leafs, junctions, edges): # Add edge self.immutable_edges.add(frozenset(edge)) self.add_edge( - node_id[i], node_id[j], xyz=xyz, radius=radii, swc_id=swc_id + node_id[i], + node_id[j], + xyz=xyz, + radius=radii, + swc_id=swc_dict["swc_id"] ) xyz_to_edge = dict((tuple(xyz), edge) for xyz in xyz) check_xyz = set(xyz_to_edge.keys()) @@ -168,7 +124,6 @@ def add_immutables(self, swc_id, swc_dict, leafs, junctions, edges): for xyz in collisions: del xyz_to_edge[xyz] self.xyz_to_edge.update(xyz_to_edge) - self.add_edges_timer += time() - t0 # Update leafs and junctions for l in leafs: @@ -178,7 +133,7 @@ def add_immutables(self, swc_id, swc_dict, leafs, junctions, edges): self.junctions.add(node_id[j]) # --- Proposal Generation --- - def generate_proposals(self, n_proposals_per_leaf=3, search_radius=25.0): + def generate_proposals(self, search_radius, n_proposals_per_leaf=3, optimize=False, optimization_depth=10): """ Generates edges for the graph. @@ -212,6 +167,7 @@ def generate_proposals(self, n_proposals_per_leaf=3, search_radius=25.0): node = j xyz = self.nodes[node]["xyz"] else: + # run complex optimization here idxs = np.where(np.all(attrs["xyz"] == xyz, axis=1))[0] node = self.add_immutable_node((i, j), attrs, idxs[0]) @@ -219,7 +175,8 @@ def generate_proposals(self, n_proposals_per_leaf=3, search_radius=25.0): self.add_edge(leaf, node, xyz=np.array([xyz_leaf, xyz])) self.mutable_edges.add(frozenset((leaf, node))) - if self.optimize_alignment or self.optimize_path: + # Check whether to optimization proposals + if optimize: self.run_optimization() def __get_proposals( @@ -352,13 +309,7 @@ def run_optimization(self): ) for edge in self.mutable_edges: xyz_1, xyz_2 = geometry_utils.optimize_alignment(self, img, edge) - proposal = [self.to_world(xyz_1)] - if self.optimize_path: - path = geometry_utils.optimize_path( - img, self.origin, xyz_1, xyz_2 - ) - proposal.append(path) - proposal.append(self.to_world(xyz_2)) + proposal = [self.to_world(xyz_1), self.to_world(xyz_2)] self.edges[edge]["xyz"] = np.vstack(proposal) def get_branch(self, xyz_or_node): @@ -390,7 +341,7 @@ def orient_edge(self, edge, i): def init_targets(self, target_neurograph): # Initializations msg = "Error: Provide swc_dir/swc_paths to initialize target edges!" - assert target_neurograph.path, msg + assert target_neurograph.swc_path, msg target_neurograph.init_densegraph() target_neurograph.init_kdtree() self.target_edges = set() @@ -569,8 +520,7 @@ def get_immutable_nbs(self, i): return nbs def compute_length(self, edge, metric="l2"): - i, j = tuple(edge) - xyz_1, xyz_2 = self.get_edge_attr("xyz", i, j) + xyz_1, xyz_2 = self.get_edge_attr("xyz", edge) return get_dist(xyz_1, xyz_2, metric=metric) def path_length(self, metric="l2"): @@ -592,14 +542,9 @@ def is_contained(self, node_or_xyz): if self.bbox: if type(node_or_xyz) == int: node_or_xyz = deepcopy(self.nodes[node_or_xyz]["xyz"]) - xyz = utils.apply_anisotropy(node_or_xyz - self.bbox["min"]) - img_shape = np.array(self.shape) - for i in range(3): - lower_bool = xyz[i] < 32 - upper_bool = xyz[i] > img_shape[i] - 32 - if lower_bool or upper_bool: - return False - return True + return utils.is_contained(self.bbox, node_or_xyz) + else: + return True def is_leaf(self, i): return True if self.immutable_degree(i) == 1 else False @@ -613,13 +558,10 @@ def creates_cycle(self, edge): self.predicted_graph.remove_edges_from([edge]) return True - def get_edge_attr(self, key, i, j): - attr_1 = self.nodes[i][key] - attr_2 = self.nodes[j][key] - return attr_1, attr_2 - - def get_center(self): - return geometry_utils.get_midpoint(self.bbox["min"], self.bbox["max"]) + def get_edge_attr(self, key, edge): + i, j = edge + xyz_arr = gutils.get_edge_attr(self, edge, key) + return xyz_arr[0], xyz_arr[1] def get_complex_proposals(self): return set([e for e in self.mutable_edges if not self.is_simple(e)]) @@ -629,10 +571,7 @@ def get_simple_proposals(self): def is_simple(self, edge): i, j = tuple(edge) - if self.immutable_degree(i) == 1 and self.immutable_degree(j) == 1: - return True - else: - return False + return True if self.is_leaf(i) and self.is_leaf(j) else False def to_img(self, node_or_xyz): if type(node_or_xyz) == int: diff --git a/src/deep_neurographs/swc_utils.py b/src/deep_neurographs/swc_utils.py index 9dcf7c1..a97df60 100644 --- a/src/deep_neurographs/swc_utils.py +++ b/src/deep_neurographs/swc_utils.py @@ -22,6 +22,15 @@ # -- io utils -- +def process_local_paths(paths, min_size, bbox=None): + swc_dicts = [] + for path in paths: + swc_dict_i = parse_local_swc(path, bbox=bbox) + swc_dict_i["swc_id"] = utils.get_swc_id(path) + swc_dicts.append(swc_dict_i) + return swc_dicts + + def parse_local_swc(path, bbox=None, min_size=0): swc_contents = read_from_local(path) parse_bool = len(swc_contents) > min_size diff --git a/src/deep_neurographs/utils.py b/src/deep_neurographs/utils.py index e10be55..0cd9b0c 100644 --- a/src/deep_neurographs/utils.py +++ b/src/deep_neurographs/utils.py @@ -19,7 +19,7 @@ import tensorstore as ts import zarr -ANISOTROPY = [0.748, 0.748, 1.0] +ANISOTROPY = np.array([0.748, 0.748, 1.0]) SUPPORTED_DRIVERS = ["neuroglancer_precomputed", "zarr"] @@ -393,15 +393,15 @@ def to_world(xyz, shift=[0, 0, 0]): return tuple([xyz[i] * ANISOTROPY[i] - shift[i] for i in range(3)]) -def to_img(xyz, shift=[0, 0, 0]): +def to_img(xyz, shift=np.array([0, 0, 0])): return apply_anisotropy(xyz - shift, return_int=True) def apply_anisotropy(xyz, return_int=False): if return_int: - return [round(xyz[i] / ANISOTROPY[i]) for i in range(3)] + return (xyz / ANISOTROPY).astype(int) else: - return [xyz[i] / ANISOTROPY[i] for i in range(3)] + return xyz / ANISOTROPY # --- math utils --- @@ -413,17 +413,28 @@ def get_avg_std(data, weights=None): def is_contained(bbox, xyz): xyz = apply_anisotropy(xyz - bbox["min"]) - dims = bbox["max"] - bbox["min"] - for i in range(3): - lower_bool = xyz[i] < 0 - upper_bool = xyz[i] >= dims[i] - if lower_bool or upper_bool: - return False - return True + shape = bbox["max"] - bbox["min"] + if any(xyz < 0) or any(xyz >= shape): + return False + else: + return True # --- miscellaneous --- -def get_id(path): +def get_bbox(origin, shape): + """ + Origin is assumed to be top, front, left corner. + + """ + if origin and shape: + origin = np.array(origin) + shape = np.array(shape) + return {"min": origin, "max": origin + shape} + else: + return None + + +def get_swc_id(path): """ Gets segment id of the swc file at "path".