fix: support for 1-channel and 16-bit images in Visualizer

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,21 +88,21 @@ 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)
 
     def render(self, image: Image.Image | np.ndarray, result: Result) -> Image.Image | np.ndarray:
         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()

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