From ae5d6475f0a45ddb72c2a5f30b563085aa73a079 Mon Sep 17 00:00:00 2001 From: "Tybulewicz, Tomasz" Date: Tue, 30 Jun 2026 10:47:41 +0200 Subject: [PATCH] fix: support for 1-channel and 16-bit images in Visualizer --- .../src/model_api/visualizer/visualizer.py | 26 +++++- .../tests/unit/visualizer/test_visualizer.py | 90 +++++++++++++++++++ 2 files changed, 113 insertions(+), 3 deletions(-) diff --git a/model_api/src/model_api/visualizer/visualizer.py b/model_api/src/model_api/visualizer/visualizer.py index f685a8b8e..8242298ae 100644 --- a/model_api/src/model_api/visualizer/visualizer.py +++ b/model_api/src/model_api/visualizer/visualizer.py @@ -51,6 +51,26 @@ def __init__(self, layout: Layout | None = None, auto_scale: bool = True) -> Non self.layout = layout self.auto_scale = auto_scale + @staticmethod + def _to_pil(image: np.ndarray) -> Image.Image: + """Convert numpy array to PIL Image, handling 16-bit and grayscale images. + + PIL doesn't support 16-bit RGB images, so we convert them to 8-bit. + Grayscale images are converted to RGB for compatibility with overlays. + + Args: + image: Input numpy array (uint8 or uint16, grayscale or RGB). + + Returns: + PIL Image in 8-bit RGB format. + """ + if image.dtype == np.uint16: + image = (image / 256).astype(np.uint8) + pil_image = Image.fromarray(image) + if pil_image.mode != "RGB": + pil_image = pil_image.convert("RGB") + return pil_image + @staticmethod def compute_scale_factor(image: Image.Image) -> float: """Compute a scale factor based on the image's longer edge relative to 720p (1280px). @@ -68,13 +88,13 @@ def compute_scale_factor(image: Image.Image) -> float: def show(self, image: Image.Image | np.ndarray, result: Result) -> None: if isinstance(image, np.ndarray): - image = Image.fromarray(image) + image = self._to_pil(image) scene = self._scene_from_result(image, result) return scene.show() def save(self, image: Image.Image | np.ndarray, result: Result, path: Path) -> None: if isinstance(image, np.ndarray): - image = Image.fromarray(image) + image = self._to_pil(image) scene = self._scene_from_result(image, result) scene.save(path) @@ -82,7 +102,7 @@ def render(self, image: Image.Image | np.ndarray, result: Result) -> Image.Image is_numpy = isinstance(image, np.ndarray) if is_numpy: - image = Image.fromarray(image) + image = self._to_pil(image) scene = self._scene_from_result(image, result) result_img: Image = scene.render() diff --git a/model_api/tests/unit/visualizer/test_visualizer.py b/model_api/tests/unit/visualizer/test_visualizer.py index d4961f2a7..09ae71025 100644 --- a/model_api/tests/unit/visualizer/test_visualizer.py +++ b/model_api/tests/unit/visualizer/test_visualizer.py @@ -97,3 +97,93 @@ def test_keypoint_scene_creation(mock_image: Image, tmpdir: Path): visualizer.save(mock_image, keypoint_result, tmpdir / "keypoint_scene.jpg") assert Path(tmpdir / "keypoint_scene.jpg").exists() + + +def test_render_16bit_image(): + """Test Visualizer.render() handles 16-bit images. + + 16-bit images are common in medical imaging and from cameras with + high dynamic range. OpenCV's cv2.imread with IMREAD_UNCHANGED flag + returns uint16 arrays for 16-bit images. + """ + # Simulate a 16-bit BGR image as returned by cv2.imread(..., cv2.IMREAD_UNCHANGED) + image_16bit = np.zeros((100, 100, 3), dtype=np.uint16) + image_16bit[25:75, 25:75] = 65535 # Max value for uint16 + + anomaly_result = AnomalyResult( + anomaly_map=np.ones((100, 100), dtype=np.uint8) * 255, + pred_boxes=np.array([[25, 25, 75, 75]]), + pred_label="Anomaly", + pred_mask=None, + pred_score=0.9, + ) + + visualizer = Visualizer() + rendered = visualizer.render(image_16bit, anomaly_result) + + assert isinstance(rendered, np.ndarray) + assert rendered.shape == image_16bit.shape[:2] + (3,) + # Output should be 8-bit for display purposes + assert rendered.dtype == np.uint8 + + +def test_show_16bit_image(monkeypatch): + """Test Visualizer.show() handles 16-bit images.""" + image_16bit = np.zeros((100, 100, 3), dtype=np.uint16) + image_16bit[25:75, 25:75] = 65535 + + anomaly_result = AnomalyResult( + anomaly_map=np.ones((100, 100), dtype=np.uint8) * 255, + pred_boxes=None, + pred_label="Anomaly", + pred_mask=None, + pred_score=0.9, + ) + + shown = [] + monkeypatch.setattr(Image.Image, "show", lambda self: shown.append(True)) + + visualizer = Visualizer() + visualizer.show(image_16bit, anomaly_result) + + assert len(shown) == 1 + + +def test_save_16bit_image(tmpdir: Path): + """Test Visualizer.save() handles 16-bit images.""" + image_16bit = np.zeros((100, 100, 3), dtype=np.uint16) + image_16bit[25:75, 25:75] = 65535 + + anomaly_result = AnomalyResult( + anomaly_map=np.ones((100, 100), dtype=np.uint8) * 255, + pred_boxes=None, + pred_label="Anomaly", + pred_mask=None, + pred_score=0.9, + ) + + visualizer = Visualizer() + output_path = tmpdir / "16bit_output.jpg" + visualizer.save(image_16bit, anomaly_result, output_path) + + assert Path(output_path).exists() + + +def test_render_grayscale_image(): + """Test Visualizer.render() handles 8-bit grayscale (1 channel) images.""" + image_gray = np.zeros((100, 100), dtype=np.uint8) + image_gray[25:75, 25:75] = 255 + + anomaly_result = AnomalyResult( + anomaly_map=np.ones((100, 100), dtype=np.uint8) * 255, + pred_boxes=np.array([[25, 25, 75, 75]]), + pred_label="Anomaly", + pred_mask=None, + pred_score=0.9, + ) + + visualizer = Visualizer() + rendered = visualizer.render(image_gray, anomaly_result) + + assert isinstance(rendered, np.ndarray) + assert rendered.dtype == np.uint8