Read IndicaLabs TIFF files (exported from HALO) and write TIFF file.

fix_indica_tiff.py

"""
Convert a TIFF file generated by IndicaLabs (HALO) to a usable TIFF file.
"""

import typing as tp
from pathlib import Path
from xml.etree import ElementTree

import tifffile
from tifffile.tiffcomment import tiffcomment
import zlib
import numpy as np


def indica_tiff_to_ome_tiff(
    indica_tiff: Path,
    output_tiff: Path | str,
    levels: list[int] = None,
    channels: list[int] = None,
    imwrite_kwargs: dict[str | tp.Any] = None,
) -> None:
    """
    Convert a TIFF file generated by IndicaLabs (HALO) to a usable TIFF file.

    Parameters
    ----------
    indica_tiff: pathlib.Path
        Input TIFF file generated by IndicaLabs.
    output_tiff: pathlib.Path
        Path to output TIFF file.
    levels: list[int]
        List of levels of pyramidal TIFF to convert.
        Default is to do all existing in the original file.
    channels: list[int]
        Integer index of channels to convert.
        Default is to do all existing in the original file.
    imwrite_kwargs: dict[str | Any]
        Additional keyword arguments to pass to `tifffile.TiffWriter.write`.
    """
    f = tifffile.TiffFile(file, is_ome=False)

    if isinstance(output_tiff, str):
        output_tiff = Path(output_tiff)

    if levels is None:
        levels = f.series[0].levels
    else:
        # convert between int and tiffifle.TiffPathSeries obj
        levels = [l for i, l in enumerate(f.series[0].levels) if i in levels]
    if channels is None:
        channels = np.arange(levels[0].shape[0])

    if imwrite_kwargs is None:
        imwrite_kwargs = dict(compression=8)
    channel_names = np.asarray(get_channel_names_from_indica_xml(file))

    with tifffile.TiffWriter(output_tiff, bigtiff=True) as tiff:
        with open(file, "rb") as fh:
            for level in levels:
                # pages = f.series[0].levels[level.index]
                stack = np.zeros(level.shape, dtype=np.float32)
                for ch, p in enumerate(level.pages):
                    if ch not in channels:
                        continue
                    for offset, _bytes, (_, (_, _, y, x, _), _) in zip(
                        p.dataoffsets, p.databytecounts, p.segments()
                    ):
                        i = f.filehandle.read_segments(offset, _bytes)

                        fh.seek(offset)
                        decompressed = zlib.decompress(fh.read(_bytes))
                        tile = np.frombuffer(decompressed, dtype="float32").reshape(
                            (p.tilewidth, p.tilelength)
                        )
                        # make sure not to go over boundaries in border tiles
                        y, x = slice(y, y + p.tilewidth), slice(x, x + p.tilelength)
                        new_shape = stack[ch, y, x].shape
                        stack[ch, y, x] = tile[: new_shape[0], : new_shape[1]]

            tiff.write(
                data=stack,
                tile=(p.tilewidth, p.tilelength),
                subifds=len(levels),
                metadata=dict(Channel=dict(Name=channel_names[channels].tolist())),
                **imwrite_kwargs,
            )


def get_channel_names_from_indica_xml(file: Path) -> list[str]:
    """
    Extract the name of channels in IndicaLabs TIFF file.

    Parameters
    ----------
    file: pathlib.Path
        Path to TIFF file.
    """
    xml = tiffcomment(file)
    root = ElementTree.fromstring(xml.replace("\n", "").replace("\t", ""))
    # channels = [x.attrib for x in root.iter() if x.tag == 'channel']
    channels = [x.get("name") for x in root.iter() if x.tag == "channel"]
    return channels