From afbc23df785c203fc3e5c355831a60323e9fc960 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sun, 18 Feb 2024 08:34:23 -0800 Subject: [PATCH] generlizing spatial bins with rtree --- python/neuroglancer/write_annotations.py | 37 +++++++++++++++++++----- 1 file changed, 30 insertions(+), 7 deletions(-) diff --git a/python/neuroglancer/write_annotations.py b/python/neuroglancer/write_annotations.py index 146727d99..3fd5dab52 100644 --- a/python/neuroglancer/write_annotations.py +++ b/python/neuroglancer/write_annotations.py @@ -20,9 +20,11 @@ import struct from collections import defaultdict from collections.abc import Sequence +from itertools import product from typing import Literal, NamedTuple, Optional, SupportsInt, Union, cast import math import numpy as np +import rtree try: import tensorstore as ts @@ -57,7 +59,6 @@ class Annotation(NamedTuple): id: int encoded: bytes relationships: Sequence[Sequence[int]] - geometry: Sequence[float] _PROPERTY_DTYPES: dict[ @@ -289,6 +290,9 @@ def __init__( shape=(self.rank,), fill_value=float("-inf"), dtype=np.float32 ) self.related_annotations = [{} for _ in self.relationships] + p = rtree.index.Property() + p.dimension = self.rank + self.rtree = rtree.index.Index(properties=p) def get_chunk_index(self, coords): return tuple((coords // self.chunk_size).astype(np.int32)) @@ -420,9 +424,15 @@ def _add_obj( id = len(self.annotations) annotation = Annotation( - id=id, encoded=encoded.tobytes(), relationships=related_ids, geometry=coords + id=id, encoded=encoded.tobytes(), relationships=related_ids ) + spatial_points = coords[: n_spatial_coords * self.rank] + spatial_points = np.reshape(spatial_points, (self.rank, n_spatial_coords)) + lower_bound = np.min(spatial_points, axis=1) + upper_bound = np.max(spatial_points, axis=1) + self.rtree.insert(id, tuple(lower_bound) + tuple(upper_bound), obj=annotation) + # for i in range(int(n_spatial_coords)): # chunk_index = self.get_chunk_index( # np.array(coords[i * self.rank : (i + 1) * self.rank]) @@ -550,11 +560,24 @@ def write(self, path: Union[str, pathlib.Path]): os.makedirs(os.path.join(path, f"rel_{relationship}"), exist_ok=True) os.makedirs(os.path.join(path, "by_id"), exist_ok=True) os.makedirs(os.path.join(path, "spatial0"), exist_ok=True) - for ann in self.annotations: - if self.annotation_type == "point": - # get the first self.rank elements of the geometry array - chunk_index = self.get_chunk_index(ann.geometry - self.lower_bound) - self.annotations_by_chunk[chunk_index].append(ann) + + # Generate all combinations of coordinates + coordinates = product(*(range(n) for n in num_chunks)) + + # Iterate over the grid + for cell in coordinates: + # Query the rtree index for annotations in the current chunk + lower_bound = self.lower_bound + np.array(cell) * self.chunk_size + upper_bound = lower_bound + self.chunk_size + coords = np.concatenate((lower_bound, upper_bound)) + chunk_annotations = self.rtree.intersection(tuple(coords), objects="raw") + self.annotations_by_chunk[cell] = list(chunk_annotations) + + # for ann in self.annotations: + # if self.annotation_type == "point": + # # get the first self.rank elements of the geometry array + # chunk_index = self.get_chunk_index(ann.geometry - self.lower_bound) + # self.annotations_by_chunk[chunk_index].append(ann) total_chunks = len(self.annotations_by_chunk) spatial_sharding_spec = choose_output_spec(