kernel: bbox/polygon geometry as a pydantic discriminated union (replaces Annotation.geometry: dict)

src/visionset/kernel/domain/__init__.py

@@ -11,6 +11,12 @@
 from visionset.kernel.domain.asset import Asset
 from visionset.kernel.domain.batch import Batch, BatchState
 from visionset.kernel.domain.dataset import Dataset
+from visionset.kernel.domain.geometry import (
+    BboxGeometry,
+    ClassificationGeometry,
+    Geometry,
+    PolygonGeometry,
+)
 from visionset.kernel.domain.ingest import IngestJob, IngestState
 from visionset.kernel.domain.project import Project
 from visionset.kernel.domain.release import Manifest, Release
@@ -29,12 +35,16 @@
     "Attribute",
     "Batch",
     "BatchState",
+    "BboxGeometry",
+    "ClassificationGeometry",
     "Dataset",
+    "Geometry",
     "GeometryType",
     "IngestJob",
     "IngestState",
     "LabelClass",
     "Manifest",
+    "PolygonGeometry",
     "Project",
     "Provenance",
     "Release",

src/visionset/kernel/domain/annotation.py

@@ -6,6 +6,8 @@
 
 from pydantic import BaseModel, Field, model_validator
 
+from visionset.kernel.domain.geometry import Geometry
+
 Provenance = Literal["human", "model", "import"]
 
 
@@ -24,7 +26,7 @@ class Annotation(BaseModel):
     asset_id: UUID
     label_class: str
     schema_version: int = Field(ge=1)
-    geometry: dict[str, object]  # refined into a discriminated union in a later session
+    geometry: Geometry
     provenance: Provenance
     model_ref: str | None = None
     confidence: float | None = Field(default=None, ge=0.0, le=1.0)

src/visionset/kernel/domain/geometry.py

@@ -0,0 +1,83 @@
+# usage: from visionset.kernel.domain import Geometry, BboxGeometry, PolygonGeometry
+"""Annotation geometry as a discriminated union.
+
+The discriminator is the ``type`` field, and its values ARE ``GeometryType``
+members — not parallel string literals. That is deliberate: a schema declares
+which geometries a class allows in terms of ``GeometryType``, so validating an
+annotation against it is a plain membership test on ``geometry.type``, with no
+translation layer in between.
+
+Adding a geometry (polyline, keypoints, mask, the 3D variants) means defining a
+model whose ``type`` is the matching ``GeometryType`` member and appending it to
+the ``Geometry`` union. Nothing about the discriminator changes shape, and no
+existing payload stops parsing. ``GeometryType`` names eight geometries; three
+are implemented here — the rest are roadmap, and a payload naming one is
+rejected until its model exists.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from visionset.kernel.domain.schema import GeometryType
+
+
+class BboxGeometry(BaseModel):
+    """An axis-aligned rectangle: top-left corner plus size.
+
+    ``width`` and ``height`` must be strictly positive — a zero-area box is as
+    meaningless as a negative one, so neither is accepted. ``x`` and ``y`` are
+    unconstrained: an annotation may legitimately start outside the asset's
+    bounds when an object is clipped by the frame edge.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    type: Literal[GeometryType.BBOX] = GeometryType.BBOX
+    x: float
+    y: float
+    width: float = Field(gt=0.0)
+    height: float = Field(gt=0.0)
+
+
+class PolygonGeometry(BaseModel):
+    """A closed polygon, as at least three ``(x, y)`` vertices.
+
+    The closing edge is implicit: the last point joins the first, and repeating
+    the first point at the end is NOT expected. Self-intersection is not
+    validated — M1 accepts any ring of three or more points, and rejecting
+    degenerate shapes is left to a later milestone.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    type: Literal[GeometryType.POLYGON] = GeometryType.POLYGON
+    points: list[tuple[float, float]] = Field(min_length=3)
+
+
+class ClassificationGeometry(BaseModel):
+    """A whole-asset tag: the annotation carries a class but no coordinates.
+
+    It exists as a variant rather than as ``geometry: None`` so that every
+    annotation has a geometry with a discriminator, and so the union stays the
+    single place that answers "what shape is this label?".
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    type: Literal[GeometryType.CLASSIFICATION_TAG] = GeometryType.CLASSIFICATION_TAG
+
+
+Geometry = Annotated[
+    BboxGeometry | PolygonGeometry | ClassificationGeometry,
+    Field(discriminator="type"),
+]
+"""Every geometry an Annotation can carry.
+
+Coordinates are ALWAYS floats in the asset's native reference frame — pixels for
+images — and are NEVER normalized. Normalization to a [0, 1] range, or to any
+other convention a format demands, is the exporter's concern and happens at the
+boundary, never in the domain.
+"""

tests/kernel/test_annotation.py

@@ -11,7 +11,7 @@ def _make(**overrides: object) -> Annotation:
         "asset_id": uuid4(),
         "label_class": "car",
         "schema_version": 1,
-        "geometry": {"type": "bbox", "x": 1.0, "y": 2.0, "w": 10.0, "h": 20.0},
+        "geometry": {"type": "bbox", "x": 1.0, "y": 2.0, "width": 10.0, "height": 20.0},
         "provenance": "human",
     }
     data.update(overrides)

tests/kernel/test_geometry.py

@@ -0,0 +1,127 @@
+from uuid import uuid4
+
+import pytest
+from pydantic import TypeAdapter, ValidationError
+
+from visionset.kernel.domain import (
+    Annotation,
+    BboxGeometry,
+    ClassificationGeometry,
+    Geometry,
+    GeometryType,
+    LabelClass,
+    PolygonGeometry,
+)
+
+geometry_adapter: TypeAdapter[Geometry] = TypeAdapter(Geometry)
+
+VARIANTS = [
+    BboxGeometry(x=1.0, y=2.0, width=10.0, height=20.0),
+    PolygonGeometry(points=[(0.0, 0.0), (10.0, 0.0), (10.0, 10.0)]),
+    ClassificationGeometry(),
+]
+
+
+def _annotation(geometry: Geometry) -> Annotation:
+    return Annotation(
+        asset_id=uuid4(),
+        label_class="car",
+        schema_version=1,
+        geometry=geometry,
+        provenance="human",
+    )
+
+
+def test_each_variant_round_trips_through_json_unchanged() -> None:
+    for geometry in VARIANTS:
+        payload = geometry_adapter.dump_json(geometry)
+        assert geometry_adapter.dump_json(geometry_adapter.validate_json(payload)) == payload
+
+
+def test_each_variant_round_trips_nested_in_an_annotation() -> None:
+    for geometry in VARIANTS:
+        annotation = _annotation(geometry)
+        payload = annotation.model_dump_json()
+        rehydrated = Annotation.model_validate_json(payload)
+        assert rehydrated == annotation
+        assert rehydrated.model_dump_json() == payload
+
+
+def test_discriminator_routes_to_the_right_variant() -> None:
+    bbox = geometry_adapter.validate_python(
+        {"type": "bbox", "x": 1.0, "y": 2.0, "width": 10.0, "height": 20.0}
+    )
+    assert isinstance(bbox, BboxGeometry)
+
+    polygon = geometry_adapter.validate_python(
+        {"type": "polygon", "points": [[0.0, 0.0], [10.0, 0.0], [10.0, 10.0]]}
+    )
+    assert isinstance(polygon, PolygonGeometry)
+
+    assert isinstance(
+        geometry_adapter.validate_python({"type": "classification_tag"}),
+        ClassificationGeometry,
+    )
+
+
+def test_bbox_rejects_non_positive_width_and_height() -> None:
+    for bad in [
+        {"width": -1.0},
+        {"height": -1.0},
+        {"width": 0.0},
+        {"height": 0.0},
+    ]:
+        data = {"type": "bbox", "x": 1.0, "y": 2.0, "width": 10.0, "height": 20.0, **bad}
+        with pytest.raises(ValidationError, match="greater than 0"):
+            geometry_adapter.validate_python(data)
+
+
+def test_polygon_needs_at_least_three_points() -> None:
+    for bad in [[], [(0.0, 0.0)], [(0.0, 0.0), (10.0, 0.0)]]:
+        with pytest.raises(ValidationError, match="at least 3 items"):
+            geometry_adapter.validate_python({"type": "polygon", "points": bad})
+
+
+def test_polygon_accepts_three_points_without_checking_self_intersection() -> None:
+    # M1 deliberately does not reject degenerate rings; documented in the model.
+    bowtie = [(0.0, 0.0), (10.0, 10.0), (10.0, 0.0), (0.0, 10.0)]
+    assert len(PolygonGeometry(points=bowtie).points) == 4
+
+
+def test_unimplemented_and_unknown_geometry_tags_are_rejected() -> None:
+    for tag in ["mask", "polyline", "keypoints", "cuboid_3d", "polyline_3d", "hexagon"]:
+        with pytest.raises(ValidationError, match="union_tag_invalid"):
+            geometry_adapter.validate_python({"type": tag})
+
+
+def test_variants_reject_fields_belonging_to_another_variant() -> None:
+    with pytest.raises(ValidationError, match="extra_forbidden"):
+        geometry_adapter.validate_python(
+            {"type": "bbox", "x": 1.0, "y": 2.0, "width": 10.0, "height": 20.0, "points": []}
+        )
+    with pytest.raises(ValidationError, match="extra_forbidden"):
+        geometry_adapter.validate_python({"type": "classification_tag", "x": 1.0})
+
+
+def test_geometry_is_immutable_once_validated() -> None:
+    with pytest.raises(ValidationError):
+        BboxGeometry(x=1.0, y=2.0, width=10.0, height=20.0).width = 30.0  # type: ignore[misc]
+
+
+def test_discriminator_values_are_geometry_type_members() -> None:
+    # The extension contract: a new variant must reuse a GeometryType member, so a
+    # schema's allowed geometries and an annotation's geometry stay directly comparable.
+    for geometry in VARIANTS:
+        assert isinstance(geometry.type, GeometryType)
+    assert {g.type for g in VARIANTS} <= set(GeometryType)
+
+
+def test_geometry_type_is_comparable_to_a_label_class_without_translation() -> None:
+    # This is the check AnnotationService (#7) performs against allowed_geometries;
+    # the union is designed so it needs no adapter layer.
+    label_class = LabelClass(name="car", geometry=GeometryType.BBOX)
+    annotation = _annotation(BboxGeometry(x=1.0, y=2.0, width=10.0, height=20.0))
+    assert annotation.geometry.type == label_class.geometry
+
+    tagged = _annotation(ClassificationGeometry())
+    assert tagged.geometry.type != label_class.geometry