Support exporting dataset

.isort.cfg

@@ -1,2 +1,2 @@
 [settings]
-known_third_party = numpy,ome_zarr,omero,omero_zarr,setuptools,skimage,zarr
+known_third_party = dask,numpy,ome_zarr,omero,omero_model_DatasetI,omero_zarr,setuptools,skimage,zarr

src/omero_zarr/cli.py

@@ -8,13 +8,15 @@
 from omero.cli import CLI, BaseControl, Parser, ProxyStringType
 from omero.gateway import BlitzGateway, BlitzObjectWrapper
 from omero.model import ImageI, PlateI
+from omero_model_DatasetI import DatasetI
 from zarr.hierarchy import open_group
 from zarr.storage import FSStore
 
 from .masks import MASK_DTYPE_SIZE, image_shapes_to_zarr, plate_shapes_to_zarr
 from .raw_pixels import (
     add_omero_metadata,
     add_toplevel_metadata,
+    dataset_to_zarr,
     image_to_zarr,
     plate_to_zarr,
 )
@@ -218,7 +220,7 @@ def _configure(self, parser: Parser) -> None:
         export.add_argument(
             "--max_workers",
             default=None,
-            help="Maximum number of workers (only for use with bioformats2raw)",
+            help="Maximum number of workers",
         )
         export.add_argument(
             "object",
@@ -257,10 +259,13 @@ def export(self, args: argparse.Namespace) -> None:
             if args.bf:
                 self._bf_export(image, args)
             else:
-                image_to_zarr(image, args)
+                image_to_zarr(image, args.output, args.cache_numpy)
         elif isinstance(args.object, PlateI):
             plate = self._lookup(self.gateway, "Plate", args.object.id)
             plate_to_zarr(plate, args)
+        elif isinstance(args.object, DatasetI):
+            dataset = self._lookup(self.gateway, "Dataset", args.object.id)
+            dataset_to_zarr(dataset, args)
 
     def _lookup(
         self, gateway: BlitzGateway, otype: str, oid: int

src/omero_zarr/raw_pixels.py

@@ -3,6 +3,7 @@
 import time
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
 
+import dask
 import numpy
 import numpy as np
 import omero.clients  # noqa
@@ -30,10 +31,10 @@ def _open_store(name: str) -> FSStore:
     )
 
 
-def image_to_zarr(image: omero.gateway.ImageWrapper, args: argparse.Namespace) -> None:
-    target_dir = args.output
-    cache_dir = target_dir if args.cache_numpy else None
-
+def image_to_zarr(
+    image: omero.gateway.ImageWrapper, target_dir: str, cache_numpy: bool
+) -> None:
+    cache_dir = target_dir if cache_numpy else None
     name = os.path.join(target_dir, "%s.zarr" % image.id)
     print(f"Exporting to {name} ({VERSION})")
     store = _open_store(name)
@@ -42,7 +43,53 @@ def image_to_zarr(image: omero.gateway.ImageWrapper, args: argparse.Namespace) -
     add_multiscales_metadata(root, axes, n_levels)
     add_omero_metadata(root, image)
     add_toplevel_metadata(root)
-    print("Finished.")
+    print("{name} Finished.")
+
+
+def image_to_zarr_threadsafe(
+    client: omero.client, image_id: int, img_root: Group, cache_dir: str
+) -> None:
+    conn = omero.gateway.BlitzGateway(client_obj=client)
+    conn.SERVICE_OPTS.setOmeroGroup("-1")
+    image = conn.getObject("Image", image_id)
+    print(f"Exporting Image {image_id}")
+    n_levels, axes = add_image(image, img_root, cache_dir=cache_dir)
+    add_multiscales_metadata(img_root, axes, n_levels)
+    add_omero_metadata(img_root, image)
+    add_toplevel_metadata(img_root)
+    print(f"{image_id} Finished.")
+    # Explicitly close RawPixelsStore
+    conn.createRawPixelsStore().close()
+
+
+def dataset_to_zarr(
+    dataset: omero.gateway.DatasetWrapper, args: argparse.Namespace
+) -> None:
+    target_dir = os.path.join(args.output, "%s.zarr" % dataset.getId())
+    cache_dir = target_dir if args.cache_numpy else None
+    print(f"Exporting to {target_dir} ({VERSION})")
+    store = _open_store(target_dir)
+    root = open_group(store)
+    jobs = []
+    collection: Dict[str, Dict] = {}
+    for image in dataset.listChildren():
+        name = image.getName()
+        # in case we get duplicate names, add _id
+        if os.path.exists(os.path.join(target_dir, name)):
+            name = f"{name}_{image.id}"
+        collection[name] = {}
+        img_root = root.create_group(name)
+        jobs.append(
+            dask.delayed(image_to_zarr_threadsafe)(
+                dataset._conn.c, image.getId(), img_root, cache_dir
+            )
+        )
+    if args.max_workers:
+        with dask.config.set(num_workers=int(args.max_workers)):
+            dask.delayed()(jobs).compute()
+    else:
+        dask.delayed()(jobs).compute()
+    root.attrs["collection"] = collection
 
 
 def add_image(