diff --git a/pyproject.toml b/pyproject.toml index 9167aac..ff3dce9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dynamic = ["version"] dependencies = [ 'boto3', + 'fastremap', 'lightning', 'more_itertools', 'networkx', diff --git a/src/deep_neurographs/graph_utils.py b/src/deep_neurographs/graph_utils.py index 361e564..c0877ec 100644 --- a/src/deep_neurographs/graph_utils.py +++ b/src/deep_neurographs/graph_utils.py @@ -60,12 +60,13 @@ def get_irreducibles(swc_dict, swc_id=None, prune=True, depth=16, smooth=True): """ # Build dense graph + swc_dict["idx"] = dict(zip(swc_dict["id"], range(len(swc_dict["id"])))) dense_graph = swc_utils.to_graph(swc_dict) if prune: dense_graph = prune_short_branches(dense_graph, depth) # Extract nodes - leafs, junctions = get_irreducible_nodes(dense_graph, swc_dict) + leafs, junctions = get_irreducible_nodes(dense_graph) if len(leafs) == 0: return False, None @@ -100,7 +101,7 @@ def get_irreducibles(swc_dict, swc_id=None, prune=True, depth=16, smooth=True): return swc_id, irreducibles -def get_irreducible_nodes(graph, swc_dict): +def get_irreducible_nodes(graph): """ Gets irreducible nodes (i.e. leafs and junctions) of a graph. @@ -315,7 +316,8 @@ def init_edge_attrs(swc_dict, i): Edge attribute dictionary. """ - return {"radius": [swc_dict["radius"][i]], "xyz": [swc_dict["xyz"][i]]} + j = swc_dict["idx"][i] + return {"radius": [swc_dict["radius"][j]], "xyz": [swc_dict["xyz"][j]]} def upd_edge_attrs(swc_dict, attrs, i): @@ -337,8 +339,9 @@ def upd_edge_attrs(swc_dict, attrs, i): Edge attribute dictionary. """ - attrs["radius"].append(swc_dict["radius"][i]) - attrs["xyz"].append(swc_dict["xyz"][i]) + j = swc_dict["idx"][i] + attrs["radius"].append(swc_dict["radius"][j]) + attrs["xyz"].append(swc_dict["xyz"][j]) return attrs @@ -403,7 +406,8 @@ def set_node_attrs(swc_dict, nodes): """ attrs = dict() for i in nodes: - attrs[i] = {"radius": swc_dict["radius"][i], "xyz": swc_dict["xyz"][i]} + j = swc_dict["idx"][i] + attrs[i] = {"radius": swc_dict["radius"][j], "xyz": swc_dict["xyz"][j]} return attrs @@ -434,7 +438,8 @@ def upd_node_attrs(swc_dict, leafs, junctions, i): Updated dictionary if "i" was contained in "junctions.keys()". """ - upd_attrs = {"radius": swc_dict["radius"][i], "xyz": swc_dict["xyz"][i]} + j = swc_dict["idx"][i] + upd_attrs = {"radius": swc_dict["radius"][j], "xyz": swc_dict["xyz"][j]} if i in leafs: leafs[i] = upd_attrs else: diff --git a/src/deep_neurographs/intake.py b/src/deep_neurographs/intake.py index 23250c3..792b479 100644 --- a/src/deep_neurographs/intake.py +++ b/src/deep_neurographs/intake.py @@ -207,6 +207,8 @@ def download_gcs_zips(bucket_name, cloud_path, min_size): cnt, t1 = report_progress( i, len(zip_paths), chunk_size, cnt, t0, t1 ) + if len(swc_dicts) > 2000: + stop return swc_dicts diff --git a/src/deep_neurographs/swc_utils.py b/src/deep_neurographs/swc_utils.py index d8e6a0c..3cd6c51 100644 --- a/src/deep_neurographs/swc_utils.py +++ b/src/deep_neurographs/swc_utils.py @@ -60,9 +60,8 @@ def parse_local_swc(path, bbox=None, min_size=0): def parse_gcs_zip(zip_file, path, min_size=0): contents = read_from_gcs_zip(zip_file, path) - swc_dict = ( - fast_parse(contents) if len(contents) > min_size else {"id": [-1]} - ) + parse_bool = len(contents) > min_size + swc_dict = fast_parse(contents) if parse_bool else {"id": [-1]} return utils.get_swc_id(path), swc_dict @@ -143,9 +142,14 @@ def fast_parse(contents): min_id = np.min(swc_dict["id"]) swc_dict["id"] -= min_id swc_dict["pid"] -= min_id + swc_dict["radius"] /= 1000.0 return swc_dict +def reindex(arr, idxs): + return arr[idxs] + + def get_contents(swc_contents): offset = [0, 0, 0] for i, line in enumerate(swc_contents): @@ -333,10 +337,12 @@ def to_graph(swc_dict, graph_id=None, set_attrs=False): def __add_attributes(swc_dict, graph): - xyz = swc_dict["xyz"] - radii = swc_dict["radius"] - attrs = [{"xyz": xyz[i], "radius": radii[i]} for i in graph.nodes] - nx.set_node_attributes(graph, dict(zip(swc_dict["id"], attrs))) + attrs = dict() + for idx, node_id in enumerate(swc_dict["id"]): + attrs[node_id] = { + "xyz": swc_dict["xyz"][idx], "radius": swc_dict["radius"][idx] + } + nx.set_node_attributes(graph, attrs) return graph diff --git a/src/deep_neurographs/utils.py b/src/deep_neurographs/utils.py index f72695a..c9d1d8d 100644 --- a/src/deep_neurographs/utils.py +++ b/src/deep_neurographs/utils.py @@ -15,6 +15,7 @@ import shutil from copy import deepcopy from io import BytesIO +from skimage.color import label2rgb from time import time from zipfile import ZipFile @@ -425,6 +426,11 @@ def write_json(path, contents): json.dump(contents, f) +def write_txt(path, contents): + f = open(path, "w") + f.write(contents) + f.close() + # --- coordinate conversions --- def world_to_img(neurograph, node_or_xyz): if type(node_or_xyz) == int: @@ -500,6 +506,12 @@ def get_img_mip(img, axis=0): return np.max(img, axis=axis) +def get_labels_mip(img, axis=0): + mip = np.max(img, axis=axis) + mip = label2rgb(mip) + return (255 * mip).astype(np.uint8) + + def normalize_img(img): img -= np.min(img) return img / np.max(img)