mosaic.py

"""
Convert a collection of tiff files generated by the LSM pipeline into a Zarr.

Example input files can be found at
https://lincbrain.org/dandiset/000004/0.240319.1924/files?location=derivatives%2F
"""

import logging

# stdlib
import os
import re
from glob import glob

# externals
import cyclopts
from tifffile import TiffFile

# internals
from linc_convert.modalities.lsm.cli import lsm
from linc_convert.utils.io.zarr import from_config
from linc_convert.utils.nifti_header import build_nifti_header
from linc_convert.utils.zarr_config import (
    GeneralConfig,
    NiftiConfig,
    ZarrConfig,
    autoconfig,
)

logger = logging.getLogger(__name__)
mosaic = cyclopts.App(name="mosaic", help_format="markdown")
lsm.command(mosaic)


@mosaic.default
@autoconfig
def convert(
    inp: str,
    *,
    voxel_size: list[float] = (1, 1, 1),
    general_config: GeneralConfig = None,
    zarr_config: ZarrConfig = None,
    nii_config: NiftiConfig = None,
) -> None:
    """
    Convert a collection of tiff files generated by the LSM pipeline into ZARR.

    Parameters
    ----------
    inp
        Path to the root directory, which contains a collection of
        subfolders named `*_z{:02d}_y{:02d}*`, each containing a
        collection of files named `*_plane{:03d}_c{:d}.tiff`.
    voxel_size
        Voxel size along the X, Y and Z dimension, in micron.
    general_config
        General configuration
    zarr_config
        Zarr related configuration
    nii_config
        NIfTI header related configuration
    """
    max_load = general_config.max_load
    if max_load % 2:
        max_load += 1

    CHUNK_PATTERN = re.compile(
        r"^(?P<prefix>\w*)"
        r"_z(?P<z>[0-9]+)"
        r"_y(?P<y>[0-9]+)"
        r"("
        r"?P<suffix>\w*)$"
    )

    all_chunks_dirnames = list(sorted(glob(os.path.join(inp, "*_z*_y*"))))
    all_chunks_info = dict(
        dirname=[],
        prefix=[],
        suffix=[],
        z=[],
        y=[],
        planes=[
            dict(
                fname=[],
                z=[],
                c=[],
                yx_shape=[],
            )
            for _ in range(len(all_chunks_dirnames))
        ],
    )

    # parse all directory names
    for dirname in all_chunks_dirnames:
        parsed = CHUNK_PATTERN.fullmatch(os.path.basename(dirname))
        all_chunks_info["dirname"].append(dirname)
        all_chunks_info["prefix"].append(parsed.group("prefix"))
        all_chunks_info["suffix"].append(parsed.group("suffix"))
        all_chunks_info["z"].append(int(parsed.group("z")))
        all_chunks_info["y"].append(int(parsed.group("y")))

    # default output name
    default_name = all_chunks_info["prefix"][0] + all_chunks_info["suffix"][0]
    general_config.set_default_name(default_name)

    # parse all individual file names
    nchunkz = max(all_chunks_info["z"])
    nchunky = max(all_chunks_info["y"])
    allshapes = [[(0, 0, 0) for _ in range(nchunky)] for _ in range(nchunkz)]
    nchannels = 0
    for zchunk in range(nchunkz):
        for ychunk in range(nchunky):
            for i in range(len(all_chunks_info["dirname"])):
                if (
                    all_chunks_info["z"][i] == zchunk + 1
                    and all_chunks_info["y"][i] == ychunk + 1
                ):
                    break
            dirname = all_chunks_info["dirname"][i]
            planes_filenames = list(sorted(glob(os.path.join(dirname, "*.tiff"))))

            PLANE_PATTERN = re.compile(
                os.path.basename(dirname) + r"_plane(?P<z>[0-9]+)"
                                            r"_c(?P<c>[0-9]+)"
                                            r".tiff$"
            )

            for fname in planes_filenames:
                parsed = PLANE_PATTERN.fullmatch(os.path.basename(fname))
                all_chunks_info["planes"][i]["fname"] += [fname]
                all_chunks_info["planes"][i]["z"] += [int(parsed.group("z"))]
                all_chunks_info["planes"][i]["c"] += [int(parsed.group("c"))]

                f = TiffFile(fname)
                yx_shape = f.pages[0].shape
                all_chunks_info["planes"][i]["yx_shape"].append(yx_shape)

            nplanes = max(all_chunks_info["planes"][i]["z"])
            nchannels = max(nchannels, max(all_chunks_info["planes"][i]["c"]))

            yx_shape = set(all_chunks_info["planes"][i]["yx_shape"])
            if not len(yx_shape) == 1:
                raise ValueError("Incompatible chunk shapes")
            yx_shape = list(yx_shape)[0]
            allshapes[zchunk][ychunk] = (nplanes, *yx_shape)

    # check that all chunk shapes are compatible
    for zchunk in range(nchunkz):
        if len(set(shape[1] for shape in allshapes[zchunk])) != 1:
            raise ValueError("Incompatible Y shapes")
    for ychunk in range(nchunky):
        if len(set(shape[ychunk][0] for shape in allshapes)) != 1:
            raise ValueError("Incompatible Z shapes")
    if len(set(shape[2] for subshapes in allshapes for shape in subshapes)) != 1:
        raise ValueError("Incompatible X shapes")

    # compute full shape
    fullshape = [0, 0, 0]
    fullshape[0] = sum(shape[0][0] for shape in allshapes)
    fullshape[1] = sum(shape[1] for shape in allshapes[0])
    fullshape[2] = allshapes[0][0][2]

    omz = from_config(general_config.out, zarr_config)
    arr = omz.create_array(
        "0", shape=[nchannels, *fullshape], dtype="float64", zarr_config=zarr_config
    )

    print("Write level 0 with shape", [nchannels, *fullshape])
    for i, dirname in enumerate(all_chunks_info["dirname"]):
        chunkz = all_chunks_info["z"][i] - 1
        chunky = all_chunks_info["y"][i] - 1
        planes = all_chunks_info["planes"][i]
        for j, fname in enumerate(planes["fname"]):
            subz = planes["z"][j] - 1
            subc = planes["c"][j] - 1
            yx_shape = planes["yx_shape"][j]

            zstart = sum(shape[0][0] for shape in allshapes[:chunkz])
            ystart = sum(shape[1] for shape in allshapes[chunkz][:chunky])
            print(
                f"Write plane "
                f"({subc}, {zstart + subz}, {ystart}:{ystart + yx_shape[0]})",
                end="\r",
            )
            slicer = (
                subc,
                zstart + subz,
                slice(ystart, ystart + yx_shape[0]),
                slice(None),
            )

            f = TiffFile(fname)
            arr[slicer] = f.asarray()
    print("")
    omz.generate_pyramid(mode="mean")

    # Write OME-Zarr multiscale metadata
    print("Write metadata")
    voxel_size = list(map(float, reversed(voxel_size)))
    omz.write_ome_metadata(["c", "z", "y", "x"], voxel_size)

    if nii_config.nii:
        header = build_nifti_header(
            zgroup=omz,
            voxel_size_zyx=tuple(voxel_size),
            unit="micrometer",
            nii_config=nii_config,
        )
        omz.write_nifti_header(header)

    logger.info("Conversion complete.")