fix backend setting (#226)

* fix backend setting

* format with ruff

* remove zarr KVStore monkey patch

Tiffslide takes care of the monkeypatching for us.
Including it here caused errors in our tests.

* attempt to fix zarr.storage issue

* expect 601 slides in linux and macos

* expect 601 patches in docker output

* install numpy, openslide, tiffslide after torch

.github/workflows/ci.yml

@@ -41,7 +41,8 @@ jobs:
           sudo apt update
           sudo apt install -y libopenslide0
           python -m pip install --upgrade pip setuptools wheel
-          python -m pip install --pre torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu openslide-python tiffslide
+          python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
+          python -m pip install openslide-python tiffslide
           python -m pip install --editable .[dev]
       - name: Run tests
         run: python -m pytest --verbose tests/
@@ -67,7 +68,7 @@ jobs:
           test -f results/run_metadata_*.json
           test -f results/patches/JP2K-33003-1.h5
           test -f results/model-outputs-csv/JP2K-33003-1.csv
-          test $(wc -l < results/model-outputs-csv/JP2K-33003-1.csv) -eq 675
+          test $(wc -l < results/model-outputs-csv/JP2K-33003-1.csv) -eq 601
 
   test-package:
     strategy:
@@ -89,7 +90,8 @@ jobs:
       - name: Install the wsinfer python package
         run: |
           python -m pip install --upgrade pip setuptools wheel
-          python -m pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu openslide-python tiffslide
+          python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+          python -m pip install numpy openslide-python tiffslide
           python -m pip install .
       - name: Run 'wsinfer run' on Unix
         if: matrix.os != 'windows-latest'
@@ -102,7 +104,7 @@ jobs:
           test -f results/run_metadata_*.json
           test -f results/patches/JP2K-33003-1.h5
           test -f results/model-outputs-csv/JP2K-33003-1.csv
-          test $(wc -l < results/model-outputs-csv/JP2K-33003-1.csv) -eq 675
+          test $(wc -l < results/model-outputs-csv/JP2K-33003-1.csv) -eq 601
       # FIXME: tissue segmentation has different outputs on Windows. The patch sizes
       # are the same but the coordinates found are different.
       - name: Run 'wsinfer run' on Windows

wsinfer/__init__.py

@@ -6,19 +6,3 @@
     from ._version import __version__
 except ImportError:
     __version__ = "0.0.unknown"
-
-
-# Patch Zarr. See:
-# https://github.com/bayer-science-for-a-better-life/tiffslide/issues/72#issuecomment-1627918238
-# https://github.com/zarr-developers/zarr-python/pull/1454
-def _patch_zarr_kvstore() -> None:
-    from zarr.storage import KVStore
-
-    def _zarr_KVStore___contains__(self, key):  # type: ignore
-        return key in self._mutable_mapping
-
-    if "__contains__" not in KVStore.__dict__:
-        KVStore.__contains__ = _zarr_KVStore___contains__
-
-
-_patch_zarr_kvstore()

wsinfer/cli/convert_csv_to_sbubmi.py

@@ -36,8 +36,8 @@
 import pandas as pd
 import tqdm
 
-from ..wsi import WSI
 from ..wsi import CanReadRegion
+from ..wsi import get_wsi_cls
 
 
 def _box_to_polygon(
@@ -353,7 +353,8 @@ def tosbu(
             click.secho(f"WSI file not found: {wsi_file}", bg="red")
             click.secho("Skipping...", bg="red")
             continue
-        slide = WSI(wsi_file)
+        wsi_reader = get_wsi_cls()
+        slide = wsi_reader(wsi_file)
 
         slide_width, slide_height = slide.level_dimensions[0]
 

wsinfer/modellib/data.py

@@ -10,7 +10,7 @@
 import torch
 from PIL import Image
 
-from wsinfer.wsi import WSI
+from wsinfer.wsi import get_wsi_cls
 
 
 def _read_patch_coords(path: str | Path) -> npt.NDArray[np.int_]:
@@ -87,7 +87,8 @@ def __init__(
 
     def worker_init(self, worker_id: int | None = None) -> None:
         del worker_id
-        self.slide = WSI(self.wsi_path)
+        wsi_reader = get_wsi_cls()
+        self.slide = wsi_reader(self.wsi_path)
 
     def __len__(self) -> int:
         return self.patches.shape[0]

wsinfer/patchlib/__init__.py

@@ -10,9 +10,9 @@
 import numpy.typing as npt
 from PIL import Image
 
-from ..wsi import WSI
 from ..wsi import _validate_wsi_directory
 from ..wsi import get_avg_mpp
+from ..wsi import get_wsi_cls
 from .patch import get_multipolygon_from_binary_arr
 from .patch import get_patch_coordinates_within_polygon
 from .segment import segment_tissue
@@ -103,7 +103,7 @@ def segment_and_patch_one_slide(
         logger.info(f"mask_path={mask_path}")
         return None
 
-    slide = WSI(slide_path)
+    slide = get_wsi_cls()(slide_path)
     mpp = get_avg_mpp(slide_path)
     logger.info(f"Slide has WxH {slide.dimensions} and MPP={mpp}")
 

wsinfer/wsi.py

@@ -3,9 +3,7 @@
 import logging
 from fractions import Fraction
 from pathlib import Path
-from typing import Literal
 from typing import Protocol
-from typing import overload
 
 import tifffile
 from PIL import Image
@@ -17,6 +15,9 @@
 
 logger = logging.getLogger(__name__)
 
+_BACKEND: str = "tiffslide"
+
+_allowed_backends = {"openslide", "tiffslide"}
 
 try:
     import openslide
@@ -46,48 +47,46 @@
     )
 
 
-@overload
-def set_backend(name: Literal["openslide"]) -> type[openslide.OpenSlide]:
-    ...
-
-
-@overload
-def set_backend(name: Literal["tiffslide"]) -> type[tiffslide.TiffSlide]:
-    ...
-
-
-def set_backend(
-    name: Literal["openslide"] | Literal["tiffslide"],
-) -> type[tiffslide.TiffSlide] | type[openslide.OpenSlide]:
-    global WSI
-    if name not in ["openslide", "tiffslide"]:
-        raise ValueError(f"Unknown backend: {name}")
-    logger.info(f"Setting backend to {name}")
-    if name == "openslide":
-        if not HAS_OPENSLIDE:
-            raise BackendNotAvailable(
-                "OpenSlide is not available. Please install the OpenSlide compiled"
-                " library and the Python package 'openslide-python'."
-                " See https://openslide.org/ for more information."
-            )
-        WSI = openslide.OpenSlide
+def set_backend(name: str) -> None:
+    global _BACKEND
+    if name not in _allowed_backends:
+        raise ValueError(f"Unknown backend: '{name}'")
+    if name == "openslide" and not HAS_OPENSLIDE:
+        raise BackendNotAvailable(
+            "OpenSlide is not available. Please install the OpenSlide compiled"
+            " library and the Python package 'openslide-python'."
+            " See https://openslide.org/ for more information."
+        )
     elif name == "tiffslide":
         if not HAS_TIFFSLIDE:
             raise BackendNotAvailable(
                 "TiffSlide is not available. Please install 'tiffslide'."
             )
-        WSI = tiffslide.TiffSlide
+
+    logger.debug(f"Set backend to {name}")
+
+    _BACKEND = name
+
+
+def get_wsi_cls() -> type[openslide.OpenSlide] | type[tiffslide.TiffSlide]:
+    if _BACKEND not in _allowed_backends:
+        raise ValueError(
+            f"Unknown backend: '{_BACKEND}'. Please contact the developer!"
+        )
+    if _BACKEND == "openslide":
+        return openslide.OpenSlide  # type: ignore
+    elif _BACKEND == "tiffslide":
+        return tiffslide.TiffSlide
     else:
-        raise ValueError(f"Unknown backend: {name}")
-    return WSI
+        raise ValueError("Contact the developer, slide backend not known")
 
 
 # Set the slide backend based on the environment.
-WSI: type[openslide.OpenSlide] | type[tiffslide.TiffSlide]
-if HAS_OPENSLIDE:
-    WSI = set_backend("openslide")
-elif HAS_TIFFSLIDE:
-    WSI = set_backend("tiffslide")
+# Prioritize TiffSlide if the user has it installed.
+if HAS_TIFFSLIDE:
+    set_backend("tiffslide")
+elif HAS_OPENSLIDE:
+    set_backend("openslide")
 else:
     raise NoBackendException("No backend found! Please install openslide or tiffslide")