Merge pull request #151 from BioimageAnalysisCoreWEHI/ome-tiff-writer

- OME-Tiff writer so we can write compressed ome-tif. 
- Modified writer class and tiffwriter class so it can be written to a single ome-tif file
- Handle zarr v2 and v3 deprecations

Author: pr4deepr

core/lls_core/models/output.py

@@ -38,7 +38,7 @@ class OutputParams(FieldAccessModel):
     )
     save_type: SaveFileType = Field(
         default=SaveFileType.h5,
-        description=f"The data type to save the result as. This will also be used to determine the file extension of the output files. Choices: `{enum_choices(SaveFileType)}`. Choices can alternatively be specifed as `str`, for example `'tiff'`."
+        description=f"The data type to save the result as. This will also be used to determine the file extension of the output files. Choices: `{enum_choices(SaveFileType)}`. Choices can alternatively be specifed as `str`, for example `'tiff'`. Note: `tiff` is saved as a compressed OME-TIFF (`.ome.tif`)."
     )
     time_range: range = Field(
         default=None,
@@ -102,7 +102,9 @@ def file_extension(self):
         elif self.save_type == SaveFileType.omezarr:
             return "ome.zarr"
         else:
-            return "tif"
+            # TIFF is written as a compressed OME-TIFF by default. The legacy
+            # uncompressed path renames itself back to plain ".tif".
+            return "ome.tif"
 
     def make_filepath(self, suffix: str) -> Path:
         """

core/lls_core/models/results.py

@@ -95,9 +95,7 @@ def save_image(self):
         Writer = self.lattice_data.get_writer()
         for roi, roi_results in groupby(self.slices, key=lambda it: it.roi_index):
             writer = Writer(self.lattice_data, roi_index=roi)
-            for slice in roi_results:
-                writer.write_slice(slice)
-            writer.close()
+            writer.write_all(roi_results)
 
 class ProcessedWorkflowOutput(BaseModel, arbitrary_types_allowed=True):
     """

core/lls_core/utils.py

@@ -15,6 +15,12 @@
 
 from . import DeskewDirection, config
 
+from packaging.version import parse as parse_version
+
+#Get zarr major version here and single source of truth for rest of the codebase
+import zarr
+ZARR_MAJOR_VERSION = parse_version(zarr.__version__).major  #get the major zarr vers
+
 if TYPE_CHECKING:
     from xml.etree.ElementTree import Element
     from dask.array.core import Array as DaskArray
@@ -349,9 +355,7 @@ def get_zarr_compression():
     Returns zarr compression settings depending on zarr version
     zarrv3 does not accept numcodecs compressors directly whereas zarrv <3 does
     """
-    from packaging.version import parse
-    import zarr
-    if parse(zarr.__version__).major >= 3:
+    if ZARR_MAJOR_VERSION >= 3:
         from zarr.codecs import BloscCodec
         return dict(
             compressors=[BloscCodec(cname="zstd", clevel=5, shuffle="shuffle")]

core/lls_core/writers.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Iterable, List, Optional
 
 from lls_core.types import ArrayLike
 
@@ -15,9 +15,26 @@
 import numpy as np
 import zarr
 
-from lls_core.utils import make_filename_suffix, get_zarr_compression
+from lls_core.utils import make_filename_suffix, get_zarr_compression, ZARR_MAJOR_VERSION
 RoiIndex = Optional[NonNegativeInt]
 
+def resolve_output_dtype(dtype: np.dtype) -> np.dtype:
+    """Pick the output dtype: keep float and small ints, cast larger ints (int32, int64) to uint16."""
+    dtype = np.dtype(dtype)
+    if np.issubdtype(dtype, np.integer):
+        return dtype if np.iinfo(dtype).max < np.iinfo(np.uint16).max else np.dtype(np.uint16)
+    if np.issubdtype(dtype, np.floating):
+        return dtype
+    raise TypeError(f"Unsupported data dtype: {dtype}")
+
+
+def to_output_dtype(array: np.ndarray, out_dtype: np.dtype) -> np.ndarray:
+    """Cast to ``out_dtype``, clipping to range only when casting to uint16."""
+    out_dtype = np.dtype(out_dtype)
+    if out_dtype == np.uint16 and array.dtype != np.uint16:
+        return np.clip(array, 0.0, 65535.0).astype(np.uint16)
+    return array.astype(out_dtype, copy=False)
+
 if TYPE_CHECKING:
     from lls_core.models.lattice_data import LatticeData
     import npy2bdv
@@ -49,6 +66,17 @@ def close(self):
         """
         pass
 
+    def write_all(self, slices: Iterable[ProcessedSlice[ArrayLike]]) -> None:
+        """
+        Write each slice, then finish the output.
+
+        By default, each slice is passed to ``write_slice`` in order, and
+        ``close`` is called at the end. Override this for custom behavior.
+        """
+        for slice in slices:
+            self.write_slice(slice)
+        self.close()
+
 @dataclass
 class BdvWriter(Writer):
     """
@@ -95,45 +123,89 @@ def close(self):
 @dataclass
 class TiffWriter(Writer):
     """
-    A writer for for TIFF output format
+    A writer for for TIFF output format.
+
+    By default write a deflate-compressed OME-TIFF (compression='zlib),
+    which keeps the compresses empty/black space at borders of deskewed images keeping
+    Fiji/Bio-Formats readability. Set compression=None to fall back to the legacy
+    uncompressed ImageJ TIFF.
     """
     pending_slices: List[ImageSlice] = field(default_factory=list, init=False)
     time: Optional[NonNegativeInt] = None
+    #: tifffile compression codec for the OME-TIFF output. ``'zlib'`` (deflate) is
+    #: the default; ``None`` writes an uncompressed ImageJ-TIFF as before.
+    compression: Optional[str] = "zlib"
 
     def __post_init__(self):
         self.pending_slices = []
-    
+
     def flush(self):
         "Write out all pending slices"
         import numpy as np
         import tifffile
-        if len(self.pending_slices) > 0:
-            first_result = self.pending_slices[0]
-            images_array = np.swapaxes(
-                np.expand_dims([result.data for result in self.pending_slices], axis=0),
-                1, 2
-            ).astype("uint16")
-            # ImageJ TIFF can only handle 16-bit uints, not 32
-            path = self.lattice.make_filepath(
-                make_filename_suffix(
-                    channel=first_result.channel,
-                    time=first_result.time,
-                    roi_index=first_result.roi_index
-                )
+        if len(self.pending_slices) == 0:
+            return
+
+        first_result = self.pending_slices[0]
+        # One buffered slice per channel for this timepoint; each is a (Z, Y, X)
+        # volume. flush() is called once the timepoint changes, so T == 1 here.
+        channel_arrays = [np.asarray(result.data) for result in self.pending_slices]
+
+        # Holds every channel for this timepoint (TCZYX / TZCYX), so
+        # name will just be by ROI and time
+        path = self.lattice.make_filepath(
+            make_filename_suffix(
+                time=first_result.time,
+                roi_index=first_result.roi_index
             )
+        )
+
+        if self.compression is None:
+            # Legacy uncompressed ImageJ-TIFF (TZCYX). ImageJ-TIFF cannot be
+            # compressed, so this path only exists as a fallback. It is not an
+            # OME-TIFF, so drop the ".ome" from the default ".ome.tif" name.
+            if path.name.endswith(".ome.tif"):
+                path = path.with_name(path.name[:-len(".ome.tif")] + ".tif")
+            images_array = np.swapaxes(
+                np.expand_dims(channel_arrays, axis=0), 1, 2
+            ).astype("uint16")  # ImageJ TIFF can only handle 16-bit uints, not 32
             tifffile.imwrite(
                 str(path),
-                data = images_array,
+                data=images_array,
                 bigtiff=True,
-                resolution=(1./self.lattice.dx, 1./self.lattice.dy),
+                resolution=(1. / self.lattice.dx, 1. / self.lattice.dy),
                 resolutionunit="MICROMETER",
                 metadata={'spacing': self.lattice.new_dz, 'unit': 'um', 'axes': 'TZCYX'},
                 imagej=True
             )
-            self.written_files.append(path)
-            
-            # Reinitialise
-            self.pending_slices = []
+        else:
+            # Compressed OME-TIFF. The full timepoint is already buffered in
+            # channel_arrays, so assemble the (T, C, Z, Y, X) array (T == 1) and
+            # write it with tifffile's native OME support, which lays out the
+            # pages and OME-XML correctly for any Z depth. imagej=True is
+            # incompatible with both compression and OME; physical pixel sizes
+            # and channel names travel in the OME-XML instead.
+
+            # Same dtype policy as OMEZarrWriter: preserve small ints and float,
+            # standardise wider integers to uint16.
+            out_dtype = resolve_output_dtype(np.result_type(*channel_arrays))
+            stack = to_output_dtype(np.stack(channel_arrays, axis=0), out_dtype)  # (C, Z, Y, X)
+            data5d = stack[np.newaxis, ...]  # (T=1, C, Z, Y, X), matches TCZYX
+
+            ome_metadata = {
+                "axes": "TCZYX",
+                "PhysicalSizeX": float(self.lattice.dx), "PhysicalSizeXUnit": "µm",
+                "PhysicalSizeY": float(self.lattice.dy), "PhysicalSizeYUnit": "µm",
+                "PhysicalSizeZ": float(self.lattice.new_dz), "PhysicalSizeZUnit": "µm",
+                "Channel": {"Name": [str(result.channel) for result in self.pending_slices]},
+            }
+            with tifffile.TiffWriter(str(path), bigtiff=True, ome=True) as tw:
+                tw.write(data5d, compression=self.compression, metadata=ome_metadata)
+
+        self.written_files.append(path)
+
+        # Reinitialise
+        self.pending_slices = []
 
     def write_slice(self, slice: ProcessedSlice[ArrayLike]):
         if slice.time != self.time:
@@ -145,6 +217,83 @@ def write_slice(self, slice: ProcessedSlice[ArrayLike]):
     def close(self):
         self.flush()
 
+    def write_all(self, slices: Iterable[ProcessedSlice[ArrayLike]]) -> None:
+        """
+        Stream every timepoint/channel into a single compressed OME-TIFF.
+
+        tifffile only makes one OME series per ``write()`` call, so we stream
+        planes from the slice iterator instead of buffering full volumes.
+        The uncompressed path keeps the old per-timepoint behavior.
+        """
+        # Legacy ImageJ-TIFF: unchanged, one file per timepoint.
+        if self.compression is None:
+            super().write_all(slices) # default behaviour
+            return
+
+        import numpy as np
+        import tifffile
+        #get whole slice iterator for write_all
+        it = iter(slices)
+        
+        first = next(it, None)
+        if first is None:
+            # Empty ROI: write nothing.
+            return
+
+        first_vol = np.asarray(first.data)
+        if first_vol.ndim != 3:
+            raise ValueError(f"Expected (Z, Y, X) slice, got shape {first_vol.shape}")
+
+        # Dtype policy matches OMEZarrWriter: fixed from the first slice.
+        out_dtype = resolve_output_dtype(first_vol.dtype)
+        z_len, y_len, x_len = (int(d) for d in first_vol.shape)
+        t_len = len(self.lattice.time_range)
+        c_len = len(self.lattice.channel_range)
+
+        # One file per ROI, named like the other writers (no timepoint/channel
+        # in the name, since the file holds them all).
+        suffix = f"_{make_filename_suffix(roi_index=str(self.roi_index))}" if self.roi_index is not None else ""
+        path = self.lattice.make_filepath(suffix)
+
+        ome_metadata = {
+            "axes": "TCZYX",
+            "PhysicalSizeX": float(self.lattice.dx), "PhysicalSizeXUnit": "µm",
+            "PhysicalSizeY": float(self.lattice.dy), "PhysicalSizeYUnit": "µm",
+            "PhysicalSizeZ": float(self.lattice.new_dz), "PhysicalSizeZUnit": "µm",
+            "Channel": {"Name": [str(ch) for ch in self.lattice.channel_range]},
+        }
+
+        def plane_generator():
+            # first slice is already loaded, so cast it once and write its
+            # Z planes. Then load and cast the rest. Slices come in time order
+            # and then channel order, so this matches the (t, c, z) order that
+            # tifffile expects for shape=(T,C,Z,Y,X).
+            first_cast = to_output_dtype(first_vol, out_dtype)
+            #load and cast rest of slices
+            for z in range(first_cast.shape[0]):
+                yield first_cast[z]
+            #stream rest of slices
+            for sl in it:
+                vol = to_output_dtype(np.asarray(sl.data), out_dtype)
+                if vol.shape != (z_len, y_len, x_len):
+                    raise ValueError(
+                        f"Inconsistent slice shape {vol.shape}; expected {(z_len, y_len, x_len)}"
+                    )
+                for z in range(vol.shape[0]):
+                    yield vol[z]
+
+        # A plane count != T*C*Z makes tw.write raise, so a partial or misordered
+        # run fails loudly instead of writing a silently wrong file.
+        with tifffile.TiffWriter(str(path), bigtiff=True, ome=True) as tw:
+            tw.write(
+                plane_generator(),
+                shape=(t_len, c_len, z_len, y_len, x_len),
+                dtype=out_dtype,
+                compression=self.compression,
+                metadata=ome_metadata,
+            )
+        self.written_files.append(path)
+
 @dataclass
 class OMEZarrWriter(Writer):
     DEFAULT_CHUNK_ZYX = (64, 256, 256)
@@ -183,7 +332,7 @@ def __init__(
         self._zyx = None
         self._t_len = None
         self._c_len = None
-        self._dtype = np.uint16 #We are enforcing 16-bit, but may change in future
+        self._dtype = np.uint16 #Placeholder; resolved per data in write_slice
 
         self._pix_z, self._pix_y, self._pix_x = (self.lattice.new_dz, self.lattice.dy, self.lattice.dx)
 
@@ -196,26 +345,21 @@ def write_slice(self, slice) -> Path:
         if self._zyx is None:
             self._zyx = (int(data3d.shape[0]), int(data3d.shape[1]), int(data3d.shape[2]))
 
-        #if dtype of data is < uint16, use the data dtype
-        if np.issubdtype(data3d.dtype, np.integer):
-            self._dtype = (data3d.dtype
-                            if np.iinfo(data3d.dtype).max < np.iinfo(np.uint16).max
-                            else np.uint16)
-        elif np.issubdtype(data3d.dtype, np.floating):
-            #float data, so preserve dtype
-            self._dtype = data3d.dtype
-        else:
-            raise TypeError(f"Unsupported data dtype: {data3d.dtype}")
-                
+        # Same dtype policy as TiffWriter: preserve small ints and float,
+        # standardise wider integers to uint16. Preserving float means label/
+        # mask images (typically float32) keep their values instead of being
+        # clipped to 16-bit.
+        self._dtype = resolve_output_dtype(data3d.dtype)
+
         t_idx = int(getattr(slice, "time_index", 0))
         c_idx = int(getattr(slice, "channel_index", 0))
         t_len, c_len = self._resolve_t_c_lengths(slice)
 
         # If it's the first slice - initialize the full zarr array size
         if self._arr is None:
             self._root_group, self._arr = self._create_store(t_len, c_len, self._zyx, self._dtype)
-        
-        self._arr[t_idx, c_idx, :, :, :] = np.clip(data3d, 0.0, 65535.0).astype(np.uint16)
+
+        self._arr[t_idx, c_idx, :, :, :] = to_output_dtype(data3d, self._arr.dtype)
         return self._root_path
 
     # Optional hook if the framework ever calls it.
@@ -238,27 +382,26 @@ def _create_store(
             import shutil
             shutil.rmtree(self._root_path)
 
-        chunks = (1, 1, *self.chunk_zyx)
-
-        compression_kwargs = get_zarr_compression()
-        #adding compatibility fix for zarr v2 and v3
-        if int(zarr.__version__.split(".")[0]) >= 3:
-            #zarr v3: group class cannot be constructed from path directly
-            root = zarr.open_group(store=str(self._root_path), mode="a")
-        else:
-            store = zarr.DirectoryStore(str(self._root_path))
-            root = zarr.group(store=store)
-
         dataset_kwargs = {
             "shape": (t_len, c_len, zyx[0], zyx[1], zyx[2]),
-            "chunks": chunks,
+            "chunks": (1, 1, *self.chunk_zyx),
             "dtype": dtype,
-            **compression_kwargs,
+            **get_zarr_compression(),
         }
-        if int(zarr.__version__.split(".")[0]) < 3:
+
+        # Single version check for both group and array creation 
+        if ZARR_MAJOR_VERSION >= 3:
+            # zarr v3: group cannot be constructed from a path directly, and
+            # create_array is the current API (create_dataset is deprecated).
+            root = zarr.open_group(store=str(self._root_path), mode="a")
+            arr = root.create_array("0", **dataset_kwargs)
+        else:
+            # zarr v2: build  group from a DirectoryStore; the group has no
+            # create_array, so create_dataset is the equivalent.
+            root = zarr.group(store=zarr.DirectoryStore(str(self._root_path)))
             dataset_kwargs["overwrite"] = self.overwrite
+            arr = root.create_dataset("0", **dataset_kwargs)
 
-        arr = root.create_dataset("0", **dataset_kwargs)
         self._write_ngff_attrs(root)
         return root, arr