simple_ome_zarr.py

import numpy as np
import zarr


def create_ome_zarr_datasets(
    f,
    shape,
    chunks,
    scale_factors,
    dtype
):
    dataset_names = []
    for scale_id, scale_factor in enumerate(scale_factors):
        ds_name = f's{scale_id}'

        internal_shape = [1, 1] + [sh // sf for sh, sf in zip(shape, scale_factor)]
        internal_chunks = [1, 1] + list(chunks)

        f.create_dataset(
            ds_name,
            shape=internal_shape,
            chunks=internal_chunks,
            dtype=dtype
        )
        dataset_names.append(ds_name)
    return dataset_names


def create_ome_zarr_metadata(
    f,
    vol_name,
    dataset_names,
    scale_factors,
    resolution,
    unit
):
    transform = {
        'axes': ['x', 'y', 'z'],
        'scale': resolution[::-1],
        'translate': [0., 0., 0.],
        'units': 3 * [unit]
    }

    pixel_res = {
        'dimensions': resolution[::-1],
        'unit': unit
    }

    ms_attributes = {
        'name': vol_name,
        'version': '0.1',
        'datasets': [{'path': name} for name in dataset_names],
        # TODO I am not sure if the axis order for the scales needs to be reversed
        # 'scales': [scale[::-1] for scale in scale_factors],
        'scales': scale_factors,
        'transform': transform,
        'pixelResolution': pixel_res
    }
    f.attrs['multiscales'] = ms_attributes


def create_ome_zarr(
    path,
    shape,
    chunks,
    resolution,
    name,
    dtype='uint8',
    scale_factors=None,
    unit='micrometer'
):
    """ Create an empty ome.zarr file for volumetric image raw data.

    Expects all inputs in axis order zyx.
    """
    assert len(shape) == len(chunks) == len(resolution) == 3, "Only support 3D data"

    if scale_factors is None:
        internal_scale_factors = [[1, 1, 1]]
    else:
        assert all(len(sf) == 3 for sf in scale_factors), "Invalid scale factors"
        internal_scale_factors = [[1, 1, 1]] + scale_factors

    f = zarr.open(path, mode='a')
    dataset_names = create_ome_zarr_datasets(f, shape, chunks, internal_scale_factors, dtype)

    create_ome_zarr_metadata(f, name, dataset_names, internal_scale_factors, resolution, unit)


if __name__ == '__main__':
    # create the ome.zarr file
    create_ome_zarr(
        'test.ome.zarr',
        shape=(100, 100, 100),
        chunks=(10, 10, 10),
        resolution=(.1, .1, .1),
        name='test',
        scale_factors=[[2, 2, 2],
                       [4, 4, 4]]
    )

    # write first slice in the first resolution layer
    f = zarr.open('test.ome.zarr')
    ds = f['s0']
    # note that the dataset is 5d with two leading singleton dimensions
    ds[0, 0, 0] = np.random.randint(0, 255, dtype='uint8', size=(100, 100))