Keep non-uniform image coordinates in NPZ transfers and HDF5 files

datalab_kernel/backends/h5coords.py

@@ -0,0 +1,78 @@
+# Copyright (c) DataLab Platform Developers, BSD 3-Clause License
+# See LICENSE file for details
+
+"""
+HDF5 Image Coordinates
+======================
+
+Storage of the pixel coordinates of images in the HDF5 workspace files written by
+the backends.
+
+An image group holds:
+
+- for uniform coordinates, the ``x0``, ``y0``, ``dx`` and ``dy`` attributes;
+- for non-uniform coordinates, the ``xcoords`` and ``ycoords`` datasets (pixel
+  center coordinates along X and Y).
+
+The ``is_uniform_coords`` attribute tells which one applies. Files written by
+earlier versions (without this attribute) have uniform coordinates.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from sigima import ImageObj
+
+UNIFORM_ATTRS = ("x0", "y0", "dx", "dy")
+
+
+def save_image_coords(grp: Any, obj: ImageObj) -> None:
+    """Save the pixel coordinates of an image to an HDF5 group.
+
+    Args:
+        grp: HDF5 group of the image.
+        obj: Image.
+    """
+    grp.attrs["is_uniform_coords"] = bool(obj.is_uniform_coords)
+    if obj.is_uniform_coords:
+        for attr in UNIFORM_ATTRS:
+            value = getattr(obj, attr)
+            if value is not None:
+                grp.attrs[attr] = value
+    else:
+        grp.create_dataset("xcoords", data=np.asarray(obj.xcoords, dtype=float))
+        grp.create_dataset("ycoords", data=np.asarray(obj.ycoords, dtype=float))
+
+
+def load_image_coords(grp: Any, obj: ImageObj) -> None:
+    """Restore the pixel coordinates of an image from an HDF5 group.
+
+    Args:
+        grp: HDF5 group of the image.
+        obj: Image, whose data is already set.
+
+    Raises:
+        ValueError: If non-uniform coordinates are missing or do not match the
+         shape of the image.
+    """
+    if bool(grp.attrs.get("is_uniform_coords", True)):
+        for attr in UNIFORM_ATTRS:
+            if attr in grp.attrs:
+                setattr(obj, attr, float(grp.attrs[attr]))
+        return
+    if "xcoords" not in grp or "ycoords" not in grp:
+        raise ValueError(
+            "Invalid image group: missing xcoords or ycoords "
+            "for non-uniform coordinates"
+        )
+    xcoords = np.array(grp["xcoords"], dtype=float)
+    ycoords = np.array(grp["ycoords"], dtype=float)
+    nrows, ncols = obj.data.shape
+    if xcoords.shape != (ncols,) or ycoords.shape != (nrows,):
+        raise ValueError(
+            f"Invalid image group: coordinates of shape {xcoords.shape} (X) and "
+            f"{ycoords.shape} (Y) do not match image shape {obj.data.shape}"
+        )
+    obj.set_coords(xcoords, ycoords)

datalab_kernel/backends/standalone.py

@@ -16,6 +16,7 @@
 import numpy as np
 from sigima import ImageObj, SignalObj
 
+from datalab_kernel.backends.h5coords import load_image_coords, save_image_coords
 from datalab_kernel.workspace import WorkspaceBackend
 
 if TYPE_CHECKING:
@@ -166,11 +167,7 @@ def _save_object_to_group(self, grp, obj: DataObject) -> None:
             # Save image data
             grp.create_dataset("data", data=obj.data)
             # Save coordinate info
-            for attr in ("x0", "y0", "dx", "dy"):
-                if hasattr(obj, attr):
-                    val = getattr(obj, attr)
-                    if val is not None:
-                        grp.attrs[attr] = val
+            save_image_coords(grp, obj)
             # Save metadata
             if hasattr(obj, "title") and obj.title:
                 grp.attrs["title"] = obj.title
@@ -253,9 +250,7 @@ def _load_image(self, grp, name: str) -> DataObject:
         obj.data = data
         obj.title = grp.attrs.get("title", name)
 
-        for attr in ("x0", "y0", "dx", "dy"):
-            if attr in grp.attrs:
-                setattr(obj, attr, float(grp.attrs[attr]))
+        load_image_coords(grp, obj)
 
         if "xlabel" in grp.attrs:
             obj.xlabel = grp.attrs["xlabel"]

datalab_kernel/backends/webapi.py

@@ -41,6 +41,7 @@
 import numpy as np
 from sigima import ImageObj, SignalObj
 
+from datalab_kernel.backends.h5coords import load_image_coords, save_image_coords
 from datalab_kernel.backends.pyfetch import (
     HttpError,
     HttpResponse,
@@ -463,11 +464,7 @@ def _save_object_to_group(self, grp, obj: DataObject) -> None:
             grp.create_dataset("data", data=obj.data)
             if hasattr(obj, "title") and obj.title:
                 grp.attrs["title"] = obj.title
-            for attr in ("x0", "y0", "dx", "dy"):
-                if hasattr(obj, attr):
-                    val = getattr(obj, attr)
-                    if val is not None:
-                        grp.attrs[attr] = val
+            save_image_coords(grp, obj)
 
     def load(self, filepath: str) -> None:
         """Load workspace from HDF5 file.
@@ -507,9 +504,7 @@ def _load_object_from_group(self, grp, name: str) -> DataObject | None:
             obj = ImageObj()
             obj.data = data
             obj.title = grp.attrs.get("title", name)
-            for attr in ("x0", "y0", "dx", "dy"):
-                if attr in grp.attrs:
-                    setattr(obj, attr, float(grp.attrs[attr]))
+            load_image_coords(grp, obj)
             return obj
 
         return None

datalab_kernel/serialization_npz.py

@@ -24,7 +24,17 @@
 
 For ImageObj:
     - ``data.npy``: Image data (2D array, preserves dtype)
+    - ``xcoords.npy``, ``ycoords.npy``: Pixel center coordinates along X and Y
+      (1D float arrays, only for non-uniform coordinates)
     - ``metadata.json``: JSON-encoded metadata
+
+Metadata JSON includes:
+    - ``type``: "signal" or "image"
+    - ``title``, ``xlabel``, ``ylabel``, etc.
+    - ``x0``, ``y0``, ``dx``, ``dy`` (for images)
+    - ``is_uniform_coords`` (for images): False when the coordinates are given by
+      ``xcoords.npy`` and ``ycoords.npy``. Archives without this key (written by
+      earlier versions) have uniform coordinates.
 """
 
 from __future__ import annotations
@@ -122,14 +132,14 @@ def _serialize_signal(obj, buffer: io.BytesIO, *, compress: bool = True) -> None
     """Serialize a SignalObj to NPZ format."""
     metadata = {
         "type": "signal",
-        "title": getattr(obj, "title", None),
-        "xlabel": getattr(obj, "xlabel", None),
-        "ylabel": getattr(obj, "ylabel", None),
-        "xunit": getattr(obj, "xunit", None),
-        "yunit": getattr(obj, "yunit", None),
+        "title": obj.title,
+        "xlabel": obj.xlabel,
+        "ylabel": obj.ylabel,
+        "xunit": obj.xunit,
+        "yunit": obj.yunit,
     }
     # Include object metadata (contains Geometry_*, Table_* results)
-    obj_metadata = getattr(obj, "metadata", None)
+    obj_metadata = obj.metadata
     if obj_metadata:
         metadata["obj_metadata"] = _serialize_obj_metadata(obj_metadata)
 
@@ -150,26 +160,30 @@ def _serialize_image(obj, buffer: io.BytesIO, *, compress: bool = True) -> None:
     """Serialize an ImageObj to NPZ format."""
     metadata = {
         "type": "image",
-        "title": getattr(obj, "title", None),
-        "xlabel": getattr(obj, "xlabel", None),
-        "ylabel": getattr(obj, "ylabel", None),
-        "zlabel": getattr(obj, "zlabel", None),
-        "xunit": getattr(obj, "xunit", None),
-        "yunit": getattr(obj, "yunit", None),
-        "zunit": getattr(obj, "zunit", None),
-        "x0": getattr(obj, "x0", 0.0),
-        "y0": getattr(obj, "y0", 0.0),
-        "dx": getattr(obj, "dx", 1.0),
-        "dy": getattr(obj, "dy", 1.0),
+        "title": obj.title,
+        "xlabel": obj.xlabel,
+        "ylabel": obj.ylabel,
+        "zlabel": obj.zlabel,
+        "xunit": obj.xunit,
+        "yunit": obj.yunit,
+        "zunit": obj.zunit,
+        "x0": obj.x0,
+        "y0": obj.y0,
+        "dx": obj.dx,
+        "dy": obj.dy,
+        "is_uniform_coords": bool(obj.is_uniform_coords),
     }
     # Include object metadata (contains Geometry_*, Table_* results)
-    obj_metadata = getattr(obj, "metadata", None)
+    obj_metadata = obj.metadata
     if obj_metadata:
         metadata["obj_metadata"] = _serialize_obj_metadata(obj_metadata)
 
     compression = zipfile.ZIP_DEFLATED if compress else zipfile.ZIP_STORED
     with zipfile.ZipFile(buffer, "w", compression) as zf:
         _write_array_to_zip(zf, "data.npy", obj.data)
+        if not obj.is_uniform_coords:
+            _write_array_to_zip(zf, "xcoords.npy", obj.xcoords)
+            _write_array_to_zip(zf, "ycoords.npy", obj.ycoords)
         zf.writestr("metadata.json", json.dumps(metadata, ensure_ascii=False))
 
 
@@ -272,10 +286,27 @@ def _deserialize_image(zf: zipfile.ZipFile, metadata: dict) -> DataObject:
     if metadata.get("zunit"):
         obj.zunit = metadata["zunit"]
 
-    obj.x0 = metadata.get("x0", 0.0)
-    obj.y0 = metadata.get("y0", 0.0)
-    obj.dx = metadata.get("dx", 1.0)
-    obj.dy = metadata.get("dy", 1.0)
+    # Set coordinate info
+    if metadata.get("is_uniform_coords", True):
+        obj.x0 = metadata.get("x0", 0.0)
+        obj.y0 = metadata.get("y0", 0.0)
+        obj.dx = metadata.get("dx", 1.0)
+        obj.dy = metadata.get("dy", 1.0)
+    else:
+        xcoords = _read_array_from_zip(zf, "xcoords.npy")
+        ycoords = _read_array_from_zip(zf, "ycoords.npy")
+        if xcoords is None or ycoords is None:
+            raise ValueError(
+                "Invalid image NPZ: missing xcoords.npy or ycoords.npy "
+                "for non-uniform coordinates"
+            )
+        nrows, ncols = data.shape
+        if xcoords.shape != (ncols,) or ycoords.shape != (nrows,):
+            raise ValueError(
+                f"Invalid image NPZ: coordinates of shape {xcoords.shape} (X) and "
+                f"{ycoords.shape} (Y) do not match image shape {data.shape}"
+            )
+        obj.set_coords(xcoords, ycoords)
 
     # Restore object metadata (Geometry_*, Table_* results)
     if metadata.get("obj_metadata"):