diff --git a/python/neuroglancer/write_annotations.py b/python/neuroglancer/write_annotations.py index 92b78d0a8..c00cd82f3 100644 --- a/python/neuroglancer/write_annotations.py +++ b/python/neuroglancer/write_annotations.py @@ -22,9 +22,9 @@ import struct from collections.abc import Sequence from typing import Literal, NamedTuple, Optional, Union, cast -from logging import warning import tensorstore as ts import numpy as np +import math from . import coordinate_space, viewer_state @@ -102,6 +102,55 @@ def choose_output_spec(total_count, total_bytes, return options +def compressed_morton_code(gridpt, grid_size): + # from cloudvolume + if hasattr(gridpt, "__len__") and len(gridpt) == 0: # generators don't have len + return np.zeros((0,), dtype=np.uint32) + + gridpt = np.asarray(gridpt, dtype=np.uint32) + single_input = False + if gridpt.ndim == 1: + gridpt = np.atleast_2d(gridpt) + single_input = True + + code = np.zeros((gridpt.shape[0],), dtype=np.uint64) + num_bits = [ math.ceil(math.log2(size)) for size in grid_size ] + j = np.uint64(0) + one = np.uint64(1) + + if sum(num_bits) > 64: + raise ValueError(f"Unable to represent grids that require more than 64 bits. Grid size {grid_size} requires {num_bits} bits.") + + max_coords = np.max(gridpt, axis=0) + if np.any(max_coords >= grid_size): + raise ValueError(f"Unable to represent grid points larger than the grid. Grid size: {grid_size} Grid points: {gridpt}") + + for i in range(max(num_bits)): + for dim in range(3): + if 2 ** i < grid_size[dim]: + bit = (((np.uint64(gridpt[:, dim]) >> np.uint64(i)) & one) << j) + code |= bit + j += one + print(gridpt, grid_size, code) + if single_input: + return code[0] + return code + +# def compressed_morton_code(position, shape): +# output_bit = 0 +# rank = len(position) +# output_num = 0 +# for bit in range(32): +# for dim in range(rank-1, -1, -1): +# if (shape[dim] - 1) >> bit: +# output_num |= ((position[dim] >> bit) & 1) << output_bit +# output_bit += 1 +# if output_bit == 64: +# # In Python, we don't have the 32-bit limitation, so we don't need to split into high and low. +# # But you can add code here to handle or signal overflow if needed. +# pass +# return output_num + def _get_dtype_for_geometry(annotation_type: AnnotationType, rank: int): geometry_size = rank if annotation_type == "point" else 2 * rank return [("geometry", "