kerchunk_tiff_references.py

# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "fsspec==2025.5.1",
#     "marimo",
#     "matplotlib==3.10.3",
#     "pooch==1.8.2",
#     "scikit-image==0.25.2",
#     "tifffile==2025.6.1",
#     "zarr==3.0.8",
# ]
#
# [tool.uv]
# exclude-newer = "2025-06-03T21:22:39.068415-04:00"
# ///

import marimo

__generated_with = "0.13.15"
app = marimo.App(width="medium")


@app.cell
def _(mo):
    import skimage

    cells = skimage.data.cells3d()

    mo.md(f"Image shape: **{', '.join(f'{k}={v}' for k, v in zip('TCYX', cells.shape))}**")
    return (cells,)


@app.cell
def _(cells, mo):
    membranes = cells[:, 0, :, :]
    nuclei = cells[:, 1, :, :]

    mo.md("Separate the two channels: **membranes** and **nuclei**.")
    return membranes, nuclei


@app.cell
def _(membranes, mo, nuclei):
    import io
    import json
    import pathlib

    import tifffile


    def dump_kerchunk_reference(file: str, directory: pathlib.Path) -> dict:
        with tifffile.TiffFile(file) as tif, io.StringIO() as f:
            store = tif.aszarr()
            store.write_fsspec(f, url=directory.as_uri())
            refs = json.loads(f.getvalue())
        return refs


    def create_combined_kerchunk_reference(
        files: list[str], directory: pathlib.Path
    ) -> dict:
        """
        Generate a kerchunk reference that stacks TIFF files along a new
        leading dimension, updating shape and chunks accordingly.
        """
        assert len(files) > 1, "Need at least two files to combine."

        refs = {}
        ref_shape = None
        ref_chunks = None

        for i, file in enumerate(files):
            inner_refs = dump_kerchunk_reference(file, directory=directory)
            zarray = json.loads(inner_refs.pop(".zarray"))
            inner_refs.pop(".zattrs", None)

            if i == 0:
                # store reference shape/chunks from first file
                ref_shape = zarray["shape"]
                ref_chunks = zarray["chunks"]

                combined_zarray = zarray.copy()
                combined_zarray["shape"] = [len(files), *ref_shape]
                combined_zarray["chunks"] = [1, *ref_chunks]

                refs[".zarray"] = json.dumps(combined_zarray)
                refs[".zattrs"] = json.dumps(
                    {"_ARRAY_DIMENSIONS": ["C", "T", "Y", "X"]}
                )
            else:
                # check shape/chunks match the first file
                assert zarray["shape"] == ref_shape, f"Shape mismatch in {file}"
                assert zarray["chunks"] == ref_chunks, f"Chunk mismatch in {file}"

            for key, value in inner_refs.items():
                # prepend index as the new axis dim for the chunk
                refs[f"{i}.{key}"] = value

        return refs


    # write some tiffs to disk
    tifffile.imwrite("membranes.tiff", membranes, photometric="minisblack")
    tifffile.imwrite("nuclei.tiff", nuclei, photometric="minisblack")

    # make the combined reference
    combined_refs = create_combined_kerchunk_reference(
        files=["membranes.tiff", "nuclei.tiff"],
        directory=pathlib.Path(__file__).parent,
    )

    # (can/should write to disk)
    mo.md("Create a combined set of kerchunk references for some tiffs.")
    return (combined_refs,)


@app.cell
def _(combined_refs):
    import zarr
    from fsspec.implementations.reference import ReferenceFileSystem

    store = zarr.storage.FsspecStore(
        ReferenceFileSystem(combined_refs, asynchronous=True)
    )

    # get "lazy" zarr view of our store
    z_arr = zarr.open(store)
    z_arr
    return (z_arr,)


@app.cell(hide_code=True)
def _(cells, mo):
    c = mo.ui.radio({"membranes": 0, "nuclei": 1}, value="membranes", label="channel")
    t = mo.ui.slider(0, cells.shape[0], value=0, label="time")
    mo.vstack([c, t])
    return c, t


@app.cell
def _(c, t, z_arr):
    import matplotlib.pyplot as plt

    plt.axis("off")
    plt.imshow(z_arr[c.value, t.value, :, :])
    return


@app.cell
def _():
    import marimo as mo
    return (mo,)


if __name__ == "__main__":
    app.run()