diff --git a/src/segmentation_skeleton_metrics/data_handling/graph_classes.py b/src/segmentation_skeleton_metrics/data_handling/graph_classes.py index 22f07b4..e914144 100644 --- a/src/segmentation_skeleton_metrics/data_handling/graph_classes.py +++ b/src/segmentation_skeleton_metrics/data_handling/graph_classes.py @@ -10,7 +10,7 @@ from collections import defaultdict, deque from io import StringIO -from scipy.spatial import distance, KDTree +from scipy.spatial import KDTree import networkx as nx import numpy as np @@ -150,7 +150,7 @@ def dist(self, i, j): float Distance between voxel coordinates of the given nodes. """ - return distance.euclidean(self.node_voxel[i], self.node_voxel[j]) + return np.linalg.norm(self.node_voxel[i] - self.node_voxel[j]) def physical_dist(self, i, j): """ @@ -170,7 +170,7 @@ def physical_dist(self, i, j): Euclidean distance between physical coordinates of the given nodes. """ - return distance.euclidean(self.node_xyz(i), self.node_xyz(j)) + return np.linalg.norm(self.node_xyz(i) - self.node_xyz(j)) def leafs(self): """ @@ -562,7 +562,7 @@ def __init__( name=None, label=None, segment_id=None, - swc_color="" + swc_color="", ): """ Instantiates a FragmentGraph object. @@ -585,9 +585,7 @@ def __init__( Color used in SWC file of the graph. Default is an empty string. """ # Call parent class - super().__init__( - anisotropy=anisotropy, name=name, swc_color=swc_color - ) + super().__init__(anisotropy=anisotropy, name=name, swc_color=swc_color) # Instance attributes self.label = label diff --git a/src/segmentation_skeleton_metrics/data_handling/graph_loading.py b/src/segmentation_skeleton_metrics/data_handling/graph_loading.py index ea4a608..fce641f 100644 --- a/src/segmentation_skeleton_metrics/data_handling/graph_loading.py +++ b/src/segmentation_skeleton_metrics/data_handling/graph_loading.py @@ -291,24 +291,30 @@ def to_graph(self, swc_dict): Dict[str, SkeletonGraph] Graph built from an SWC file. """ - # Initialize graph + # Initializations graph = self._init_graph(swc_dict) - key = graph.name if self.is_groundtruth else graph.label - - # Build graph structure - id_lookup = dict() - for i, id_i in enumerate(swc_dict["id"]): - id_lookup[id_i] = i - if swc_dict["pid"][i] != -1: - parent = id_lookup[swc_dict["pid"][i]] - graph.add_edge(i, parent) - graph.run_length += graph.dist(i, parent) - - # Apply voxel coordinate conversion (if applicable) + ids = swc_dict["id"] + pids = swc_dict["pid"] + voxels = swc_dict["voxel"] + + # Build id->index lookup + id_lookup = {id_i: i for i, id_i in enumerate(ids)} + + # Vectorized edge building + children = np.where(pids != -1)[0] + parents = np.array([id_lookup[pids[i]] for i in children], dtype=int) + edges = list(zip(children.tolist(), parents.tolist())) + graph.add_edges_from(edges) + + # Vectorized run length + graph.run_length = np.linalg.norm( + voxels[children] - voxels[parents], axis=1 + ).sum() + if self.use_anisotropy: graph.node_voxel = (graph.node_voxel / self.anisotropy).astype(int) graph.node_voxel[:, [0, 2]] = graph.node_voxel[:, [2, 0]] - return {key: graph} + return {graph.name: graph} def _init_graph(self, swc_dict): """ diff --git a/src/segmentation_skeleton_metrics/data_handling/swc_loading.py b/src/segmentation_skeleton_metrics/data_handling/swc_loading.py index 4f0aeaf..da761cb 100644 --- a/src/segmentation_skeleton_metrics/data_handling/swc_loading.py +++ b/src/segmentation_skeleton_metrics/data_handling/swc_loading.py @@ -100,7 +100,7 @@ def __call__(self, swc_pointer): # Local SWC files paths = util.read_paths(swc_pointer, extension=".swc") if len(paths) > 0: - return self.read_swcs(paths, self.read_swc) + return self.read_swcs(paths) raise Exception("Directory is Invalid!") @@ -144,14 +144,12 @@ def read_swc(self, path): else: return None - def read_swcs(self, swc_paths, read_fn): + def read_swcs(self, swc_paths): """ Reads SWC files stored in a GCS or S3 bucket. Parameters ---------- - bucket_name : str - Name of bucket containing SWC files. swc_paths : List[str] List of paths to SWC files to be read. @@ -166,7 +164,7 @@ def read_swcs(self, swc_paths, read_fn): threads = set() for path in swc_paths: if self.confirm_read(os.path.basename(path)): - threads.add(executor.submit(read_fn, path)) + threads.add(executor.submit(self.read_swc, path)) # Store results swc_dicts = deque() @@ -283,22 +281,22 @@ def read_from_cloud(self, path): Dictionaries whose keys and values are the attribute names and values from an SWC file. """ - # Extact info + # Extract info assert util.is_s3_path(path) or util.is_gcs_path(path) use_s3 = util.is_s3_path(path) - # List filenames + # List paths swc_paths = util.list_cloud_paths(path, ".swc") zip_paths = util.list_cloud_paths(path, ".zip") # Call reader if swc_paths: - return self.read_swcs(swc_paths, self.read_swc) + return self.read_swcs(swc_paths) elif zip_paths: read_fn = self.read_s3_zip if use_s3 else self.read_gcs_zip return self.read_zips(zip_paths, read_fn) - - raise Exception(f"SWC Pointer is invalid {path}") + else: + return list() def read_gcs_swc(self, path): """ @@ -341,33 +339,14 @@ def read_gcs_zip(self, path): Dictionaries whose keys and values are the attribute names and values from an SWC file. """ - # Download ZIP - bucket_name, path = util.parse_cloud_path(path) + bucket_name, key = util.parse_cloud_path(path) bucket = storage.Client().bucket(bucket_name) try: - zip_content = bucket.blob(path).download_as_bytes() + zip_content = bucket.blob(key).download_as_bytes() except TransportError: - print(f"Failed to read {zip_path}!") + print(f"Failed to read {path}!") return deque() - - # Parse ZIP contents - swc_dicts = deque() - with ZipFile(BytesIO(zip_content), "r") as zf: - with ThreadPoolExecutor() as executor: - # Assign threads - threads = set() - for name in zf.namelist(): - if self.confirm_read(name): - threads.add( - executor.submit(self.read_zipped_swc, zf, name) - ) - - # Process results - for thread in as_completed(threads): - result = thread.result() - if result: - swc_dicts.append(result) - return swc_dicts + return self._parse_zip_bytes(zip_content) def read_s3_zip(self, path): """ @@ -385,29 +364,10 @@ def read_s3_zip(self, path): Dictionaries whose keys and values are the attribute names and values from an SWC file. """ - # Initialize cloud reader bucket, key = util.parse_cloud_path(path) s3 = boto3.client("s3", config=Config(signature_version=UNSIGNED)) zip_content = s3.get_object(Bucket=bucket, Key=key)["Body"].read() - - # Parse ZIP - with ZipFile(BytesIO(zip_content), "r") as zf: - with ThreadPoolExecutor() as executor: - # Assign threads - threads = set() - for name in zf.namelist(): - if self.confirm_read(name): - threads.add( - executor.submit(self.read_zipped_swc, zf, name) - ) - - # Store results - swc_dicts = deque() - for thread in as_completed(threads): - result = thread.result() - if result: - swc_dicts.append(result) - return swc_dicts + return self._parse_zip_bytes(zip_content) def confirm_read(self, path): """ @@ -427,23 +387,19 @@ def confirm_read(self, path): name = os.path.splitext(os.path.basename(path))[0] return name in self.swc_names if self.swc_names else True - # -- Process Text --- - def iterator(self, iterator): - """ - Gets an iterator that optionally displays a progress bar. - - Parameters - ---------- - iterator : iterable - Object to be iterated over. - - Returns - ------- - tqdm.tqdm - Iterator that is optionally wrapped in a progress bar. - """ - return tqdm(iterator, desc="Read SWCs") if self.verbose else iterator + def _parse_zip_bytes(self, zip_content): + with ZipFile(BytesIO(zip_content), "r") as zf: + names = [f for f in zf.namelist() if f.endswith(".swc")] + with ThreadPoolExecutor() as executor: + threads = { + executor.submit(self.read_zipped_swc, zf, name) + for name in names + } + return deque( + t.result() for t in as_completed(threads) if t.result() + ) + # -- Process Text --- def manual_progress_bar(self, total): """ Gets progress bar that needs to be updated manually. @@ -521,14 +477,15 @@ def process_content(self, content): offset = self.read_coordinate(parts[2:5]) if not line.startswith("#") and len(line.strip()) > 0: return content[i:], offset + return [], offset - def read_coordinate(self, xyz_str, offset=(0, 0, 0)): + def read_coordinate(self, coord_str, offset=(0, 0, 0)): """ Reads a coordinate from a string and converts it to voxel coordinates. Parameters ---------- - xyz_str : str + coord_str : str Coordinate stored as a string. offset : Tuple[int] Offset of coordinates in SWC file. Default is (0, 0, 0). @@ -538,4 +495,4 @@ def read_coordinate(self, xyz_str, offset=(0, 0, 0)): Tuple[int] xyz coordinates of an entry from an SWC file. """ - return [float(xyz_str[i]) + offset[i] for i in range(3)] + return [float(coord_str[i]) + offset[i] for i in range(3)] diff --git a/src/segmentation_skeleton_metrics/evaluate.py b/src/segmentation_skeleton_metrics/evaluate.py index 5220254..cc90c42 100644 --- a/src/segmentation_skeleton_metrics/evaluate.py +++ b/src/segmentation_skeleton_metrics/evaluate.py @@ -299,16 +299,14 @@ def save_fragments(self, gt_graphs, fragment_graphs): for key, graph in gt_graphs.items(): # Create zip writer zip_path = os.path.join(output_dir, f"{graph.name}.zip") - zip_writer = ZipFile(zip_path, "a") + zf = ZipFile(zip_path, "a") # Save skeletons - graph.to_zipped_swcs(zip_writer) - self.save_intersecting_fragments( - graph, fragment_graphs, zip_writer - ) + graph.to_zipped_swcs(zf) + self.save_intersecting_fragments(graph, fragment_graphs, zf) @staticmethod - def save_intersecting_fragments(gt_graph, fragment_graphs, zip_writer): + def save_intersecting_fragments(gt_graph, fragment_graphs, zf): """ Saves SWC files for all fragment graphs whose label intersects with the given ground-truth graph. @@ -319,13 +317,13 @@ def save_intersecting_fragments(gt_graph, fragment_graphs, zip_writer): Graphs built from ground truth SWC files. fragment_graphs : Dict[str, FragmentGraph] Graphs built from skeletons obtained from a segmentation. - zip_writer : zipfile.ZipFile + zf : zipfile.ZipFile Open ZIP file handle used to write fragments. """ intersecting_labels = gt_graph.node_labels() for key, graph in fragment_graphs.items(): if graph.label in intersecting_labels: - graph.to_zipped_swcs(zip_writer) + graph.to_zipped_swcs(zf) def save_merge_results(self, gt_graphs, fragment_graphs): """ @@ -339,39 +337,39 @@ def save_merge_results(self, gt_graphs, fragment_graphs): fragment_graphs : Dict[str, FragmentsGraph] Graphs built from skeletons obtained from a segmentation. """ - # Initialize a writer - filename = f"{self.prefix}fragments_with_merges.zip" - zip_path = os.path.join(self.output_dir, filename) - util.rm_file(zip_path) - zip_writer = ZipFile(zip_path, "a") - - # Save SWC files - self.save_merge_sites(zip_writer) - self.save_skeletons_with_merge(gt_graphs, fragment_graphs, zip_writer) - zip_writer.close() - - # Save CSV report - path = os.path.join(self.output_dir, f"{self.prefix}merge_sites.csv") - self.metrics["# Merges"].merge_sites.to_csv(path, index=True) - - def save_merge_sites(self, zip_writer): + if len(self.metrics["# Merges"].merge_sites) > 0: + # Initialize a writer + filename = f"{self.prefix}fragments_with_merges.zip" + zip_path = os.path.join(self.output_dir, filename) + util.rm_file(zip_path) + zf = ZipFile(zip_path, "a") + + # Save SWC files + self.save_merge_sites(zf) + self.save_skeletons_with_merge(gt_graphs, fragment_graphs, zf) + zf.close() + + # Save CSV report + filename = f"{self.prefix}merge_sites.csv" + path = os.path.join(self.output_dir, filename) + self.metrics["# Merges"].merge_sites.to_csv(path, index=True) + + def save_merge_sites(self, zf): """ Saves merge site coordinates into a ZIP archive. Parameters ---------- - zip_writer : zipfile.ZipFile + zf : zipfile.ZipFile Open ZIP file handle used to write merge sites. """ merge_sites = self.metrics["# Merges"].merge_sites for i in range(len(merge_sites)): filename = merge_sites.index[i] xyz = merge_sites["World"].iloc[i] - util.to_zipped_point(zip_writer, filename, xyz) + util.to_zipped_point(zf, filename, xyz) - def save_skeletons_with_merge( - self, gt_graphs, fragment_graphs, zip_writer - ): + def save_skeletons_with_merge(self, gt_graphs, fragment_graphs, zf): """ Saves ground truth and fragment skeletons containing merge sites into a ZIP archive. @@ -382,14 +380,19 @@ def save_skeletons_with_merge( Graphs built from ground truth SWC files. fragment_graphs : Dict[str, FragmentsGraph] Graphs built from skeletons obtained from a segmentation. - zip_writer : zipfile.ZipFile + zf : zipfile.ZipFile Open ZIP file handle used to write SWC data. """ # Save ground truth skeletons keys = self.metrics["# Merges"].merge_sites["GroundTruth_ID"] for key in keys.unique(): - gt_graphs[key].to_zipped_swcs(zip_writer) + gt_graphs[key].to_zipped_swcs(zf) # Save fragments + name2label = {g.name: g.label for g in fragment_graphs.values()} for key in self.metrics["# Merges"].fragments_with_merge: - fragment_graphs[key].to_zipped_swcs(zip_writer) + if key in fragment_graphs: + fragment_graphs[key].to_zipped_swcs(zf) + else: + label = name2label[key] + fragment_graphs[label].to_zipped_swcs(zf) diff --git a/src/segmentation_skeleton_metrics/skeleton_metrics.py b/src/segmentation_skeleton_metrics/skeleton_metrics.py index 666e6b6..cb031c5 100644 --- a/src/segmentation_skeleton_metrics/skeleton_metrics.py +++ b/src/segmentation_skeleton_metrics/skeleton_metrics.py @@ -497,7 +497,7 @@ def verify_site(self, gt_graph, fragment_graph, gt_node, fragment_node): xyz = fragment_graph.node_xyz(fragment_node) gt_graph.labels_with_merge.add(fragment_graph.label) - self.fragments_with_merge.add(fragment_graph.label) + self.fragments_with_merge.add(fragment_graph.name) self.merge_sites.append( { "Fragment_Name": fragment_graph.name, @@ -877,17 +877,21 @@ def __call__(self, gt_graphs, fragment_graphs, merge_sites): pair_to_length = dict() for i in self.get_iterator(merge_sites.index): # Extract site info - label = merge_sites["Label"][i] gt_id = merge_sites["GroundTruth_ID"][i] + label = merge_sites["Label"][i] + name = merge_sites["Fragment_Name"][i] pair_id = (label, gt_id) - # Check wheter to visit + # Check whether to visit if pair_id in pair_to_length: merge_sites.loc[i, self.name] = pair_to_length[pair_id] else: # Get graphs gt_graph = gt_graphs[gt_id] - fragment_graph = deepcopy(fragment_graphs[label]) + if label in fragment_graphs: + fragment_graph = deepcopy(fragment_graphs[label]) + elif name in fragment_graphs: + fragment_graph = deepcopy(fragment_graphs[name]) # Compute metric pair_to_length[pair_id] = self.compute_added_length(