Merge pull request #43 from WEHI-SODA-Hub/fix/label-mask-dtype-overflow

fix: label masks silently truncate past 65,535 cells

Author: mikemcka

.github/workflows/linting.yml

@@ -29,8 +29,10 @@ jobs:
 
   pytest:
     # Unit tests for the bin/ scripts. Deliberately dependency-light: these cover
-    # the pure logic (channel naming, exclusion, the marker report), so they need
-    # no torch, no GPU and no gated model weights, and run in seconds.
+    # the pure logic (channel naming, exclusion, the marker report, the label
+    # dtype), so they need no torch, no GPU and no gated model weights, and
+    # run in seconds. The geo stack is here only because bin/parquet_to_tiff.py
+    # imports it at module level; all of it is plain wheels.
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@0ad4b8fadaa221de15dcec353f45205ec38ea70b # v4
@@ -41,7 +43,7 @@ jobs:
           python-version: "3.12"
 
       - name: Install test dependencies
-        run: pip install pytest numpy tifffile
+        run: pip install pytest numpy tifffile rasterio shapely pyarrow typer
 
       - name: Run pytest
         run: pytest tests/python -q

bin/cellsam_segment.py

@@ -11,6 +11,13 @@
 import cellSAM
 
 
+# This is an intermediate mask, smoothed downstream, so compress it: label masks
+# are long runs of a repeated ID and deflate shrinks them ~90x. Only the final
+# smoothed mask is left uncompressed for viewers.
+MASK_COMPRESSION = "zlib"
+MASK_COMPRESSION_ARGS = {"level": 1}
+
+
 def _local_tag_name(tag):
     '''Return XML tag name without namespace.'''
     return tag.split('}', 1)[-1]
@@ -397,7 +404,10 @@ def main():
         mask = remove_small_cells(mask, args.min_area)
 
     # Save segmentation mask
-    tifffile.imwrite(args.output, mask.astype(np.uint32))
+    tifffile.imwrite(
+        args.output, mask.astype(np.uint32),
+        compression=MASK_COMPRESSION, compressionargs=MASK_COMPRESSION_ARGS,
+    )
     print(f"Segmentation mask saved to {args.output}")
 
 

bin/parquet_to_tiff.py

@@ -23,6 +23,17 @@
 from typing import Annotated
 
 
+# Label IDs are pixel values, so the dtype has to index every cell. uint32
+# covers 4.29e9 cells -- orders of magnitude past any real slide -- so one wide
+# dtype avoids the silent truncation uint16 hits beyond 65535 cells (rasterize()
+# caps at the dtype maximum, it never raises). rasterize() has no uint64.
+LABEL_DTYPE = np.uint32
+
+# Label masks are long runs of a repeated ID, so deflate shrinks them ~90x.
+MASK_COMPRESSION = "zlib"
+MASK_COMPRESSION_ARGS = {"level": 1}
+
+
 def get_image_dimensions(tiff_path: Path) -> tuple[int, int]:
     """
     Extracts image dimensions from the OME-TIFF metadata.
@@ -67,6 +78,15 @@ def main(
     # Create incrementing IDs for each geometry (1-based as 0 is background)
     ids = [id for id in range(1, len(geometries) + 1)]
 
+    # Guard the (physically implausible) case of more cells than uint32 can
+    # index, so we never silently truncate the way uint16 did.
+    if len(geometries) > np.iinfo(LABEL_DTYPE).max:
+        raise ValueError(
+            f"{len(geometries)} geometries exceeds the "
+            f"{np.iinfo(LABEL_DTYPE).max} label limit of "
+            f"{np.dtype(LABEL_DTYPE).name} masks."
+        )
+
     # Get X and Y dimensions from the TIFF file
     (size_x, size_y) = get_image_dimensions(tiff_path)
 
@@ -84,12 +104,25 @@ def main(
             width=size_x,
             height=size_y
         ),
-        dtype=np.uint16,
+        dtype=LABEL_DTYPE,
         all_touched=True  # Fill shapes
     )
 
+    # stdout carries the TIFF, so diagnostics go to stderr.
+    max_label = int(mask.max())
+    print(
+        f"Rasterised {len(geometries)} geometries as {mask.dtype.name}; "
+        f"max label {max_label}",
+        file=sys.stderr,
+    )
+
     # Write TIFF output
-    imwrite(sys.stdout.buffer, np.flipud(mask))
+    imwrite(
+        sys.stdout.buffer,
+        np.flipud(mask),
+        compression=MASK_COMPRESSION,
+        compressionargs=MASK_COMPRESSION_ARGS,
+    )
 
 
 if __name__ == "__main__":

modules/local/cellmeasurement/environment.yml

@@ -8,4 +8,4 @@ dependencies:
   - uv
   - pip
   - pip:
-      - git+https://github.com/WEHI-SODA-Hub/cellmeasurement-py@0.2.1
+      - git+https://github.com/WEHI-SODA-Hub/cellmeasurement-py@0.2.2

modules/local/cellmeasurement/main.nf

@@ -3,7 +3,7 @@ process CELLMEASUREMENT {
     label 'process_multi'
 
     conda "${moduleDir}/environment.yml"
-    container 'ghcr.io/wehi-soda-hub/cellmeasurement-py:0.2.1'
+    container 'ghcr.io/wehi-soda-hub/cellmeasurement-py:0.2.2'
 
     input:
     tuple val(meta),

tests/python/test_parquet_to_tiff.py

@@ -0,0 +1,90 @@
+"""Unit tests for label-dtype selection in ``bin/parquet_to_tiff.py``.
+
+Label IDs are stored as pixel values, so a mask dtype too narrow for the largest
+ID silently truncates the segmentation. This is worth pinning because there is
+no error to notice: ``rasterio.features.rasterize`` does not raise on overflow,
+it caps at the dtype maximum and every cell past that point vanishes. It only
+bites on whole-slide images, which is exactly where nobody is watching.
+
+Pure python -- no pipeline run required.
+"""
+
+import numpy as np
+import rasterio.features
+from shapely.geometry import box
+
+import tifffile
+
+from parquet_to_tiff import (
+    LABEL_DTYPE,
+    MASK_COMPRESSION,
+    MASK_COMPRESSION_ARGS,
+)
+
+UINT16_MAX = np.iinfo(np.uint16).max  # 65535
+UINT32_MAX = np.iinfo(np.uint32).max
+
+
+def _label_shapes(n):
+    """``n`` disjoint 1px cells with 1-based IDs, laid out in a 100-wide grid."""
+    return ((box(i % 100, i // 100, i % 100 + 0.9, i // 100 + 0.9), i) for i in range(1, n + 1))
+
+
+def _rasterize(n, dtype):
+    height = (n // 100) + 2
+    return rasterio.features.rasterize(
+        _label_shapes(n), out_shape=(height, 110), fill=0, dtype=dtype, all_touched=True
+    )
+
+
+# ----------------------------------------------------------------------------
+# the label dtype
+# ----------------------------------------------------------------------------
+
+
+def test_label_dtype_indexes_far_more_cells_than_any_real_slide():
+    """uint32 is the wide fixed choice: 4.29e9 IDs, well past any real mask."""
+    assert np.dtype(LABEL_DTYPE).name == "uint32"
+    assert np.iinfo(LABEL_DTYPE).max == UINT32_MAX
+
+
+# ----------------------------------------------------------------------------
+# the truncation this guards against
+# ----------------------------------------------------------------------------
+
+
+def test_uint16_silently_truncates_beyond_its_maximum():
+    """Documents the failure mode uint32 avoids: no exception, missing cells."""
+    mask = _rasterize(UINT16_MAX + 2, np.uint16)
+
+    assert mask.max() == UINT16_MAX
+    assert len(np.unique(mask)) - 1 == UINT16_MAX  # two labels lost, silently
+
+
+def test_label_dtype_preserves_every_label_past_uint16():
+    n = UINT16_MAX + 2
+    mask = _rasterize(n, LABEL_DTYPE)
+
+    assert mask.dtype == np.uint32
+    assert mask.max() == n
+    assert len(np.unique(mask)) - 1 == n
+
+
+# ----------------------------------------------------------------------------
+# compression
+# ----------------------------------------------------------------------------
+
+
+def test_mask_compression_is_lossless_and_smaller(tmp_path):
+    """These masks are written to and re-read from shared scratch, so raw byte
+    count is the cost. Compression must not change a single label."""
+    mask = _rasterize(5000, np.uint32)
+    path = tmp_path / "mask.tiff"
+
+    tifffile.imwrite(
+        path, mask,
+        compression=MASK_COMPRESSION, compressionargs=MASK_COMPRESSION_ARGS,
+    )
+
+    assert path.stat().st_size < mask.nbytes
+    assert np.array_equal(tifffile.imread(path), mask)