Add API for getting well images

src/ome_zarr_models/v04/well.py

@@ -3,13 +3,15 @@
 """
 
 from collections import defaultdict
+from collections.abc import Generator
 from typing import Annotated, Literal
 
 from pydantic import AfterValidator, Field
 from pydantic_zarr.v2 import ArraySpec, GroupSpec
 
 from ome_zarr_models._utils import _AlphaNumericConstraint, _unique_items_validator
 from ome_zarr_models.base import Base
+from ome_zarr_models.v04.image import Image
 
 # WellGroup is defined one level higher
 __all__ = ["Well", "WellAttrs", "WellImage"]
@@ -71,4 +73,30 @@ class WellAttrs(Base):
 
 
 class WellGroup(GroupSpec[WellAttrs, ArraySpec | GroupSpec]):
-    pass
+    def get_image(self, i: int) -> Image:
+        """
+        Get a single image from this well.
+        """
+        image = self.attributes.well.images[i]
+        image_path = image.path
+        image_path_parts = image_path.split("/")
+        group = self
+        for part in image_path_parts:
+            group = group.members[part]
+
+        return Image(attributes=group.attributes, members=group.members)
+
+    @property
+    def n_images(self) -> int:
+        """
+        Number of images.
+        """
+        return len(self.attributes.well.images)
+
+    @property
+    def images(self) -> Generator[Image, None, None]:
+        """
+        Generator for all images in this well.
+        """
+        for i in range(self.n_images):
+            yield self.get_image(i)

tests/v04/test_hcs.py

@@ -3,7 +3,12 @@
 import zarr
 
 from ome_zarr_models.v04 import HCS
+from ome_zarr_models.v04.axes import Axis
+from ome_zarr_models.v04.coordinate_transformations import VectorScale
 from ome_zarr_models.v04.hcs import HCSAttrs
+from ome_zarr_models.v04.image import ImageAttrs
+from ome_zarr_models.v04.multiscales import Dataset, Multiscale
+from ome_zarr_models.v04.omero import Channel, Omero, Window
 from ome_zarr_models.v04.plate import Acquisition, Column, Plate, Row, WellInPlate
 from ome_zarr_models.v04.well import Well, WellImage
 
@@ -34,6 +39,74 @@ def test_example_hcs() -> None:
 
     well_groups = list(hcs.well_groups)
     assert len(well_groups) == 1
-    assert well_groups[0].attributes.well == Well(
+    well_group = well_groups[0]
+    assert well_group.attributes.well == Well(
         images=[WellImage(path="0", acquisition=None)], version="0.4"
     )
+
+    images = list(well_group.images)
+    assert len(images) == 1
+
+    assert images[0].attributes == ImageAttrs(
+        multiscales=[
+            Multiscale(
+                axes=[
+                    Axis(name="c", type="channel", unit=None),
+                    Axis(name="z", type="space", unit="micrometer"),
+                    Axis(name="y", type="space", unit="micrometer"),
+                    Axis(name="x", type="space", unit="micrometer"),
+                ],
+                datasets=[
+                    Dataset(
+                        path="0",
+                        coordinateTransformations=(
+                            VectorScale(type="scale", scale=[1.0, 1.0, 0.1625, 0.1625]),
+                        ),
+                    ),
+                    Dataset(
+                        path="1",
+                        coordinateTransformations=(
+                            VectorScale(type="scale", scale=[1.0, 1.0, 0.325, 0.325]),
+                        ),
+                    ),
+                    Dataset(
+                        path="2",
+                        coordinateTransformations=(
+                            VectorScale(type="scale", scale=[1.0, 1.0, 0.65, 0.65]),
+                        ),
+                    ),
+                    Dataset(
+                        path="3",
+                        coordinateTransformations=(
+                            VectorScale(type="scale", scale=[1.0, 1.0, 1.3, 1.3]),
+                        ),
+                    ),
+                    Dataset(
+                        path="4",
+                        coordinateTransformations=(
+                            VectorScale(type="scale", scale=[1.0, 1.0, 2.6, 2.6]),
+                        ),
+                    ),
+                ],
+                version="0.4",
+                coordinateTransformations=None,
+                metadata=None,
+                name=None,
+                type=None,
+            )
+        ],
+        omero=Omero(
+            channels=[
+                Channel(
+                    color="00FFFF",
+                    window=Window(max=65535.0, min=0.0, start=110.0, end=800.0),
+                    label="DAPI",
+                    wavelength_id="A01_C01",
+                )
+            ],
+            id=1,
+            name="TBD",
+            version="0.4",
+        ),
+        labels=None,
+    )