From 768aeb74a874e98d64028f277facae537525a9f9 Mon Sep 17 00:00:00 2001 From: Armando Anaya Date: Sun, 26 Jul 2026 02:45:18 -0700 Subject: [PATCH] feat(kernel): bbox/polygon geometry as a pydantic discriminated union MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces Annotation.geometry: dict[str, object] with a real union discriminated on `type`: BboxGeometry, PolygonGeometry and ClassificationGeometry. The discriminator values are GeometryType members rather than parallel string literals, so a schema's allowed geometries and an annotation's geometry are directly comparable — AnnotationService (#7) needs no translation layer. Variants are frozen and forbid extra fields, so a bbox payload carrying polygon keys is rejected instead of silently validating. Non-positive width/height and polygons with fewer than three points are rejected at the model level; self-intersection is explicitly not validated in M1. --- src/visionset/kernel/domain/__init__.py | 10 ++ src/visionset/kernel/domain/annotation.py | 4 +- src/visionset/kernel/domain/geometry.py | 83 ++++++++++++++ tests/kernel/test_annotation.py | 2 +- tests/kernel/test_geometry.py | 127 ++++++++++++++++++++++ 5 files changed, 224 insertions(+), 2 deletions(-) create mode 100644 src/visionset/kernel/domain/geometry.py create mode 100644 tests/kernel/test_geometry.py diff --git a/src/visionset/kernel/domain/__init__.py b/src/visionset/kernel/domain/__init__.py index 14d62f37..6d2c039c 100644 --- a/src/visionset/kernel/domain/__init__.py +++ b/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", diff --git a/src/visionset/kernel/domain/annotation.py b/src/visionset/kernel/domain/annotation.py index 17450787..693331df 100644 --- a/src/visionset/kernel/domain/annotation.py +++ b/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) diff --git a/src/visionset/kernel/domain/geometry.py b/src/visionset/kernel/domain/geometry.py new file mode 100644 index 00000000..0f8b8ef4 --- /dev/null +++ b/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. +""" diff --git a/tests/kernel/test_annotation.py b/tests/kernel/test_annotation.py index 54c2ca1d..b75a2778 100644 --- a/tests/kernel/test_annotation.py +++ b/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) diff --git a/tests/kernel/test_geometry.py b/tests/kernel/test_geometry.py new file mode 100644 index 00000000..e16b4e6c --- /dev/null +++ b/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