Rename BaseVolume to Volume and support Metadata.

PiperOrigin-RevId: 666854804

connectomics/volume/base.py

@@ -21,6 +21,7 @@
 
 from connectomics.common import array
 from connectomics.common import bounding_box
+from connectomics.volume import metadata
 from connectomics.volume import subvolume
 import numpy as np
 
@@ -34,10 +35,10 @@ def slice_to_bbox(ind: array.CanonicalSlice) -> bounding_box.BoundingBox:
 
 class VolumeIndexer:
   """Interface for indexing supporting point lookups and slices."""
-  _volume: 'BaseVolume'
+  _volume: 'Volume'
   slices: array.CanonicalSlice
 
-  def __init__(self, volume: 'BaseVolume'):
+  def __init__(self, volume: 'Volume'):
     self._volume = volume
 
   def __getitem__(
@@ -64,9 +65,14 @@ def __getitem__(self, ind: array.IndexExpOrPointLookups) -> np.ndarray:
 
 # TODO(timblakely): Make generic-typed so it exposes both VolumeInfo and
 # Tensorstore via .descriptor.
-class BaseVolume:
+class Volume:
   """Common interface to multiple volume backends for Decorators."""
 
+  meta: metadata.VolumeMetadata
+
+  def __init__(self, meta: metadata.VolumeMetadata):
+    self.meta = meta
+
   def __getitem__(
       self, ind: array.IndexExpOrPointLookups) -> Union[np.ndarray, Subvolume]:
     return VolumeIndexer(self)[ind]
@@ -110,32 +116,33 @@ def write(self, subvol: subvolume.Subvolume):
   @property
   def volume_size(self) -> array.Tuple3i:
     """Volume size in voxels, XYZ."""
-    raise NotImplementedError
+    return self.meta.volume_size
 
   @property
-  def voxel_size(self) -> array.Tuple3f:
+  def pixel_size(self) -> array.Tuple3f:
     """Size of an individual voxels in physical dimensions (Nanometers)."""
-    raise NotImplementedError
+    return self.meta.pixel_size
 
   @property
   def shape(self) -> array.Tuple4i:
     """Shape of the volume in voxels, CZYX."""
-    raise NotImplementedError
+    return (self.meta.num_channels,) + self.volume_size[::-1]
 
   @property
   def ndim(self) -> int:
     """Number of dimensions in this volume."""
-    raise NotImplementedError
+    # TODO(timblakely): Support 3D volumes?
+    return 4
 
   @property
   def dtype(self) -> np.dtype:
     """Datatype of the underlying data."""
-    raise NotImplementedError
+    return self.meta.dtype
 
   @property
   def bounding_boxes(self) -> list[bounding_box.BoundingBox]:
     """List of bounding boxes contained in this volume."""
-    raise NotImplementedError
+    return self.meta.bounding_boxes
 
   @property
   def chunk_size(self) -> array.Tuple4i:
@@ -156,7 +163,7 @@ def clip_box_to_volume(
 
 
 def get_bounding_boxes_or_full(
-    volume: BaseVolume,
+    volume: Volume,
     bounding_boxes: Optional[Sequence[bounding_box.BoundingBoxBase]] = None,
     clip: bool = False,
 ) -> list[bounding_box.BoundingBox]:

connectomics/volume/base_test.py

@@ -15,37 +15,32 @@
 """Tests for base."""
 
 from absl.testing import absltest
-from connectomics.common import array
 from connectomics.common import bounding_box
 from connectomics.volume import base
+from connectomics.volume import metadata
 import numpy as np
 import numpy.testing as npt
 
 Box = bounding_box.BoundingBox
 
 
-class ShimVolume(base.BaseVolume):
+class ShimVolume(base.Volume):
 
-  def __init__(self, *args, **kwargs):
-    super(*args, **kwargs)
-    self.called = False
+  default_metadata = metadata.VolumeMetadata(
+      volume_size=(10, 11, 12),
+      pixel_size=(1, 2, 3),
+      bounding_boxes=[Box([0, 0, 0], [10, 20, 30])],
+      num_channels=1,
+      dtype=np.float32,
+  )
 
-  @property
-  def shape(self) -> array.Tuple4i:
-    return (1, 12, 11, 10)
+  def __init__(self):
+    super().__init__(self.default_metadata)
+    self.called = False
 
 
 class BaseVolumeTest(absltest.TestCase):
 
-  def test_not_implemented(self):
-    v = base.BaseVolume()
-
-    for field in [
-        'volume_size', 'voxel_size', 'shape', 'ndim', 'dtype', 'bounding_boxes'
-    ]:
-      with self.assertRaises(NotImplementedError):
-        _ = getattr(v, field)
-
   def test_get_points(self):
     tself = self
 

connectomics/volume/descriptor.py

@@ -60,7 +60,7 @@ def load_descriptor(spec: Union[str, VolumeDescriptor]) -> VolumeDescriptor:
 
 def open_descriptor(
     spec: Union[str, VolumeDescriptor],
-    context: Optional[dict[str, Any]] = None) -> base.BaseVolume:
+    context: Optional[dict[str, Any]] = None) -> base.Volume:
   """Open a volume from a volume descriptor.
 
   Args:

connectomics/volume/metadata.py

@@ -21,6 +21,8 @@
 from connectomics.common import bounding_box
 from connectomics.common import file
 import dataclasses_json
+import numpy as np
+import numpy.typing as npt
 
 
 @dataclasses_json.dataclass_json
@@ -32,13 +34,20 @@ class VolumeMetadata:
     volume_size: Volume size in voxels. XYZ order.
     pixel_size: Pixel size in nm. XYZ order.
     bounding_boxes: Bounding boxes associated with the volume.
+    num_channels: Number of channels in the volume.
+    dtype: Datatype of the volume. Must be numpy compatible.
   """
   volume_size: tuple[int, int, int]
   pixel_size: tuple[float, float, float]
   bounding_boxes: list[bounding_box.BoundingBox]
-  # TODO(timblakely): In the event we want to enforce the assumption that volumes
-  # are XYZC (i.e. processing happens differently for spatial and channel axes),
-  # add num_channels to this class to record any changes in channel counts.
+  num_channels: int = 1
+  dtype: npt.DTypeLike = dataclasses.field(
+      metadata=dataclasses_json.config(
+          decoder=np.dtype,
+          encoder=lambda x: np.dtype(x).name,
+      ),
+      default=np.uint8,
+  )
 
   def scale(
       self, scale_factors: float | Sequence[float]

connectomics/volume/metadata_test.py

@@ -20,6 +20,7 @@
 from absl.testing import absltest
 from connectomics.common import bounding_box
 from connectomics.volume import metadata
+import numpy as np
 
 
 FLAGS = flags.FLAGS
@@ -83,6 +84,7 @@ def test_volume_save_metadata(self):
         volume_size=(100, 100, 100),
         pixel_size=(8, 8, 30),
         bounding_boxes=[BBOX([10, 10, 10], [100, 100, 100])],
+        dtype=np.uint64,
     )
     temp_path = pathlib.Path(self.create_tempdir().full_path)
     vol = metadata.Volume(path=temp_path / 'foo.volinfo', meta=meta)

connectomics/volume/tensorstore.py

@@ -59,7 +59,7 @@ class TensorstoreConfig(utils.NPDataClassJsonMixin):
           decoder=file.dataclass_loader(TensorstoreMetadata)))
 
 
-class TensorstoreVolume(base.BaseVolume):
+class TensorstoreVolume(base.Volume):
   """Tensorstore-backed Volume."""
 
   _store: ts.TensorStore

connectomics/volume/tsv_decorator.py

@@ -34,7 +34,7 @@
 class DecoratorFactory:
   """Constructs a VolumeDecorator based on a name and arguments."""
 
-  def make_decorator(self, wrapped_volume: base.BaseVolume, name: str,
+  def make_decorator(self, wrapped_volume: base.Volume, name: str,
                      *args: list[Any],
                      **kwargs: dict[str, Any]) -> 'VolumeDecorator':
     raise NotImplementedError()
@@ -43,14 +43,14 @@ def make_decorator(self, wrapped_volume: base.BaseVolume, name: str,
 class GlobalsDecoratorFactory:
   """Loads VolumeDecorators from globals()."""
 
-  def make_decorator(self, wrapped_volume: base.BaseVolume, name: str,
+  def make_decorator(self, wrapped_volume: base.Volume, name: str,
                      *args: list[Any],
                      **kwargs: dict[str, Any]) -> 'VolumeDecorator':
     decorator_ctor = globals()[name]
     return decorator_ctor(wrapped_volume, *args, **kwargs)
 
 
-def from_specs(volume: base.BaseVolume,
+def from_specs(volume: base.Volume,
                specs: list[DecoratorSpec],
                decorator_factory: Optional[DecoratorFactory] = None):
   """Decorates the given volume from the given specs.
@@ -85,12 +85,12 @@ def from_specs(volume: base.BaseVolume,
   return volume
 
 
-class VolumeDecorator(base.BaseVolume):
+class VolumeDecorator(base.Volume):
   """Delegates to wrapped volumes, optionally applying transforms."""
 
-  wrapped: base.BaseVolume
+  wrapped: base.Volume
 
-  def __init__(self, wrapped: base.BaseVolume):
+  def __init__(self, wrapped: base.Volume):
     self._wrapped = wrapped
 
   def get_points(self, points: array.PointLookups) -> np.ndarray:
@@ -105,7 +105,7 @@ def volume_size(self) -> array.Tuple3i:
 
   @property
   def voxel_size(self) -> array.Tuple3f:
-    return self._wrapped.voxel_size
+    return self._wrapped.pixel_size
 
   @property
   def shape(self) -> array.Tuple4i:
@@ -133,7 +133,7 @@ class Upsample(VolumeDecorator):
 
   scale_zyx: np.ndarray
 
-  def __init__(self, wrapped: base.BaseVolume, scale: array.ArrayLike3d):
+  def __init__(self, wrapped: base.Volume, scale: array.ArrayLike3d):
     """Initializes the wrapper.
 
     Args:
@@ -154,7 +154,7 @@ def volume_size(self) -> array.Tuple3i:
 
   @property
   def voxel_size(self) -> array.Tuple3i:
-    return tuple(self._wrapped.voxel_size / self.scale_zyx[::-1])
+    return tuple(self._wrapped.pixel_size / self.scale_zyx[::-1])
 
   @property
   def shape(self) -> array.Tuple4i:

connectomics/volume/tsv_decorator_test.py

@@ -15,15 +15,17 @@
 """Tests for tsv_decorator."""
 
 import typing
-from typing import Any, Sequence, Tuple
+from typing import Any
 
 from absl.testing import absltest
 from connectomics.common import array
 from connectomics.common import bounding_box
 from connectomics.volume import base as base_volume
 from connectomics.volume import descriptor as vd
+from connectomics.volume import metadata
 from connectomics.volume import tsv_decorator
 import numpy as np
+import numpy.typing as nptyping
 import numpy.testing as npt
 
 BBox = bounding_box.BoundingBox
@@ -33,13 +35,24 @@
 
 # TODO(timblakely): Create an common in-memory volume implementation. Would be
 # useful in both tests and in temporary volume situations.
-class DummyVolume(base_volume.BaseVolume):
-
-  def __init__(self, volume_size: Sequence[int], voxel_size: Sequence[int],
-               bounding_boxes: list[BBox], data: np.ndarray):
-    self._volume_size = tuple(volume_size)
-    self._voxel_size = tuple(voxel_size)
-    self._bounding_boxes = bounding_boxes
+class DummyVolume(base_volume.Volume):
+
+  def __init__(
+      self,
+      volume_size: tuple[int, int, int],
+      voxel_size: tuple[int, int, int],
+      bounding_boxes: list[BBox],
+      data: np.ndarray,
+      dtype: nptyping.DTypeLike,
+  ):
+    super().__init__(
+        metadata.VolumeMetadata(
+            volume_size=volume_size,
+            pixel_size=voxel_size,
+            bounding_boxes=bounding_boxes,
+            dtype=dtype,
+        )
+    )
     self._data = data
 
   def __getitem__(self, ind):
@@ -56,39 +69,17 @@ def get_points(self, points: array.PointLookups) -> np.ndarray:
   def get_slices(self, slices: array.CanonicalSlice) -> np.ndarray:
     return self._data[slices]
 
-  @property
-  def volume_size(self) -> array.Tuple3i:
-    return self._volume_size
-
-  @property
-  def voxel_size(self) -> array.Tuple3i:
-    return self._voxel_size
-
-  @property
-  def shape(self) -> array.Tuple4i:
-    return (1,) + tuple(self._volume_size[::-1])
-
-  @property
-  def ndim(self) -> int:
-    return len(self._data.shape)
-
-  @property
-  def dtype(self) -> np.dtype:
-    return self._data.dtype
-
-  @property
-  def bounding_boxes(self) -> list[BBox]:
-    return self._bounding_boxes
-
 
-def _make_dummy_vol() -> Tuple[DummyVolume, BBox, np.ndarray]:
+def _make_dummy_vol() -> tuple[DummyVolume, BBox, np.ndarray]:
   bbox = BBox([100, 200, 300], [20, 50, 100])
   data = np.zeros(bbox.size)
   data[0] = 1
   data = np.cumsum(
       np.cumsum(np.cumsum(data, axis=0), axis=1), axis=2, dtype=np.uint64)
   data = data[np.newaxis]
-  vol = DummyVolume([3000, 2000, 1000], (8, 8, 33), [bbox], data)
+  vol = DummyVolume(
+      (3000, 2000, 1000), (8, 8, 33), [bbox], data, dtype=np.uint64
+  )
   return vol, bbox, data
 
 
@@ -99,7 +90,7 @@ def test_dummy_volume(self):
     self.assertEqual((3000, 2000, 1000), vol.volume_size)
     self.assertLen(vol.bounding_boxes, 1)
     self.assertEqual([bbox], vol.bounding_boxes)
-    self.assertEqual((8, 8, 33), vol.voxel_size)
+    self.assertEqual((8, 8, 33), vol.pixel_size)
     self.assertEqual(np.uint64, vol.dtype)
     self.assertEqual(4, vol.ndim)
     self.assertEqual((1, 1000, 2000, 3000), vol.shape)
@@ -137,7 +128,7 @@ class CustomDecoratorFactory(tsv_decorator.DecoratorFactory):
   def __init__(self, *args, **kwargs):
     self.called = False
 
-  def make_decorator(self, wrapped_volume: base_volume.BaseVolume, name: str,
+  def make_decorator(self, wrapped_volume: base_volume.Volume, name: str,
                      *args: list[Any],
                      **kwargs: dict[str, Any]) -> tsv_decorator.VolumeDecorator:
     if name == 'CustomDecorator':