fix: handle incomplete DICOM pyramid when getting thumbnail

src/aignostics/wsi/_gui.py

@@ -12,6 +12,7 @@
 
 from loguru import logger
 
+from ._openslide_handler import DEFAULT_MAX_SAFE_DIMENSION
 from ._service import Service
 
 
@@ -23,11 +24,13 @@ def register_pages() -> None:
         app.add_static_files("/wsi_assets", Path(__file__).parent / "assets")
 
         @app.get("/thumbnail")
-        def thumbnail(source: str) -> Response:
+        def thumbnail(source: str, max_safe_dimension: int = DEFAULT_MAX_SAFE_DIMENSION) -> Response:
             """Serve a thumbnail for a given source reference.
 
             Args:
                 source (str): The source of the slide pointing to a file on the filesystem.
+                max_safe_dimension (int): Maximum dimension (width or height) of smallest pyramid level
+                    before considering the pyramid incomplete.
 
             Returns:
                 fastapi.Response: HTTP response containing the thumbnail or fallback image.
@@ -36,7 +39,10 @@ def thumbnail(source: str) -> Response:
             from fastapi.responses import RedirectResponse  # noqa: PLC0415
 
             try:
-                return Response(content=Service().get_thumbnail_bytes(Path(source)), media_type="image/png")
+                return Response(
+                    content=Service().get_thumbnail_bytes(Path(source), max_safe_dimension=max_safe_dimension),
+                    media_type="image/png",
+                )
             except ValueError:
                 logger.warning("Error generating thumbnail on bad request or invalid image input")
                 return RedirectResponse("/wsi_assets/fallback.png")

src/aignostics/wsi/_openslide_handler.py

@@ -10,6 +10,7 @@
 from PIL.Image import Image
 
 TIFF_IMAGE_DESCRIPTION = "tiff.ImageDescription"
+DEFAULT_MAX_SAFE_DIMENSION = 4096  # Maximum safe pyramid level dimension (pixels)
 
 
 class OpenSlideHandler:
@@ -62,15 +63,34 @@ def _detect_format(self) -> str | None:
 
         return base_format
 
-    def get_thumbnail(self) -> Image:
+    def get_thumbnail(self, max_safe_dimension: int = DEFAULT_MAX_SAFE_DIMENSION) -> Image:
         """Get thumbnail of the slide.
 
+        Args:
+            max_safe_dimension (int): Maximum dimension (width or height) of smallest pyramid level
+                before considering the pyramid incomplete.
+
         Returns:
             Image: Thumbnail image of the slide.
 
         Raises:
-            OpenSlideError: If the slide cannot be opened or if thumbnail generation fails.
+            RuntimeError: If the slide has an incomplete pyramid and thumbnail generation
+                would require excessive memory.
         """
+        # Detect incomplete pyramid by checking smallest level
+        smallest_level_idx = self.slide.level_count - 1
+        smallest_width, smallest_height = self.slide.level_dimensions[smallest_level_idx]
+
+        if max(smallest_width, smallest_height) > max_safe_dimension:
+            msg = (
+                f"Cannot generate thumbnail: incomplete pyramid detected. "
+                f"Smallest available level (Level {smallest_level_idx}) is "
+                f"{smallest_width}x{smallest_height} pixels, which exceeds safe "
+                f"threshold of {max_safe_dimension}x{max_safe_dimension}. "
+                f"This file appears to be missing lower-resolution pyramid levels."
+            )
+            raise RuntimeError(msg)
+
         return self.slide.get_thumbnail((256, 256))
 
     def _parse_xml_image_description(self, xml_string: str) -> dict[str, Any]:  # noqa: C901, PLR6301

src/aignostics/wsi/_service.py

@@ -10,6 +10,8 @@
 from aignostics import WSI_SUPPORTED_FILE_EXTENSIONS
 from aignostics.utils import BaseService, Health
 
+from ._openslide_handler import DEFAULT_MAX_SAFE_DIMENSION
+
 TIMEOUT = 60  # 1 minutes
 
 
@@ -38,11 +40,13 @@ def health(self) -> Health:  # noqa: PLR6301
         )
 
     @staticmethod
-    def get_thumbnail(path: Path) -> "PIL.Image.Image":  # type: ignore # noqa: F821
+    def get_thumbnail(path: Path, max_safe_dimension: int = DEFAULT_MAX_SAFE_DIMENSION) -> "PIL.Image.Image":  # type: ignore # noqa: F821
         """Get thumbnail as PIL image.
 
         Args:
             path (Path): Path to the image.
+            max_safe_dimension (int): Maximum dimension (width or height) of smallest pyramid level
+                before considering the pyramid incomplete.
 
         Returns:
             PIL.Image.Image: Thumbnail of the image.
@@ -62,18 +66,20 @@ def get_thumbnail(path: Path) -> "PIL.Image.Image":  # type: ignore # noqa: F821
             logger.warning(message)
             raise ValueError(message)
         try:
-            return OpenSlideHandler.from_file(path).get_thumbnail()
+            return OpenSlideHandler.from_file(path).get_thumbnail(max_safe_dimension=max_safe_dimension)
         except Exception as e:
             message = f"Error processing file {path}: {e!s}"
             logger.exception(message)
             raise RuntimeError(message) from e
 
     @staticmethod
-    def get_thumbnail_bytes(path: Path) -> bytes:
+    def get_thumbnail_bytes(path: Path, max_safe_dimension: int = DEFAULT_MAX_SAFE_DIMENSION) -> bytes:
         """Get thumbnail of a image as bytes.
 
         Args:
             path (Path): Path to the image.
+            max_safe_dimension (int): Maximum dimension (width or height) of smallest pyramid level
+                before considering the pyramid incomplete.
 
         Returns:
             bytes: Thumbnail of the image.
@@ -82,7 +88,7 @@ def get_thumbnail_bytes(path: Path) -> bytes:
             ValueError: If the file type is not supported.
             RuntimeError: If there is an error processing the file.
         """
-        thumbnail_image = Service.get_thumbnail(path)
+        thumbnail_image = Service.get_thumbnail(path, max_safe_dimension=max_safe_dimension)
         buffer = io.BytesIO()
         thumbnail_image.save(buffer, format="PNG")
         return buffer.getvalue()

tests/aignostics/application/gui_test.py

@@ -29,6 +29,7 @@
 
 
 @pytest.mark.e2e
+@pytest.mark.flaky(retries=2, delay=5, only_on=[AssertionError])
 @pytest.mark.timeout(timeout=30)
 async def test_gui_index(user: User, silent_logging, record_property) -> None:
     """Test that the user sees the index page, and sees the intro."""

tests/aignostics/wsi/service_test.py

@@ -109,6 +109,32 @@ def test_serve_thumbnail_for_tiff(user: User, record_property) -> None:
     assert image.height > 0
 
 
+@pytest.mark.integration
+def test_serve_thumbnail_fails_on_incomplete_pyramid(user: User, silent_logging, record_property) -> None:
+    """Test that thumbnail generation fails gracefully for DICOM with incomplete pyramid.
+
+    The small-pyramidal.dcm test file has only 1 pyramid level at 2054x1529 pixels,
+    with no smaller resolution levels available. By setting max_safe_dimension=1024
+    via query parameter, we simulate the condition where the smallest available pyramid
+    level is too large for safe thumbnail generation, which would normally cause OOM errors.
+    """
+    record_property("tested-item-id", "SPEC-WSI-SERVICE")
+
+    client = TestClient(app)
+
+    test_dir = Path(__file__).parent
+    resources_dir = test_dir.parent.parent / "resources" / "run"
+    test_file_path = resources_dir / "small-pyramidal.dcm"
+
+    # Use low max_safe_dimension (1024) to trigger incomplete pyramid detection
+    # The file has dimensions 2054x1529, which exceeds the threshold
+    response = client.get(f"/thumbnail?source={test_file_path.absolute()}&max_safe_dimension=1024")
+
+    # Should return 200 with fallback image (not crash with 500 or OOM)
+    assert response.status_code == 200
+    assert int(response.headers["Content-Length"]) == CONTENT_LENGTH_FALLBACK
+
+
 @pytest.mark.integration
 @pytest.mark.timeout(timeout=60)
 def test_serve_tiff_to_jpeg_fails_on_broken_url(user: User, record_property) -> None: