Memory layout optimized TIFF reader (#120)

* Add experimental tiff reader

* Add multipage image support for experimental reader

* Dynamically calculate tile scale

* Add unit tests for experimental tiff reader

* Extract logic for finding memory layout of chucks

* Refactor experimental reader

* Update tests

* Address PR comments

* Added missing import

* Remove erroneous clear function call on metadata

* REplace the optimal reader implementation with a caching approach for tiles that are read out of order

* Limit the number of pending tiles for writing to bound memory usage

* Add support for 5D images

* lint fixes

* Increase CI timeout

* Exclude experimental reader tests from Windows CI

* Add support for depth ordered images

* Expose experimental reader parameters

* Lint fix

Author: XanthosXanthopoulos

.github/workflows/ci.yml

@@ -6,7 +6,7 @@ jobs:
   build:
     name: ${{ matrix.sys.os }}
     runs-on: ${{ matrix.sys.os }}
-    timeout-minutes: 20
+    timeout-minutes: 25
     strategy:
       matrix:
         sys:
@@ -88,7 +88,7 @@ jobs:
         AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
       run: |
         mkdir test_temp
-        micromamba run -n test pytest --basetemp=test_temp -v --cov=tiledb --cov-report=term-missing --durations=0 tests/
+        micromamba run -n test pytest --basetemp=test_temp -v --cov=tiledb --cov-report=term-missing --durations=0 tests/ --ignore=tests/integration/converters/test_ome_tiff_experimental.py
       if: ${{ matrix.sys.os == 'windows-latest' }}
 
     - name: Run notebook examples

.pre-commit-config.yaml

@@ -1,13 +1,13 @@
 repos:
   - repo: https://github.com/ambv/black
-    rev: 23.1.0
+    rev: 24.4.2
     hooks:
     - id: black
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.0.241
+    rev: v0.4.2
     hooks:
     - id: ruff
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.991
+    rev: v1.10.0
     hooks:
       - id: mypy

pyproject.toml

@@ -9,7 +9,7 @@ module = "tests.*"
 ignore_errors = true
 
 [tool.ruff]
-ignore = ["E501"]
-extend-select = ["I001"]
+lint.ignore = ["E501"]
+lint.extend-select = ["I001"]
 exclude = ["__init__.py"]
 fix = true

setup.cfg

@@ -33,7 +33,7 @@ classifiers =
 [options]
 zip_safe = False
 packages = find_namespace:
-python_requires = >=3.7
+python_requires = >=3.8
 
 
 [options.packages.find]

tests/integration/converters/test_ome_tiff_experimental.py

@@ -0,0 +1,174 @@
+import numpy as np
+import pytest
+import tifffile
+
+import tiledb
+from tests import assert_image_similarity, get_path
+from tiledb.bioimg.converters.ome_tiff import OMETiffConverter
+from tiledb.bioimg.helpers import open_bioimg
+from tiledb.bioimg.openslide import TileDBOpenSlide
+from tiledb.cc import WebpInputFormat
+
+
+# We need to expand on the test files. Most of the test files we have currently are not memory
+# contiguous and the ones that are not RGB files to test the different compressors
+@pytest.mark.parametrize("filename,num_series", [("UTM2GTIF.tiff", 1)])
+@pytest.mark.parametrize("preserve_axes", [False, True])
+@pytest.mark.parametrize("chunked,max_workers", [(False, 0), (True, 0), (True, 4)])
+@pytest.mark.parametrize(
+    "compressor",
+    [tiledb.ZstdFilter(level=0)],
+)
+def test_ome_tiff_converter_exclude_original_metadata(
+    tmp_path, filename, num_series, preserve_axes, chunked, max_workers, compressor
+):
+    # The image is not RGB to use the WbeP compressor
+    if isinstance(compressor, tiledb.WebpFilter) and filename == "UTM2GTIF.tiff":
+        pytest.skip(f"WebPFilter cannot be applied to {filename}")
+
+    input_path = get_path(filename)
+    tiledb_path = tmp_path / "to_tiledb"
+    exprimental_path = tmp_path / "experimental"
+    OMETiffConverter.to_tiledb(
+        input_path,
+        str(tiledb_path),
+        preserve_axes=preserve_axes,
+        chunked=chunked,
+        max_workers=max_workers,
+        compressor=compressor,
+        log=False,
+        exclude_metadata=True,
+    )
+
+    OMETiffConverter.to_tiledb(
+        input_path,
+        str(exprimental_path),
+        preserve_axes=preserve_axes,
+        chunked=True,
+        max_workers=max_workers,
+        compressor=compressor,
+        log=False,
+        exclude_metadata=True,
+        experimental_reader=True,
+    )
+
+    with TileDBOpenSlide(str(tiledb_path)) as t:
+        with TileDBOpenSlide(str(exprimental_path)) as e:
+            assert t.level_count == e.level_count
+
+            for level in range(t.level_count):
+                np.testing.assert_array_equal(t.read_level(level), e.read_level(level))
+
+
+@pytest.mark.parametrize("filename,num_series", [("UTM2GTIF.tiff", 1)])
+@pytest.mark.parametrize("preserve_axes", [False, True])
+@pytest.mark.parametrize("chunked,max_workers", [(True, 0), (True, 4)])
+@pytest.mark.parametrize(
+    "compressor",
+    [tiledb.ZstdFilter(level=0)],
+)
+def test_ome_tiff_converter_roundtrip(
+    tmp_path, filename, num_series, preserve_axes, chunked, max_workers, compressor
+):
+    if isinstance(compressor, tiledb.WebpFilter) and filename == "UTM2GTIF.tiff":
+        pytest.skip(f"WebPFilter cannot be applied to {filename}")
+
+    input_path = get_path(filename)
+    tiledb_path = tmp_path / "to_tiledb"
+    output_path = tmp_path / "from_tiledb"
+    OMETiffConverter.to_tiledb(
+        input_path,
+        str(tiledb_path),
+        preserve_axes=preserve_axes,
+        chunked=chunked,
+        max_workers=max_workers,
+        compressor=compressor,
+        log=False,
+        experimental_reader=True,
+        reader_kwargs=dict(
+            extra_tags=(
+                "ModelPixelScaleTag",
+                "ModelTiepointTag",
+                "GeoKeyDirectoryTag",
+                "GeoAsciiParamsTag",
+            )
+        ),
+    )
+    # Store it back to NGFF Zarr
+    OMETiffConverter.from_tiledb(str(tiledb_path), str(output_path))
+
+    with tifffile.TiffFile(input_path) as t1, tifffile.TiffFile(output_path) as t2:
+        compare_tiff(t1, t2, lossless=False)
+
+
+@pytest.mark.parametrize(
+    "filename,dims",
+    [
+        ("single-channel.ome.tif", "YX"),
+        ("z-series.ome.tif", "ZYX"),
+        ("multi-channel.ome.tif", "CYX"),
+        ("time-series.ome.tif", "TYX"),
+        ("multi-channel-z-series.ome.tif", "CZYX"),
+        ("multi-channel-time-series.ome.tif", "TCYX"),
+        ("4D-series.ome.tif", "TZYX"),
+        ("multi-channel-4D-series.ome.tif", "TCZYX"),
+    ],
+)
+@pytest.mark.parametrize("tiles", [{}, {"X": 128, "Y": 128, "Z": 2, "C": 1, "T": 3}])
+def test_ome_tiff_converter_artificial_rountrip(tmp_path, filename, dims, tiles):
+    input_path = get_path(f"artificial-ome-tiff/{filename}")
+    tiledb_path = tmp_path / "to_tiledb"
+    experimental_path = tmp_path / "_experimental"
+    output_path = tmp_path / "from_tiledb"
+
+    OMETiffConverter.to_tiledb(input_path, str(tiledb_path), tiles=tiles)
+    OMETiffConverter.to_tiledb(
+        input_path,
+        str(experimental_path),
+        tiles=tiles,
+        experimental_reader=True,
+        chunked=True,
+        max_workers=16,
+    )
+
+    with TileDBOpenSlide(str(experimental_path)) as t:
+        assert len(tiledb.Group(str(experimental_path))) == t.level_count == 1
+
+    with open_bioimg(str(experimental_path / "l_0.tdb")) as A:
+        assert "".join(dim.name for dim in A.domain) == dims
+        assert A.dtype == np.int8
+        assert A.dim("X").tile == tiles.get("X", 439)
+        assert A.dim("Y").tile == tiles.get("Y", 167)
+        if A.domain.has_dim("Z"):
+            assert A.dim("Z").tile == tiles.get("Z", 1)
+        if A.domain.has_dim("C"):
+            assert A.dim("C").tile == tiles.get("C", 3)
+        if A.domain.has_dim("T"):
+            assert A.dim("T").tile == tiles.get("T", 1)
+
+    OMETiffConverter.from_tiledb(str(tiledb_path), str(output_path))
+    with tifffile.TiffFile(input_path) as t1, tifffile.TiffFile(output_path) as t2:
+        compare_tiff(t1, t2, lossless=True)
+
+
+def compare_tiff(t1: tifffile.TiffFile, t2: tifffile.TiffFile, lossless: bool = True):
+    assert len(t1.series[0].levels) == len(t2.series[0].levels)
+
+    for l1, l2 in zip(t1.series[0].levels, t2.series[0].levels):
+        assert l1.axes.replace("S", "C") == l2.axes.replace("S", "C")
+        assert l1.shape == l2.shape
+        assert l1.dtype == l2.dtype
+        assert l1.nbytes == l2.nbytes
+
+        if lossless:
+            np.testing.assert_array_equal(l1.asarray(), l2.asarray())
+        else:
+            assert_image_similarity(l1.asarray(), l2.asarray(), channel_axis=0)
+
+
+compressors = [
+    None,
+    tiledb.ZstdFilter(level=0),
+    tiledb.WebpFilter(WebpInputFormat.WEBP_RGB, lossless=False),
+    tiledb.WebpFilter(WebpInputFormat.WEBP_RGB, lossless=True),
+]

tests/unit/test_tiles.py

@@ -11,10 +11,18 @@
 
 def test_dim_range():
     dims = list(domain)
-    assert dim_range(dims[0]) == range(0, 10, 3)
-    assert dim_range(dims[1]) == range(0, 15, 5)
-    assert dim_range(dims[2]) == range(0, 7, 4)
-    assert dim_range(dims[3]) == range(0, 3, 3)
+    assert dim_range(
+        (int(dims[0].domain[0]), int(dims[0].domain[1]), int(dims[0].tile))
+    ) == range(0, 10, 3)
+    assert dim_range(
+        (int(dims[1].domain[0]), int(dims[1].domain[1]), int(dims[1].tile))
+    ) == range(0, 15, 5)
+    assert dim_range(
+        (int(dims[2].domain[0]), int(dims[2].domain[1]), int(dims[2].tile))
+    ) == range(0, 7, 4)
+    assert dim_range(
+        (int(dims[3].domain[0]), int(dims[3].domain[1]), int(dims[3].tile))
+    ) == range(0, 3, 3)
 
 
 def test_iter_slices():