Get rid of LabelImage, add image_label field directly to Image

src/yaozarrs/v04/__init__.py

@@ -24,7 +24,6 @@
 from ._label import (
     ImageLabel,
     LabelColor,
-    LabelImage,
     LabelProperty,
     LabelSource,
 )
@@ -44,7 +43,6 @@
     "Image",
     "ImageLabel",
     "LabelColor",
-    "LabelImage",
     "LabelProperty",
     "LabelSource",
     "Multiscale",

src/yaozarrs/v04/_image.py

@@ -1,14 +1,22 @@
-from typing import Annotated, Literal, TypeAlias
+from typing import Annotated, Any, Literal, TypeAlias
 
 from annotated_types import Len, MinLen
-from pydantic import AfterValidator, Field, WrapValidator, model_validator
+from pydantic import (
+    AfterValidator,
+    Field,
+    WrapValidator,
+    model_serializer,
+    model_validator,
+)
 from typing_extensions import Self
 
 from yaozarrs._base import ZarrGroupModel, _BaseModel
 from yaozarrs._omero import Omero
 from yaozarrs._types import UniqueList
 from yaozarrs._units import SpaceUnits, TimeUnits
 
+from ._label import ImageLabel
+
 # ------------------------------------------------------------------------------
 # Axis model
 # ------------------------------------------------------------------------------
@@ -250,3 +258,19 @@ def ndim(self) -> int:
 class Image(ZarrGroupModel):
     multiscales: Annotated[UniqueList[Multiscale], MinLen(1)]
     omero: Omero | None = None
+    image_label: ImageLabel | None = Field(
+        default=None,
+        alias="image-label",
+        description=(
+            "Label-specific metadata for segmentation images. "
+            "When present, indicates this is a label/segmentation image."
+        ),
+    )
+
+    @model_serializer(mode="wrap")
+    def _serialize_model(self, serializer: Any, info: Any) -> dict[str, Any]:
+        """Custom serializer to exclude image-label when None or empty."""
+        data = serializer(self)
+        if "image-label" in data and not data["image-label"]:
+            del data["image-label"]
+        return data

src/yaozarrs/v04/_label.py

@@ -6,8 +6,6 @@
 from yaozarrs._base import _BaseModel
 from yaozarrs._types import UniqueList
 
-from ._image import Image
-
 # ------------------------------------------------------------------------------
 # Color model
 # ------------------------------------------------------------------------------
@@ -83,14 +81,3 @@ class LabelsGroup(_BaseModel):
     labels: Annotated[list[str], MinLen(1)] = Field(
         description="Array of paths to labeled multiscale images"
     )
-
-
-# ------------------------------------------------------------------------------
-# Label model (top-level for individual label images)
-# ------------------------------------------------------------------------------
-
-
-class LabelImage(Image):
-    """Model for individual label images with multiscales + image-label metadata."""
-
-    image_label: ImageLabel = Field(alias="image-label")

src/yaozarrs/v04/_zarr_json.py

@@ -226,16 +226,14 @@
 
 from ._bf2raw import Bf2Raw
 from ._image import Image
-from ._label import LabelImage, LabelsGroup
+from ._label import LabelsGroup
 from ._plate import Plate
 from ._series import Series
 from ._well import Well
 
 
 def _discriminate_ome_v04_metadata(v: Any) -> str | None:
     if isinstance(v, dict):
-        if "image-label" in v:
-            return "label-image"
         if "multiscales" in v:
             return "image"
         if "plate" in v:
@@ -249,8 +247,6 @@ def _discriminate_ome_v04_metadata(v: Any) -> str | None:
         if "series" in v:
             return "series"
     elif isinstance(v, BaseModel):
-        if isinstance(v, LabelImage):
-            return "label-image"
         if isinstance(v, Image):
             return "image"
         if isinstance(v, Plate):
@@ -270,8 +266,7 @@ def _discriminate_ome_v04_metadata(v: Any) -> str | None:
 # these are ALL also ZarrGroupModels (i.e. have a "uri" attribute)
 OMEZarrGroupJSON: TypeAlias = Annotated[
     (
-        Annotated[LabelImage, Tag("label-image")]
-        | Annotated[Image, Tag("image")]
+        Annotated[Image, Tag("image")]
         | Annotated[Plate, Tag("plate")]
         | Annotated[Bf2Raw, Tag("bf2raw")]
         | Annotated[Well, Tag("well")]

src/yaozarrs/v05/__init__.py

@@ -22,7 +22,6 @@
 from ._label import (
     ImageLabel,
     LabelColor,
-    LabelImage,
     LabelProperty,
     LabelsGroup,
     LabelSource,
@@ -43,7 +42,6 @@
     "Image",
     "ImageLabel",
     "LabelColor",
-    "LabelImage",
     "LabelProperty",
     "LabelSource",
     "LabelsGroup",
@@ -61,7 +59,6 @@
     "Row",
     "ScaleTransformation",
     "Series",
-    "Series",
     "SpaceAxis",
     "TimeAxis",
     "TranslationTransformation",

src/yaozarrs/v05/_image.py

@@ -1,14 +1,22 @@
-from typing import Annotated, Literal, TypeAlias
+from typing import Annotated, Any, Literal, TypeAlias
 
 from annotated_types import Len, MinLen
-from pydantic import AfterValidator, Field, WrapValidator, model_validator
+from pydantic import (
+    AfterValidator,
+    Field,
+    WrapValidator,
+    model_serializer,
+    model_validator,
+)
 from typing_extensions import Self
 
 from yaozarrs._base import _BaseModel
 from yaozarrs._omero import Omero
 from yaozarrs._types import UniqueList
 from yaozarrs._units import SpaceUnits, TimeUnits
 
+from ._label import ImageLabel
+
 # ------------------------------------------------------------------------------
 # Axis model
 # ------------------------------------------------------------------------------
@@ -321,6 +329,27 @@ def ndim(self) -> int:
 
 
 class Image(_BaseModel):
+    """Top-level OME-NGFF image metadata.
+
+    When `image_label` is present, this represents a label/segmentation image.
+    """
+
     version: Literal["0.5"] = "0.5"
     multiscales: Annotated[UniqueList[Multiscale], MinLen(1)]
     omero: Omero | None = None
+    image_label: ImageLabel | None = Field(
+        default=None,
+        alias="image-label",
+        description=(
+            "Label-specific metadata for segmentation images. "
+            "When present, indicates this is a label/segmentation image."
+        ),
+    )
+
+    @model_serializer(mode="wrap")
+    def _serialize_model(self, serializer: Any, info: Any) -> dict[str, Any]:
+        """Custom serializer to exclude image-label when None or empty."""
+        data = serializer(self)
+        if "image-label" in data and not data["image-label"]:
+            del data["image-label"]
+        return data

src/yaozarrs/v05/_label.py

@@ -6,8 +6,6 @@
 from yaozarrs._base import _BaseModel
 from yaozarrs._types import UniqueList
 
-from ._image import Image
-
 # ------------------------------------------------------------------------------
 # Color model
 # ------------------------------------------------------------------------------
@@ -92,14 +90,3 @@ class LabelsGroup(_BaseModel):
     labels: Annotated[list[str], MinLen(1)] = Field(
         description="Array of paths to labeled multiscale images"
     )
-
-
-# ------------------------------------------------------------------------------
-# Label model (top-level for individual label images)
-# ------------------------------------------------------------------------------
-
-
-class LabelImage(Image):
-    """Model for individual label images with multiscales + image-label metadata."""
-
-    image_label: ImageLabel = Field(alias="image-label")

src/yaozarrs/v05/_storage.py

@@ -18,7 +18,7 @@
 from yaozarrs._zarr import ZarrArray, ZarrGroup
 from yaozarrs.v05._bf2raw import Bf2Raw
 from yaozarrs.v05._image import Image, Multiscale
-from yaozarrs.v05._label import LabelImage, LabelsGroup
+from yaozarrs.v05._label import LabelsGroup
 from yaozarrs.v05._plate import Plate
 from yaozarrs.v05._series import Series
 from yaozarrs.v05._well import Well
@@ -117,9 +117,7 @@ def validate_group(
         loc_prefix = ("ome",)
 
         # Dispatch to appropriate visitor method based on metadata type
-        if isinstance(ome_metadata, LabelImage):
-            return validator.visit_label_image(zarr_group, ome_metadata, loc_prefix)
-        elif isinstance(ome_metadata, Image):
+        if isinstance(ome_metadata, Image):
             return validator.visit_image(zarr_group, ome_metadata, loc_prefix)
         elif isinstance(ome_metadata, LabelsGroup):
             return validator.visit_labels_group(zarr_group, ome_metadata, loc_prefix)
@@ -137,26 +135,30 @@ def validate_group(
             )
 
     def visit_label_image(
-        self, zarr_group: ZarrGroup, label_image_model: LabelImage, loc_prefix: Loc
+        self, zarr_group: ZarrGroup, image_model: Image, loc_prefix: Loc
     ) -> ValidationResult:
-        """Validate a LabelImage group."""
+        """Validate label-specific metadata for an Image with image_label.
+
+        This should only be called for images where image_model.image_label is not None.
+        """
         result = ValidationResult()
 
         # The value of the source key MUST be a JSON object containing information
         # about the original image from which the label image derives. This object
         # MAY include a key image, whose value MUST be a string specifying the
         # relative path to a Zarr image group.
-        src = label_image_model.image_label.source
-        if src is not None and (src_img := src.image) is not None:
+        if image_model.image_label is not None:
+            src = image_model.image_label.source
+            if src is not None and (src_img := src.image) is not None:
+                result = result.merge(
+                    self._validate_labels_image_source(zarr_group, src_img, loc_prefix)
+                )
+
+            # For label images, validate integer data types
             result = result.merge(
-                self._validate_labels_image_source(zarr_group, src_img, loc_prefix)
+                self._validate_label_data_types(image_model, zarr_group, loc_prefix)
             )
 
-        # For label images, validate integer data types
-        result = result.merge(
-            self._validate_label_data_types(label_image_model, zarr_group, loc_prefix)
-        )
-
         return result
 
     def visit_image(
@@ -181,6 +183,12 @@ def visit_image(
                 self._visit_multiscale_no_prefetch(zarr_group, multiscale, ms_loc)
             )
 
+        # If this is a label image, validate label-specific aspects
+        if image_model.image_label is not None:
+            result = result.merge(
+                self.visit_label_image(zarr_group, image_model, loc_prefix)
+            )
+
         # Check whether this image has a labels group, and validate if so
         lbls_check = self._check_for_labels_group(zarr_group, loc_prefix)
         result = result.merge(lbls_check.result)
@@ -301,20 +309,23 @@ def visit_labels_group(
                             },
                         )
 
-            if isinstance(label_image_model, LabelImage):
-                # Recursively validate the label image
-                result = result.merge(
-                    self.visit_label_image(label_group, label_image_model, label_loc)
-                )
-            else:
+            # Check if this is actually a label image (has image-label metadata)
+            if label_image_model.image_label is None:
                 # TODO: should it just be a warning?
                 result.add_error(
                     StorageErrorType.label_image_invalid,
                     label_loc,
                     f"Label path '{label_path}' contains Image metadata, "
-                    "but is not a LabelImage (missing 'image-label' metadata?)",
+                    "but is missing 'image-label' metadata",
                     ctx={"path": label_path, "type": type(label_image_model).__name__},
                 )
+                continue
+
+            # Recursively validate the label image
+            # Note: visit_image will validate multiscales AND label-specific aspects
+            result = result.merge(
+                self.visit_image(label_group, label_image_model, label_loc)
+            )
 
         return result
 
@@ -833,7 +844,7 @@ def _validate_labels_image_source(
         return result
 
     def _validate_label_data_types(
-        self, image_model: LabelImage, zarr_group: ZarrGroup, loc_prefix: Loc
+        self, image_model: Image, zarr_group: ZarrGroup, loc_prefix: Loc
     ) -> ValidationResult:
         """Validate that label arrays contain only integer data types."""
         result = ValidationResult()

src/yaozarrs/v05/_zarr_json.py

@@ -275,16 +275,14 @@
 from yaozarrs.v05._bf2raw import Bf2Raw
 
 from ._image import Image
-from ._label import LabelImage, LabelsGroup
+from ._label import LabelsGroup
 from ._plate import Plate
 from ._series import Series
 from ._well import Well
 
 
 def _discriminate_ome_v05_metadata(v: Any) -> str | None:
     if isinstance(v, dict):
-        if "image-label" in v:
-            return "label-image"
         if "multiscales" in v:
             return "image"
         if "plate" in v:
@@ -298,8 +296,6 @@ def _discriminate_ome_v05_metadata(v: Any) -> str | None:
         if "series" in v:
             return "series"
     elif isinstance(v, BaseModel):
-        if isinstance(v, LabelImage):
-            return "label-image"
         if isinstance(v, Image):
             return "image"
         if isinstance(v, Plate):
@@ -317,8 +313,7 @@ def _discriminate_ome_v05_metadata(v: Any) -> str | None:
 
 OMEMetadata: TypeAlias = Annotated[
     (
-        Annotated[LabelImage, Tag("label-image")]
-        | Annotated[Image, Tag("image")]
+        Annotated[Image, Tag("image")]
         | Annotated[Plate, Tag("plate")]
         | Annotated[Bf2Raw, Tag("bf2raw")]
         | Annotated[Well, Tag("well")]