FIX: AddImageVideoConverter handles grayscale and 16-bit images, keeps fractional FPS, and raises for unreadable input or an unopenable writer

pyrit/converter/add_image_to_video_converter.py

@@ -89,7 +89,9 @@ async def _add_image_to_video_async(self, image_path: str) -> bytes:
 
         Raises:
             ModuleNotFoundError: If OpenCV is not installed.
-            ValueError: If the image path is invalid or unsupported video format.
+            ValueError: If the image path is invalid, the input video format is unsupported, the input
+                video cannot be read, the output video cannot be written, or the overlay image cannot
+                be decoded or has an unsupported bit depth.
         """
         try:
             import cv2  # noqa: F401
@@ -132,8 +134,9 @@ def _add_image_to_video_sync(
             bytes: The converted video data.
 
         Raises:
-            ValueError: If the input video format is unsupported or the overlay image cannot
-                be decoded.
+            ValueError: If the input video format is unsupported, the input video cannot be read,
+                the output video cannot be written, or the overlay image cannot be decoded or has an
+                unsupported bit depth.
         """
         import cv2
 
@@ -153,26 +156,44 @@ def _add_image_to_video_sync(
         try:
             try:
                 cap = cv2.VideoCapture(str(input_path))
+                if not cap.isOpened():
+                    raise ValueError("Could not read the input video")
 
-                # Get video properties
-                fps = int(cap.get(cv2.CAP_PROP_FPS))
+                # Get video properties; the frame rate can be fractional (29.97, 23.976)
+                fps = cap.get(cv2.CAP_PROP_FPS)
                 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                 height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-                video_char_code = cv2.VideoWriter.fourcc(*video_encoding_map[file_extension])
-                output_video = cv2.VideoWriter(str(output_path), video_char_code, fps, (width, height))
 
                 # Load and resize the overlay image
 
                 image_np_arr = np.frombuffer(image_bytes, np.uint8)
                 decoded = cv2.imdecode(image_np_arr, cv2.IMREAD_UNCHANGED)
                 if decoded is None:
                     raise ValueError("Failed to decode overlay image")
+                # Frames are 8-bit BGR: scale 16-bit images down, reject other depths (a float32 TIFF
+                # would be silently truncated), and give a grayscale image (no channel axis, or a
+                # single channel) three channels
+                if decoded.dtype == np.uint16:
+                    decoded = cv2.convertScaleAbs(decoded, alpha=255 / 65535)
+                elif decoded.dtype != np.uint8:
+                    raise ValueError(f"Unsupported overlay image depth: {decoded.dtype}")
+                if decoded.ndim == 2 or decoded.shape[2] == 1:
+                    decoded = cv2.cvtColor(decoded, cv2.COLOR_GRAY2BGR)
                 overlay = cv2.resize(decoded, self._img_resize_size)
 
+                video_char_code = cv2.VideoWriter.fourcc(*video_encoding_map[file_extension])
+                output_video = cv2.VideoWriter(str(output_path), video_char_code, fps, (width, height))
+                if not output_video.isOpened():
+                    raise ValueError(
+                        f"Could not open a video writer for '{file_extension}' "
+                        f"(codec {video_encoding_map[file_extension]}) at {fps} FPS, {width}x{height}"
+                    )
+
                 # Get overlay image dimensions
                 image_height, image_width, _ = overlay.shape
                 x, y = self._img_position  # Position where the overlay will be placed
 
+                frames_written = 0
                 while cap.isOpened():
                     ret, frame = cap.read()
                     if not ret:
@@ -197,6 +218,10 @@ def _add_image_to_video_sync(
 
                     # Write the modified frame to the output video
                     output_video.write(frame)
+                    frames_written += 1
+
+                if frames_written == 0:
+                    raise ValueError("Could not read the input video: it has no frames")
             finally:
                 if cap is not None:
                     cap.release()
@@ -222,7 +247,10 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "imag
             ConverterResult: The result containing filename of the converted video.
 
         Raises:
-            ValueError: If the input type is not supported.
+            ModuleNotFoundError: If OpenCV is not installed.
+            ValueError: If the input type is not supported, the image path is invalid, the input video
+                format is unsupported, the input video cannot be read, the output video cannot be
+                written, or the overlay image cannot be decoded.
         """
         if not self.input_supported(input_type):
             raise ValueError("Input type not supported")

tests/unit/converter/test_add_image_video_converter.py

@@ -159,3 +159,163 @@ def test_add_image_video_converter_rejects_output_path(video_converter_sample_vi
             video_path=video_converter_sample_video,
             output_path=tmp_path / "output.mp4",  # type: ignore[call-arg]
         )
+
+
+def _write_video(path: Path, *, fps: float, frames: int = 5, size: tuple[int, int] = (128, 96)) -> bytes:
+    import cv2
+
+    writer = cv2.VideoWriter(str(path), cv2.VideoWriter.fourcc(*"mp4v"), fps, size)
+    assert writer.isOpened()
+    try:
+        for _ in range(frames):
+            writer.write(np.zeros((size[1], size[0], 3), dtype=np.uint8))
+    finally:
+        writer.release()
+    return path.read_bytes()
+
+
+def _png(image: np.ndarray) -> bytes:
+    import cv2
+
+    success, encoded = cv2.imencode(".png", image)
+    assert success
+    return bytes(encoded)
+
+
+def _first_frame_and_fps(video_bytes: bytes, path: Path) -> tuple[np.ndarray, float]:
+    import cv2
+
+    path.write_bytes(video_bytes)
+    capture = cv2.VideoCapture(str(path))
+    try:
+        ok, frame = capture.read()
+        assert ok
+        return frame, capture.get(cv2.CAP_PROP_FPS)
+    finally:
+        capture.release()
+
+
+@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
+def test_add_image_to_video_accepts_grayscale_image(tmp_path: Path) -> None:
+    video_bytes = _write_video(tmp_path / "input.mp4", fps=10)
+    converter = AddImageVideoConverter(video_path=tmp_path / "input.mp4", img_position=(8, 8), img_resize_size=(32, 24))
+    gray = np.full((24, 32), 200, dtype=np.uint8)
+
+    with patch("pyrit.converter.add_image_to_video_converter.DB_DATA_PATH", tmp_path):
+        output = converter._add_image_to_video_sync(video_bytes=video_bytes, image_bytes=_png(gray))
+
+    frame, _ = _first_frame_and_fps(output, tmp_path / "output.mp4")
+    # The overlay's middle is the image's gray; the rest of the frame stays black.
+    assert abs(int(frame[20, 24].mean()) - 200) < 20
+    assert int(frame[80, 100].mean()) < 20
+
+
+@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
+def test_add_image_to_video_scales_16_bit_image(tmp_path: Path) -> None:
+    video_bytes = _write_video(tmp_path / "input.mp4", fps=10)
+    converter = AddImageVideoConverter(video_path=tmp_path / "input.mp4", img_position=(8, 8), img_resize_size=(32, 24))
+    # 40000 of 65535 is 156 of 255; cast without scaling, it would wrap to 64
+    gray16 = np.full((24, 32), 40000, dtype=np.uint16)
+
+    with patch("pyrit.converter.add_image_to_video_converter.DB_DATA_PATH", tmp_path):
+        output = converter._add_image_to_video_sync(video_bytes=video_bytes, image_bytes=_png(gray16))
+
+    frame, _ = _first_frame_and_fps(output, tmp_path / "output.mp4")
+    assert abs(int(frame[20, 24].mean()) - 156) < 20
+
+
+@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
+def test_add_image_to_video_rejects_float_image(tmp_path: Path) -> None:
+    import cv2
+
+    video_bytes = _write_video(tmp_path / "input.mp4", fps=10)
+    converter = AddImageVideoConverter(video_path=tmp_path / "input.mp4", img_position=(8, 8), img_resize_size=(32, 24))
+    success, encoded = cv2.imencode(".tiff", np.full((24, 32), 0.5, dtype=np.float32))
+    assert success
+
+    with patch("pyrit.converter.add_image_to_video_converter.DB_DATA_PATH", tmp_path):
+        with pytest.raises(ValueError, match="Unsupported overlay image depth: float32"):
+            converter._add_image_to_video_sync(video_bytes=video_bytes, image_bytes=encoded.tobytes())
+
+
+@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
+def test_add_image_to_video_keeps_fractional_frame_rate(tmp_path: Path) -> None:
+    video_bytes = _write_video(tmp_path / "input.mp4", fps=29.97)
+    converter = AddImageVideoConverter(video_path=tmp_path / "input.mp4", img_position=(8, 8), img_resize_size=(32, 24))
+
+    with patch("pyrit.converter.add_image_to_video_converter.DB_DATA_PATH", tmp_path):
+        output = converter._add_image_to_video_sync(
+            video_bytes=video_bytes, image_bytes=_png(np.zeros((24, 32, 3), dtype=np.uint8))
+        )
+
+    _, fps = _first_frame_and_fps(output, tmp_path / "output.mp4")
+    assert fps == pytest.approx(29.97, abs=0.01)
+
+
+@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
+@pytest.mark.parametrize("video_bytes", [b"", b"not a video"])
+def test_add_image_to_video_raises_for_unreadable_video(tmp_path: Path, video_bytes: bytes) -> None:
+    converter = AddImageVideoConverter(video_path=tmp_path / "input.mp4")
+    files_before = set(tmp_path.iterdir())
+
+    with patch("pyrit.converter.add_image_to_video_converter.DB_DATA_PATH", tmp_path):
+        with pytest.raises(ValueError, match="Could not read the input video"):
+            converter._add_image_to_video_sync(
+                video_bytes=video_bytes, image_bytes=_png(np.zeros((24, 32, 3), dtype=np.uint8))
+            )
+
+    assert set(tmp_path.iterdir()) == files_before
+
+
+@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
+def test_add_image_to_video_raises_when_video_has_no_frames(tmp_path: Path) -> None:
+    import cv2
+
+    converter = AddImageVideoConverter(video_path=tmp_path / "input.mp4")
+    files_before = set(tmp_path.iterdir())
+    # The container opens but holds no decodable frames
+    capture = MagicMock(spec=cv2.VideoCapture)
+    capture.isOpened.return_value = True
+    capture.read.return_value = (False, None)
+    capture.get.side_effect = lambda prop: {
+        cv2.CAP_PROP_FPS: 10.0,
+        cv2.CAP_PROP_FRAME_WIDTH: 128,
+        cv2.CAP_PROP_FRAME_HEIGHT: 96,
+    }[prop]
+
+    with (
+        patch("pyrit.converter.add_image_to_video_converter.DB_DATA_PATH", tmp_path),
+        patch.object(cv2, "VideoCapture", return_value=capture),
+    ):
+        with pytest.raises(ValueError, match="it has no frames"):
+            converter._add_image_to_video_sync(
+                video_bytes=b"container", image_bytes=_png(np.zeros((24, 32, 3), dtype=np.uint8))
+            )
+
+    capture.release.assert_called_once()
+    assert set(tmp_path.iterdir()) == files_before
+
+
+@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
+def test_add_image_to_video_raises_when_writer_cannot_open(tmp_path: Path) -> None:
+    import cv2
+
+    video_bytes = _write_video(tmp_path / "input.mp4", fps=10)
+    (tmp_path / "input.mp4").unlink()
+    converter = AddImageVideoConverter(video_path=tmp_path / "input.mp4")
+    files_before = set(tmp_path.iterdir())
+    writer = MagicMock(spec=cv2.VideoWriter)
+    writer.isOpened.return_value = False
+
+    with (
+        patch("pyrit.converter.add_image_to_video_converter.DB_DATA_PATH", tmp_path),
+        patch.object(cv2, "VideoWriter", return_value=writer) as video_writer,
+    ):
+        video_writer.fourcc = cv2.VideoWriter_fourcc
+        with pytest.raises(ValueError, match="Could not open a video writer"):
+            converter._add_image_to_video_sync(
+                video_bytes=video_bytes, image_bytes=_png(np.zeros((24, 32, 3), dtype=np.uint8))
+            )
+
+    writer.write.assert_not_called()
+    assert set(tmp_path.iterdir()) == files_before