diff --git a/src/aignostics/wsi/_gui.py b/src/aignostics/wsi/_gui.py index 154b4ac7b..b1628495b 100644 --- a/src/aignostics/wsi/_gui.py +++ b/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") diff --git a/src/aignostics/wsi/_openslide_handler.py b/src/aignostics/wsi/_openslide_handler.py index 46c74259a..f118f6c96 100644 --- a/src/aignostics/wsi/_openslide_handler.py +++ b/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 diff --git a/src/aignostics/wsi/_service.py b/src/aignostics/wsi/_service.py index ca6f7166a..1e77f37fd 100644 --- a/src/aignostics/wsi/_service.py +++ b/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() diff --git a/tests/aignostics/application/gui_test.py b/tests/aignostics/application/gui_test.py index 48e885ad4..ae9e82205 100644 --- a/tests/aignostics/application/gui_test.py +++ b/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.""" diff --git a/tests/aignostics/wsi/service_test.py b/tests/aignostics/wsi/service_test.py index b79213ca8..76f9177b5 100644 --- a/tests/aignostics/wsi/service_test.py +++ b/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: