Add ngff example data script WIP

elf/io/ngff.py

@@ -6,27 +6,56 @@
 AXES_NAMES = {"t", "c", "z", "y", "x"}
 
 
-def _get_chunks(ndim):
-    return (256, 256) if ndim == 2 else (64, 64, 64)
+def _get_chunks(axes_names):
+    if len(axes_names) == 2:
+        return (256, 256)
+    elif len(axes_names) == 3:
+        return 3*(64,) if axes_names[0] == 'z' else (1, 256, 256)
+    elif len(axes_names) == 4:
+        return (1, 1, 256, 256) if axes_names[:2] == ('t', 'c') else (1, 64, 64, 64)
+    else:
+        return (1, 1, 64, 64, 64)
 
 
-def write_ome_zarr(data, path, name, n_scales,
+def _validate_axes_names(ndim, axes_names):
+    assert len(axes_names) == ndim
+    val_axes = tuple(axes_names)
+    if ndim == 2:
+        assert val_axes == ('y', 'x')
+    elif ndim == 3:
+        assert val_axes in [('z', 'y', 'x'), ('c', 'y', 'x'), ('t', 'y', 'x')]
+    elif ndim == 4:
+        assert val_axes in [('t', 'z', 'y', 'x'), ('c', 'z' 'y', 'x'), ('t', 'c', 'y', 'x')]
+    else:
+        assert val_axes == ('t', 'c', 'z', 'y', 'x')
+
+
+# TODO downscale only in spatial dimensions
+def _downscale(data, downscaler, kwargs):
+    data = downscaler(data, **kwargs).astype(data.dtype)
+    return data
+
+
+# TODO expose dimension separator as param
+def write_ome_zarr(data, path, axes_names, name, n_scales,
                    key=None, chunks=None,
                    downscaler=skimage.transform.rescale,
                    kwargs={"scale": (0.5, 0.5, 0.5), "order": 0, "preserve_range": True}):
     """Write numpy data to ome.zarr format.
     """
-    assert data.ndim in (2, 3)
-    chunks = _get_chunks(data.ndim) if chunks is None else chunks
-    axes_names = ["y", "x"] if data.ndim == 2 else ["z", "y", "x"]
+
+    assert 2 <= data.ndim <= 5
+    _validate_axes_names(data.ndim, axes_names)
+
+    chunks = _get_chunks(axes_names) if chunks is None else chunks
     store = zarr.NestedDirectoryStore(path, dimension_separator="/")
+
     with zarr.open(store, mode='a') as f:
         g = f if key is None else f.require_group(key)
         g.create_dataset('s0', data=data, chunks=chunks, dimension_separator="/")
-        if n_scales > 1:
-            for ii in range(1, n_scales):
-                data = downscaler(data, **kwargs).astype(data.dtype)
-                g.create_dataset(f's{ii}', data=data, chunks=chunks, dimension_separator="/")
+        for ii in range(1, n_scales):
+            data = _downscale(data, downscaler, kwargs)
+            g.create_dataset(f's{ii}', data=data, chunks=chunks, dimension_separator="/")
         function_name = f'{downscaler.__module__}.{downscaler.__name__}'
         create_ngff_metadata(g, name, axes_names,
                              type_=function_name, metadata=kwargs)
@@ -61,4 +90,8 @@ def create_ngff_metadata(g, name, axes_names, type_=None, metadata=None):
     metadata = g.attrs.get("multiscales", [])
     metadata.append(ms_entry)
     g.attrs["multiscales"] = metadata
-    g.attrs["_ARRAY_DIMENSIONS"] = axes_names
+
+    # write the array dimensions for compat with xarray:
+    # https://xarray.pydata.org/en/stable/internals/zarr-encoding-spec.html?highlight=zarr
+    for ds in g.values():
+        ds.attrs["_ARRAY_DIMENSIONS"] = axes_names

example/io/ngff_examples.py

@@ -0,0 +1,85 @@
+import argparse
+import os
+
+import imageio
+import numpy as np
+from elf.io import open_file
+from elf.io.ngff import write_ome_zarr
+
+
+def _kwargs_2d():
+    return {"scale": (0.5, 0.5), "order": 0, "preserve_range": True}
+
+
+def _kwargs_3d():
+    return {"scale": (0.5, 0.5, 0.5), "order": 0, "preserve_range": True}
+
+
+def _load_data(path, key, bb):
+    if key is None:
+        data = imageio.imread(path)
+        data = data[bb]
+    else:
+        with open_file(path, 'r') as f:
+            data = f[key][bb]
+    return data
+
+
+def _create_example(in_path, folder, axes, key=None, bb=np.s_[:]):
+    data = _load_data(in_path, key, bb)
+    assert data.ndim == len(axes)
+    ax_name = ''.join(axes)
+    out_path = os.path.join(folder, f"{ax_name}.ome.zr")
+    kwargs = _kwargs_3d() if axes[-3:] == ('z', 'y', 'x') else _kwargs_2d()
+    write_ome_zarr(data, out_path, axes, ax_name,
+                   n_scales=3, kwargs=kwargs)
+
+#
+# create ngff ome.zarr example data
+#
+# all the filepath are hard-coded to the EMBL kreshuk group share
+
+
+# axes: yx
+def create_2d_example(folder):
+    in_path = os.path.join("/g/kreshuk/data/covid/covid-data-vibor/20200405_test_images",
+                           "WellC01_PointC01_0000_ChannelDAPI,WF_GFP,TRITC,WF_Cy5_Seq0216.tiff")
+    _create_example(in_path, folder, axes=("y", "x"), bb=np.s_[0, :, :])
+
+
+# add linked labels for the zyx example
+# axes: zyx, cyx, tyx
+def create_3d_examples(folder):
+    pass
+
+
+# axes: tcyx, tzyx, czyx
+def create_4d_examples(folder):
+    pass
+
+
+# axes: tczyx
+def create_5d_example(folder):
+    pass
+
+
+# using '.' dimension separator
+def create_flat_example(folder):
+    pass
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-o', '--output_root', required=True)
+    parser.add_argument('--version', default="v0.3")
+    args = parser.parse_args()
+
+    output_folder = os.path.join(args.output_root, args.version)
+    os.makedirs(output_folder, exist_ok=True)
+
+    create_2d_example(output_folder)
+    create_3d_examples(output_folder)
+    create_4d_examples(output_folder)
+    create_5d_example(output_folder)
+    create_flat_example(output_folder)
+    # TODO copy README