Add function blueprints for ome zarr and cryo et portal loading

Author: constantinpape

synapse_net/file_utils.py

@@ -5,6 +5,16 @@
 import numpy as np
 import pooch
 
+try:
+    import cryoet_data_portal as cdp
+except ImportError:
+    cdp = None
+
+try:
+    import zarr
+except ImportError:
+    zarr = None
+
 
 def get_cache_dir() -> str:
     """Get the cache directory of synapse net.
@@ -88,3 +98,47 @@ def read_mrc(path: str) -> Tuple[np.ndarray, Dict[str, float]]:
     # Transpose the data to match python axis order.
     data = np.flip(data, axis=1) if data.ndim == 3 else np.flip(data, axis=0)
     return data, voxel_size
+
+
+def read_ome_zarr(uri: str, scale_level: int = 0) -> Tuple[np.ndarray, Dict[str, float]]:
+    """Read data and voxel size from an ome.zarr file.
+
+    Args:
+        uri: Path or url to the ome.zarr file.
+        scale_level: The level of the multi-scale image pyramid to load.
+
+    Returns:
+        The data read from the file.
+        The voxel size read from the file.
+    """
+    if zarr is None:
+        raise RuntimeError("The zarr library is required to read ome.zarr files.")
+
+    # TODO handle URLs / make sure that zarr parses it correctly.
+    with zarr.open(uri, "r") as f:
+        multiscales = f.attrs["multiscales"][0]
+        # TODO double check that the metadata is correct and transform the voxel size to a dict.
+        # TODO voxel size is given in Angstrom, divide by 10 to get nanometer
+        internal_path = multiscales["dataset"][scale_level]
+        data = f[internal_path][:]
+        transformation = multiscales["transformation"][scale_level]
+        voxel_size = transformation["scale"]
+
+    return data, voxel_size
+
+
+def read_data_from_cryo_et_portal_run(
+    run_id: int, output_path: Optional[str] = None
+) -> Tuple[np.ndarray, Dict[str, float]]:
+    """Read data and voxel size from a CryoET Data Portal run.
+
+    Args:
+        run_id: The ID of the experiment run.
+        output_path: The path for saving the data. The data will be streamed if the path is not given.
+
+    Returns:
+        The data read from the run.
+        The voxel size read from the run
+    """
+    if cdp is None:
+        raise RuntimeError("The CryoET Data portal library is required to read data from the portal.")