patchworks 2.1.1__tar.gz → 2.1.2__tar.gz

{patchworks-2.1.1 → patchworks-2.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchworks
-Version: 2.1.1
+Version: 2.1.2
 Summary: Tiled processing of arbitrarily large images with globally consistent labels
 Project-URL: Homepage, https://github.com/imcf/patchworks
 Project-URL: Issues, https://github.com/imcf/patchworks/issues

{patchworks-2.1.1 → patchworks-2.1.2}/src/patchworks/plugins/ome_zarr.py

@@ -49,6 +49,7 @@ import glob
 import logging
 import math
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import nullcontext as _nullcontext
 from itertools import product as _iproduct
 from pathlib import Path
 from typing import Union
@@ -171,22 +172,41 @@ def _strides(axes: str, downscale: int) -> tuple[int, ...]:
     return tuple(downscale if a in _DOWNSAMPLE_AXES else 1 for a in axes)
 
 
-def _default_chunks(shape: tuple[int, ...], axes: str) -> tuple[int, ...]:
+def _default_chunks(
+    shape: tuple[int, ...],
+    axes: str,
+    source_chunks: Union[tuple[int, ...], None] = None,
+) -> tuple[int, ...]:
     """Bounded chunk shape so writing a level never blows up RAM.
 
+    With *source_chunks*, the result never spans more than one source chunk on
+    any axis. That matters when the source's own granularity is coarse: a
+    folder of stitched TIFFs gives one chunk per **file**, so a plane can be
+    gigabytes. Grouping 16 of those into one output chunk (the ``z`` cap) means
+    holding 16 whole planes to write a 32 MB chunk -- tens of GB for the
+    smallest unit of work, which is how conversion gets OOM-killed.
+
+    Splitting a source chunk is fine (that is a slice of something already
+    read); combining several is not.
+
     Parameters
     ----------
     shape : tuple of int
         Array shape.
     axes : str
         One letter per axis (selects the per-axis cap).
+    source_chunks : tuple of int, optional
+        The input array's own chunking.
 
     Returns
     -------
     tuple of int
-        Per-axis chunk size, capped by ``_CHUNK_CAP``.
+        Per-axis chunk size, capped by ``_CHUNK_CAP`` and by *source_chunks*.
     """
-    return tuple(min(s, _CHUNK_CAP.get(a, s)) for s, a in zip(shape, axes))
+    caps = [min(s, _CHUNK_CAP.get(a, s)) for s, a in zip(shape, axes)]
+    if source_chunks is not None:
+        caps = [min(c, max(1, int(sc))) for c, sc in zip(caps, source_chunks)]
+    return tuple(caps)
 
 
 ShardSpec = Union[bool, tuple[int, ...]]
@@ -446,6 +466,70 @@ def _stream_strided_level(
             pass
 
 
+# One chunk is one file without sharding. A shared cluster filesystem starts
+# to hurt well before a million of them, and the cost lands on every later
+# read too, not just the write.
+_CHUNK_COUNT_WARN = 200_000
+
+
+def _warn_if_many_chunks(
+    shape: tuple[int, ...], chunks: tuple[int, ...], shard: ShardSpec
+) -> None:
+    """Warn when a level will be written as a very large number of files.
+
+    A coarse source forces fine output chunks (see :func:`_default_chunks`),
+    and fine chunks mean many files. Sharding packs them into far fewer
+    without changing the chunking, so the fix is cheap -- but only if someone
+    notices before the store is written.
+    """
+    if shard:
+        return
+    count = int(np.prod([math.ceil(s / c) for s, c in zip(shape, chunks)]))
+    if count < _CHUNK_COUNT_WARN:
+        return
+    logger.warning(
+        "level 0 will be %s chunks, i.e. that many files: the source's own "
+        "granularity forced chunks of %s. Set shard=True (config: "
+        "shard: true) to pack them into far fewer files -- same chunking, "
+        "same memory, and kinder to a shared filesystem.",
+        f"{count:,}",
+        chunks,
+    )
+
+
+def _bounded_scheduler(arr: da.Array):
+    """Cap dask's threads by the memory one **source** chunk costs.
+
+    Dask keeps a source chunk alive while every task that reads it runs, so
+    peak memory tracks the *input* granularity, not the output's. A folder of
+    stitched TIFFs gives one chunk per file -- gigabytes each -- and the
+    default worker count is the machine's core count, so 32 threads x a 3.6 GB
+    plane is over 100 GB before anything is written.
+
+    Sizing from the source chunk against the job's real budget keeps that
+    bounded. Returns a context manager; a no-op when a distributed client is
+    driving, since it schedules with its own memory awareness.
+    """
+    import dask
+
+    from .._chunks import cpu_allocation, safe_worker_count
+    from .._cluster import _distributed_client
+
+    if _distributed_client() is not None:
+        return _nullcontext()
+    chunk_nbytes = int(np.prod(arr.chunksize)) * arr.dtype.itemsize
+    # Two live copies per worker: the source chunk plus the slice being written.
+    workers = min(
+        cpu_allocation(), safe_worker_count(chunk_nbytes, fn_overhead=2)
+    )
+    logger.info(
+        "writing with %d dask worker(s): source chunks are %.2f GB each",
+        workers,
+        chunk_nbytes / 1024**3,
+    )
+    return dask.config.set(scheduler="threads", num_workers=workers)
+
+
 def _progress_ctx(progress: bool, label: str):
     """Return a progress-bar context manager.
 
@@ -648,7 +732,12 @@ def _write_pyramid(
     strides = _strides(axes, downscale)
 
     if write_base:
-        base_chunks = chunks or _default_chunks(arr.shape, axes)
+        # Cap by the source's own chunking: combining several source chunks
+        # into one output chunk can mean holding gigabytes to write megabytes.
+        base_chunks = chunks or _default_chunks(
+            arr.shape, axes, source_chunks=arr.chunksize
+        )
+        _warn_if_many_chunks(arr.shape, base_chunks, shard)
         _to_zarr_level(
             arr.rechunk(base_chunks), group_path, base_name, shard, progress
         )
@@ -872,16 +961,75 @@ _UNIT_TO_UM = {
 _RESUNIT_TO_UM = {2: 25400.0, 3: 10000.0}
 
 
+def _ome_xml_pixel_size(xml: str) -> PixelSize:
+    """Voxel size from an OME-XML ``Pixels`` element.
+
+    Stitching and acquisition software commonly writes OME-XML into
+    ``ImageDescription`` and nothing else -- no ImageJ block, and resolution
+    tags left at their defaults. Readers that only look at the TIFF tags then
+    report an uncalibrated image even though the size is right there.
+
+    Parameters
+    ----------
+    xml : str
+        The OME-XML document.
+
+    Returns
+    -------
+    dict
+        ``{axis: micrometers}`` for whichever axes carry a size.
+    """
+    from xml.etree import ElementTree
+
+    try:
+        root = ElementTree.fromstring(xml)
+    except ElementTree.ParseError:
+        return {}
+    # The OME namespace is versioned, so match on the local tag name.
+    pixels = next(
+        (el for el in root.iter() if el.tag.rsplit("}", 1)[-1] == "Pixels"),
+        None,
+    )
+    if pixels is None:
+        return {}
+    out: PixelSize = {}
+    for axis in "zyx":
+        value = pixels.get(f"PhysicalSize{axis.upper()}")
+        if not value:
+            continue
+        unit = pixels.get(f"PhysicalSize{axis.upper()}Unit", "µm")
+        factor = _UNIT_TO_UM.get(str(unit).lower())
+        if factor is None:
+            logger.warning(
+                "OME-XML PhysicalSize%s has unit %r, which is not recognised; "
+                "ignoring it.",
+                axis.upper(),
+                unit,
+            )
+            continue
+        try:
+            out[axis] = float(value) * factor
+        except ValueError:
+            continue
+    return out
+
+
 def _tiff_pixel_size(path: str) -> PixelSize:
     """Read physical voxel size from a TIFF file's own metadata.
 
-    Z comes from ImageJ metadata (``spacing`` + ``unit``), if present. X/Y
-    come from the page's ``XResolution``/``YResolution`` tags (pixels per
-    unit); the unit itself is taken from the ``ResolutionUnit`` tag when set
-    (plain TIFFs), or falls back to ImageJ metadata's ``unit`` — ImageJ
-    itself writes ``ResolutionUnit=NONE`` and keeps the unit as a string in
-    its metadata instead. Unrecognized units are ignored (treated as
-    uncalibrated).
+    Checks every place the size is commonly written, most explicit first:
+
+    1. **OME-XML** ``PhysicalSizeX/Y/Z`` (+ their units) in
+       ``ImageDescription`` — what stitched and OME-TIFF output usually
+       carries, and the only one of the three that can give ``z`` for a
+       sequence of single planes.
+    2. **ImageJ metadata** — ``spacing`` (+ ``unit``) for ``z``.
+    3. **Resolution tags** — ``XResolution``/``YResolution`` in pixels per
+       unit, the unit from ``ResolutionUnit`` (plain TIFFs) or from ImageJ's
+       ``unit`` string (ImageJ writes ``ResolutionUnit=NONE``).
+
+    Earlier sources win per axis, so a file carrying both keeps the explicit
+    OME value. Unrecognized units are ignored rather than guessed at.
 
     Parameters
     ----------
@@ -898,9 +1046,18 @@ def _tiff_pixel_size(path: str) -> PixelSize:
 
     pixel_size: PixelSize = {}
     with tifffile.TiffFile(path) as tif:
+        description = tif.pages[0].tags.get("ImageDescription")
+        xml = tif.ome_metadata or (
+            description.value
+            if description and "<OME" in str(description.value)
+            else None
+        )
+        if xml:
+            pixel_size.update(_ome_xml_pixel_size(str(xml)))
+
         ij = tif.imagej_metadata or {}
         spacing = ij.get("spacing")
-        if spacing:
+        if spacing and "z" not in pixel_size:
             factor = _UNIT_TO_UM.get(str(ij.get("unit", "micron")).lower(), 1.0)
             pixel_size["z"] = float(spacing) * factor
 
@@ -911,6 +1068,8 @@ def _tiff_pixel_size(path: str) -> PixelSize:
         ) or _UNIT_TO_UM.get(str(ij.get("unit", "")).lower())
         if um_per_unit is not None:
             for axis, tag_name in (("y", "YResolution"), ("x", "XResolution")):
+                if axis in pixel_size:
+                    continue
                 tag = page.tags.get(tag_name)
                 if tag and tag.value[0]:
                     num, den = tag.value
@@ -1014,6 +1173,18 @@ def _open_tiff_sequence(
         arr.shape,
         pixel_size,
     )
+    missing = [a for a in axes if a in _SPATIAL_AXES and a not in pixel_size]
+    if missing:
+        logger.warning(
+            "no voxel size for %s in %s -- checked OME-XML PhysicalSize*, "
+            "ImageJ spacing/unit, and the resolution tags. The store will be "
+            "uncalibrated on %s, which also leaves anything deriving physical "
+            "units from it (e.g. deconvolution voxel sizes) with nothing to "
+            "read. Pass pixel_size= to supply it.",
+            "/".join(missing),
+            Path(ts[0]).name,
+            "/".join(missing),
+        )
     return arr, axes, pixel_size
 
 
@@ -1361,17 +1532,18 @@ def to_ome_zarr(
 
     out = str(out_path)
     zarr.open_group(out, mode="w" if overwrite else "w-")
-    datasets = _write_pyramid(
-        arr,
-        axes,
-        out,
-        n_levels=n_levels,
-        downscale=downscale,
-        chunks=chunks,
-        base_scale=base_scale,
-        shard=shard,
-        progress=progress,
-    )
+    with _bounded_scheduler(arr):
+        datasets = _write_pyramid(
+            arr,
+            axes,
+            out,
+            n_levels=n_levels,
+            downscale=downscale,
+            chunks=chunks,
+            base_scale=base_scale,
+            shard=shard,
+            progress=progress,
+        )
     _write_multiscales(out, axes, datasets, Path(out).stem, calibrated=bool(ps))
     return out
 

{patchworks-2.1.1 → patchworks-2.1.2}/tests/test_ome_zarr.py

@@ -415,3 +415,164 @@ def test_glob_without_sequence_pattern_says_so(tmp_path):
     with pytest.raises(Exception) as exc:
         to_ome_zarr(str(tmp_path / "scan.ims"), tmp_path / "o2.zarr")
     assert "sequence_pattern" not in str(exc.value)
+
+
+def test_base_chunks_never_group_source_chunks(tmp_path):
+    """One output chunk must not require several source chunks.
+
+    A folder of stitched TIFFs gives one dask chunk per file, so a plane can
+    be gigabytes. The z cap of 16 would group sixteen of them into one output
+    chunk -- ~58 GB held to write 32 MB, against a 64 GB job. Splitting a
+    source chunk is fine; combining several is not.
+    """
+    from patchworks.plugins.ome_zarr import _default_chunks
+
+    shape = (4, 126, 45961, 42072)
+    # one chunk per file: (c=1, z=1, whole plane)
+    assert _default_chunks(
+        shape, "czyx", source_chunks=(1, 1, 45961, 42072)
+    ) == (
+        1,
+        1,
+        1024,
+        1024,
+    )
+    # a normally-chunked source is unaffected by the cap
+    assert _default_chunks(shape, "czyx", source_chunks=shape) == (
+        1,
+        16,
+        1024,
+        1024,
+    )
+    # and omitting it keeps the previous behaviour
+    assert _default_chunks(shape, "czyx") == (1, 16, 1024, 1024)
+
+
+def test_tiff_sequence_keeps_one_plane_per_chunk(tmp_path):
+    """End to end: a per-file source must not be z-grouped on write."""
+    tifffile = pytest.importorskip("tifffile")
+    n_z, n_c, size = 4, 2, 8
+    for z in range(n_z):
+        for c in range(n_c):
+            tifffile.imwrite(
+                tmp_path / f"s_Z{z:03d}_C{c}_V0.tif",
+                np.full((size, size), z * 10 + c, "uint16"),
+            )
+
+    out = tmp_path / "seq.zarr"
+    to_ome_zarr(
+        str(tmp_path / "*.tif"),
+        out,
+        sequence_pattern=r"_Z(?P<Z>\d+)_C(?P<C>\d+)_V\d+",
+        n_levels=1,
+        progress=False,
+    )
+    level0 = zarr.open_array(str(out) + "/0", mode="r")
+    assert level0.chunks[1] == 1, (
+        f"z was grouped into {level0.chunks[1]} planes per chunk; each source "
+        "chunk is a whole file, so that multiplies the read"
+    )
+    # and the data still round-trips
+    result = np.asarray(load_ome_zarr(out, channel=None))
+    assert result.shape == (n_c, n_z, size, size)
+    assert (
+        result[:, :, 0, 0]
+        == [[z * 10 + c for z in range(n_z)] for c in range(n_c)]
+    ).all()
+
+
+def _ome_xml(x, y, z, unit="um"):
+    return (
+        '<?xml version="1.0"?>'
+        '<OME xmlns="http://www.openmicroscopy.org/Schemas/OME/2016-06">'
+        f'<Image><Pixels PhysicalSizeX="{x}" PhysicalSizeXUnit="{unit}"'
+        f' PhysicalSizeY="{y}" PhysicalSizeYUnit="{unit}"'
+        f' PhysicalSizeZ="{z}" PhysicalSizeZUnit="{unit}"/></Image></OME>'
+    )
+
+
+def test_pixel_size_from_ome_xml(tmp_path):
+    """OME-XML PhysicalSize* must be read, not just the resolution tags.
+
+    Stitched output commonly writes OME-XML into ImageDescription and leaves
+    the TIFF tags at their defaults, so a tag-only reader calls a perfectly
+    calibrated image uncalibrated -- and it is the only source that can give
+    z for a sequence of single planes.
+    """
+    tifffile = pytest.importorskip("tifffile")
+    from patchworks.plugins.ome_zarr import _tiff_pixel_size
+
+    path = tmp_path / "ome.tif"
+    tifffile.imwrite(
+        path,
+        np.zeros((8, 8), "uint16"),
+        description=_ome_xml(0.325, 0.325, 1.5),
+    )
+    assert _tiff_pixel_size(str(path)) == pytest.approx(
+        {"z": 1.5, "y": 0.325, "x": 0.325}
+    )
+
+    # mm is converted, not taken at face value
+    mm = tmp_path / "mm.tif"
+    tifffile.imwrite(
+        mm,
+        np.zeros((8, 8), "uint16"),
+        description=_ome_xml(0.001, 0.001, 0.002, "mm"),
+    )
+    assert _tiff_pixel_size(str(mm)) == pytest.approx(
+        {"z": 2.0, "y": 1.0, "x": 1.0}
+    )
+
+
+def test_ome_xml_wins_over_the_resolution_tags(tmp_path):
+    """The explicit source wins per axis when a file carries both."""
+    tifffile = pytest.importorskip("tifffile")
+    from patchworks.plugins.ome_zarr import _tiff_pixel_size
+
+    path = tmp_path / "both.tif"
+    tifffile.imwrite(
+        path,
+        np.zeros((8, 8), "uint16"),
+        description=_ome_xml(0.325, 0.325, 1.5),
+        resolution=(20000.0, 20000.0),  # would say 0.5 µm
+        resolutionunit="CENTIMETER",
+    )
+    got = _tiff_pixel_size(str(path))
+    assert got["x"] == pytest.approx(0.325), "OME-XML must win over the tag"
+    assert got["z"] == pytest.approx(1.5)
+
+
+def test_uncalibrated_tiff_reports_nothing(tmp_path):
+    """No metadata gives an empty dict rather than an invented default."""
+    tifffile = pytest.importorskip("tifffile")
+    from patchworks.plugins.ome_zarr import _tiff_pixel_size
+
+    path = tmp_path / "bare.tif"
+    tifffile.imwrite(path, np.zeros((8, 8), "uint16"))
+    assert _tiff_pixel_size(str(path)) == {}
+
+
+def test_many_chunks_suggests_sharding(caplog):
+    """A coarse source forces fine chunks; fine chunks mean many files.
+
+    That trade-off is invisible until a shared filesystem is carrying a
+    million of them, so it has to be said at write time -- with the remedy.
+    """
+    import logging
+
+    from patchworks.plugins.ome_zarr import _warn_if_many_chunks
+
+    real = (4, 126, 45961, 42072)  # ~950k chunks at (1, 1, 1024, 1024)
+    with caplog.at_level(logging.WARNING, logger="patchworks.plugins.ome_zarr"):
+        _warn_if_many_chunks(real, (1, 1, 1024, 1024), shard=False)
+    assert "shard: true" in caplog.text
+    assert "952,560" in caplog.text
+
+    # Sharding already asked for, or a modest chunk count: nothing to say.
+    caplog.clear()
+    with caplog.at_level(logging.WARNING, logger="patchworks.plugins.ome_zarr"):
+        _warn_if_many_chunks(real, (1, 1, 1024, 1024), shard=True)
+        _warn_if_many_chunks(
+            (4, 126, 1024, 1024), (1, 16, 1024, 1024), shard=False
+        )
+    assert caplog.text == ""