feat: precomp attrs and write_annotations refinements

tests/unit/db_annotations/test_precomp_annotations.py

@@ -18,11 +18,11 @@ def test_round_trip():
     file_dir = os.path.join(temp_dir, "round_trip")
 
     lines = [
-        LineAnnotation(line_id=1, start=(1640.0, 1308.0, 61.0), end=(1644.0, 1304.0, 57.0)),
-        LineAnnotation(line_id=2, start=(1502.0, 1709.0, 589.0), end=(1498.0, 1701.0, 589.0)),
-        LineAnnotation(line_id=3, start=(254.0, 68.0, 575.0), end=(258.0, 62.0, 575.0)),
-        LineAnnotation(line_id=4, start=(1061.0, 657.0, 507.0), end=(1063.0, 653.0, 502.0)),
-        LineAnnotation(line_id=5, start=(1298.0, 889.0, 315.0), end=(1295.0, 887.0, 314.0)),
+        LineAnnotation(id=1, start=(1640.0, 1308.0, 61.0), end=(1644.0, 1304.0, 57.0)),
+        LineAnnotation(id=2, start=(1502.0, 1709.0, 589.0), end=(1498.0, 1701.0, 589.0)),
+        LineAnnotation(id=3, start=(254.0, 68.0, 575.0), end=(258.0, 62.0, 575.0)),
+        LineAnnotation(id=4, start=(1061.0, 657.0, 507.0), end=(1063.0, 653.0, 502.0)),
+        LineAnnotation(id=5, start=(1298.0, 889.0, 315.0), end=(1295.0, 887.0, 314.0)),
     ]
     # Note: line 2 above, with the chunk_sizes below, will span 2 chunks, and so will
     # be written out to both of them.
@@ -74,8 +74,8 @@ def test_round_trip():
 
     # Test replacing only the two lines in that bounds.
     new_lines = [
-        LineAnnotation(line_id=104, start=(1061.5, 657.0, 507.0), end=(1062.5, 653.0, 502.0)),
-        LineAnnotation(line_id=105, start=(1298.5, 889.0, 315.0), end=(1294.5, 887.0, 314.0)),
+        LineAnnotation(id=104, start=(1061.5, 657.0, 507.0), end=(1062.5, 653.0, 502.0)),
+        LineAnnotation(id=105, start=(1298.5, 889.0, 315.0), end=(1294.5, 887.0, 314.0)),
     ]
     sf.write_annotations(new_lines, clearing_bbox=roi)
     lines_read = sf.read_in_bounds(roi, strict=False)
@@ -110,7 +110,7 @@ def test_resolution_changes():
     sf.clear()
 
     # writing with voxel size 10, 10, 80
-    lines = [LineAnnotation(line_id=1, start=(100, 500, 50), end=(200, 600, 60))]
+    lines = [LineAnnotation(id=1, start=(100, 500, 50), end=(200, 600, 60))]
     sf.write_annotations(lines, Vec3D(10, 10, 80))
 
     # pull those back out at file native resolution, i.e. (20, 20, 40)
@@ -132,11 +132,11 @@ def test_single_level():
     file_dir = os.path.join(temp_dir, "single_level")
 
     lines = [
-        LineAnnotation(line_id=1, start=(1640.0, 1308.0, 61.0), end=(1644.0, 1304.0, 57.0)),
-        LineAnnotation(line_id=2, start=(1502.0, 1709.0, 589.0), end=(1498.0, 1701.0, 589.0)),
-        LineAnnotation(line_id=3, start=(254.0, 68.0, 575.0), end=(258.0, 62.0, 575.0)),
-        LineAnnotation(line_id=4, start=(1061.0, 657.0, 507.0), end=(1063.0, 653.0, 502.0)),
-        LineAnnotation(line_id=5, start=(1298.0, 889.0, 315.0), end=(1295.0, 887.0, 314.0)),
+        LineAnnotation(id=1, start=(1640.0, 1308.0, 61.0), end=(1644.0, 1304.0, 57.0)),
+        LineAnnotation(id=2, start=(1502.0, 1709.0, 589.0), end=(1498.0, 1701.0, 589.0)),
+        LineAnnotation(id=3, start=(254.0, 68.0, 575.0), end=(258.0, 62.0, 575.0)),
+        LineAnnotation(id=4, start=(1061.0, 657.0, 507.0), end=(1063.0, 653.0, 502.0)),
+        LineAnnotation(id=5, start=(1298.0, 889.0, 315.0), end=(1295.0, 887.0, 314.0)),
     ]
     # Note: line 2 above, with the chunk_sizes below, will span 2 chunks, and so will
     # be written out to both of them.
@@ -157,6 +157,14 @@ def test_single_level():
     chunk_path = os.path.join(file_dir, "spatial0", "2_1_1")
     assert precomp_annotations.count_lines_in_file(chunk_path) == 2
 
+    with pytest.raises(ValueError):
+        out_of_bounds_lines = [
+            LineAnnotation(id=1, start=(1640.0, 1308.0, 61.0), end=(1644.0, 1304.0, 57.0)),
+            LineAnnotation(id=666, start=(-100, 0, 0), end=(50, 50, 50)),
+        ]
+        roi = BBox3D.from_coords((25, 25, 25), (250, 250, 250), resolution=(10, 10, 40))
+        sf.write_annotations(out_of_bounds_lines, clearing_bbox=roi)
+
 
 def test_edge_cases():
     with pytest.raises(ValueError):

web_api/app/precomputed_annotations.py

@@ -69,20 +69,22 @@ async def add_multiple(
     """
     The PUT endpoint replaces all data within the given bounds in the given file
     (which may or may not exist yet) with the given new set of lines.
+
+    Note that any lines not entirely contained within the bbox are ignored.
     """
-    lines = []
-    for entry in annotations:
-        annotation = AnnotationDBEntry.from_dict(entry["id"], entry)
-        line = annotation.ng_annotation
-        lines.append(
-            precomp_annotations.LineAnnotation(
-                int(annotation.id),
-                line.point_a.tolist(),
-                line.point_b.tolist(),
-            )
-        )
     resolution_vec = Vec3D(*resolution)
     index = VolumetricIndex.from_coords(bbox_start, bbox_end, resolution_vec)
+    lines = []
+    for d in annotations:
+        db_entry = AnnotationDBEntry.from_dict(d["id"], d)
+        ng_annotation = db_entry.ng_annotation
+        line_annotation = precomp_annotations.LineAnnotation(
+            int(db_entry.id),
+            ng_annotation.point_a.tolist(),
+            ng_annotation.point_b.tolist(),
+        )
+        if line_annotation.in_bounds(index, strict=True):
+            lines.append(line_annotation)
     layer = build_annotation_layer(path, index=index, mode="replace")
     layer.write_annotations(lines, annotation_resolution=resolution_vec, clearing_bbox=index.bbox)
 

zetta_utils/db_annotations/precomp_annotations.py

@@ -18,8 +18,9 @@
 import struct
 from math import ceil
 from random import shuffle
-from typing import IO, Literal, Optional, Sequence
+from typing import IO, Any, Literal, Optional, Sequence, Tuple
 
+import attrs
 from cloudfiles import CloudFile, CloudFiles
 
 from zetta_utils import builder, log, mazepa
@@ -47,38 +48,27 @@ def path_join(*paths: str):
         return os.path.join(*paths)
 
 
-class LineAnnotation:
-    BYTES_PER_ENTRY = 24  # start (3 floats), end (3 floats)
-
-    def __init__(self, line_id: int, start: Sequence[float], end: Sequence[float]):
-        """
-        Initialize a LineAnnotation instance.
+def to_3_tuple(value: Any) -> Tuple[float, ...]:
+    # Ensure the value is a 3-element tuple of floats
+    if len(value) != 3:
+        raise ValueError("3-element sequence expected")
+    return tuple(float(x) for x in value)
 
-        :param id: An integer representing the ID of the annotation.
-        :param start: A tuple of three floats representing the start coordinate (x, y, z).
-        :param end: A tuple of three floats representing the end coordinate (x, y, z).
 
-        Coordinates are in units defined by "dimensions" in the info file, or some
-        other resolution specified by the user.  (Like a Vec3D, context is needed to
-        interpret these coordinates.)
-        """
-        self.start = start
-        self.end = end
-        self.id = line_id
+@attrs.define
+class LineAnnotation:
+    """
+    LineAnnotation represents a Neuroglancer line annotation.  Start and end
+    points are in voxels -- i.e., the coordinates are in units defined by
+    "dimensions" in the info file, or some other resolution specified by the
+    user.  (Like a Vec3D, context is needed to interpret these coordinates.)
+    """
 
-    def __repr__(self):
-        """
-        Return a string representation of the LineAnnotation instance.
-        """
-        return f"LineAnnotation(line_id={self.id}, start={self.start}, end={self.end})"
+    BYTES_PER_ENTRY = 24  # start (3 floats), end (3 floats)
 
-    def __eq__(self, other):
-        return (
-            isinstance(other, LineAnnotation)
-            and self.id == other.id
-            and self.start == other.start
-            and self.end == other.end
-        )
+    id: int
+    start: Tuple[float, ...] = attrs.field(converter=to_3_tuple)
+    end: Tuple[float, ...] = attrs.field(converter=to_3_tuple)
 
     def write(self, output: IO[bytes]):
         """
@@ -100,11 +90,14 @@ def read(in_stream: IO[bytes]):
             struct.unpack("<3f", in_stream.read(12)),
         )
 
-    def in_bounds(self, bounds: VolumetricIndex):
+    def in_bounds(self, bounds: VolumetricIndex, strict=False):
         """
-        Return whether either end of this line is in the given bounds.
-        (Assumes our coordinates match that of the given VolumetricIndex.)
+        Check whether this line at all crosses (when strict=False), or
+        is entirely within (when strict=True), the given VolumetricIndex.
+        (Assumes our coordinates match that of the index.)
         """
+        if strict:
+            return bounds.contains(self.start) and bounds.contains(self.end)
         return bounds.line_intersects(self.start, self.end)
 
     def convert_coordinates(self, from_res: Vec3D, to_res: Vec3D):
@@ -419,6 +412,7 @@ def subdivide(data, bounds: VolumetricIndex, chunk_sizes, write_to_dir=None, lev
 
 
 @builder.register("AnnotationLayer")
+@attrs.define(init=False)
 class AnnotationLayer:
     """
     This class represents a spatial (precomputed annotation) file.  It knows its data
@@ -433,6 +427,10 @@ class AnnotationLayer:
     have to worry about it.
     """
 
+    index: VolumetricIndex
+    chunk_sizes: Sequence[Sequence[int]]
+    path: str = ""
+
     def __init__(
         self,
         path: str,
@@ -474,12 +472,6 @@ def __init__(
         self.index = index
         self.chunk_sizes = chunk_sizes
 
-    def __repr__(self):
-        return (
-            f"AnnotationLayer(path='{self.path}', index={self.index}, "
-            f"chunk_sizes={self.chunk_sizes})"
-        )
-
     def exists(self) -> bool:
         """
         Return whether this spatial file (more specifically, its info file) already exists.
@@ -561,7 +553,7 @@ def write_annotations(
         annotation_resolution: Optional[Vec3D] = None,
         all_levels: bool = True,
         clearing_bbox: Optional[BBox3D] = None,
-    ):
+    ):  # pylint: disable=too-many-branches
         """
         Write a set of line annotations to the file, adding to any already there.
 
@@ -570,7 +562,8 @@ def write_annotations(
         if not specified, assumes native coordinates (i.e. self.index.resolution)
         :param all_levels: if true, write to all spatial levels (chunk sizes).
             If false, write only to the lowest level (smallest chunks).
-        :param clearing_bbox: if given, clear any existing data within these bounds.
+        :param clearing_bbox: if given, clear any existing data within these bounds;
+            annotations must be entirely within these bounds.
         """
         if not annotations:
             logger.info("write_annotations called with 0 annotations to write")
@@ -591,6 +584,8 @@ def write_annotations(
                 round(clearing_bbox.end / self.index.resolution),
                 self.index.resolution,
             )
+            if not all(map(lambda x: x.in_bounds(clearing_idx, strict=True), annotations)):
+                raise ValueError("All annotations must be strictly within clearing_bbox")
 
         for level in levels:
             limit = 0
@@ -683,7 +678,9 @@ def write_annotations(
                         old_data = read_lines(anno_file_path)
                         if clearing_idx:
                             old_data = list(
-                                filter(lambda d: not d.in_bounds(clearing_idx), old_data)
+                                filter(
+                                    lambda d: not d.in_bounds(clearing_idx, strict=True), old_data
+                                )
                             )
                         chunk_data += old_data
                         limit = max(limit, len(chunk_data))