Internal cleanup

PiperOrigin-RevId: 665400591

connectomics/volume/mask.py

@@ -20,6 +20,7 @@
 
 from connectomics.common import bounding_box
 from connectomics.volume import descriptor
+from connectomics.volume import metadata
 import dataclasses_json
 import numpy as np
 from scipy import ndimage
@@ -69,7 +70,7 @@ class MaskChannelConfig(dataclasses_json.DataClassJsonMixin):
 @dataclasses.dataclass
 class VolumeMaskOptions(dataclasses_json.DataClassJsonMixin):
   """Configuration for a volume mask."""
-  mask: descriptor.VolumeDescriptor
+  mask: metadata.DecoratedVolume
   channels: list[MaskChannelConfig] = dataclasses.field(default_factory=list)
 
 

connectomics/volume/metadata.py

@@ -26,10 +26,14 @@
 @dataclasses_json.dataclass_json
 @dataclasses.dataclass(frozen=True)
 class VolumeMetadata:
-  """Metadata associated with a Volume."""
-  # Volume size in voxels. XYZ order.
+  """Metadata associated with a Volume.
+
+  Attributes:
+    volume_size: Volume size in voxels. XYZ order.
+    pixel_size: Pixel size in nm. XYZ order.
+    bounding_boxes: Bounding boxes associated with the volume.
+  """
   volume_size: tuple[int, int, int]
-  # Pixel size in nm. XYZ order.
   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
@@ -71,6 +75,12 @@ def scale_xy(self, factor: float) -> 'VolumeMetadata':
 
 
 class Volume:
+  """Path to a volume with associated metadata.
+
+  Attributes:
+    path: The path to the volume.
+    meta: The volume metadata.
+  """
   path: pathlib.Path
   meta: VolumeMetadata
 
@@ -84,3 +94,18 @@ def save_metadata(self, kvdriver: str = 'file'):
         self.path.parent / f'{self.path.stem}.metadata.json',
         kvdriver=kvdriver,
     )
+
+
+@dataclasses_json.dataclass_json
+@dataclasses.dataclass(frozen=True)
+class DecoratedVolume:
+  """A volume with additional metadata.
+
+  Attributes:
+    path: The path to the volume.
+    decorator_specs: A JSON string of decorator specs.
+  """
+
+  path: pathlib.Path
+  # TODO(timblakely): This should be a list of DecoratorSpec dataclasses.
+  decorator_specs: str