Fix UVQ 1.5 frame selection and speed up CPU inference

README.md

@@ -111,6 +111,8 @@ Gaming_1080P-0ce6_orig.mp4,3.880362033843994
 *   `--output OUTPUT`: Path to save the output scores to a file. Scores will be saved in JSON format.
 *   `--device DEVICE`: Device to run inference on (e.g., `cpu` or `cuda`).
 *   `--fps FPS`: (UVQ 1.5 only) Frames per second to sample. Default is 1. Use -1 to sample all frames.
+*   `--batch_size BATCH_SIZE`: (UVQ 1.5 only) Maximum sampled frames per inference batch. Defaults to 1 on CPU and 24 on CUDA.
+*   `--threads THREADS`: Number of PyTorch CPU inference threads. CPU inference defaults to the smaller of four and the number of CPUs available to the process.
 *   `--output_all_stats`: If specified, print all stats in JSON format to stdout.
 *   `--ffmpeg_path`: Path to FFmpeg executable (default: `ffmpeg`).
 *   `--ffprobe_path`: Path to FFprobe executable (default: `ffprobe`).

tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for UVQ inference utilities."""

tests/test_video_reader.py

@@ -0,0 +1,95 @@
+"""Tests for the bounded UVQ 1.5 video reader."""
+
+import io
+import unittest
+from unittest import mock
+
+import numpy as np
+
+from utils import video_reader
+
+
+class _FakeProcess:
+  """Minimal subprocess.Popen stand-in for raw-video reader tests."""
+
+  def __init__(self, output: bytes, return_code: int = 0):
+    self.stdout = io.BytesIO(output)
+    self._configured_return_code = return_code
+    self.returncode = None
+
+  def poll(self):
+    return self.returncode
+
+  def terminate(self):
+    self.returncode = -15
+
+  def wait(self):
+    if self.returncode is None:
+      self.returncode = self._configured_return_code
+    return self.returncode
+
+
+class VideoReaderTest(unittest.TestCase):
+
+  def test_ffmpeg_samples_before_scaling_without_shell_command(self):
+    path = "video path with spaces;not-a-command.mp4"
+    command = video_reader._video_1p5_command(
+        path,
+        transpose=True,
+        video_fps=1,
+        video_height=1080,
+        video_width=1920,
+        ffmpeg_path="ffmpeg",
+    )
+
+    self.assertIn(path, command)
+    video_filter = command[command.index("-vf") + 1]
+    self.assertLess(video_filter.index("fps="), video_filter.index("scale="))
+    self.assertLess(
+        video_filter.index("fps="), video_filter.index("transpose=")
+    )
+    self.assertIn("start_time=0:round=up", video_filter)
+
+  @mock.patch("utils.video_reader.subprocess.Popen")
+  def test_batches_are_float32_and_truncated_input_is_padded(self, popen):
+    first_frame = bytes([0, 64, 128, 192, 224, 255])
+    popen.return_value = _FakeProcess(first_frame)
+
+    batches = list(
+        video_reader.iter_video_batches_1p5(
+            "input.mp4",
+            video_length=2,
+            video_fps=1,
+            video_height=1,
+            video_width=2,
+            batch_size=1,
+        )
+    )
+
+    self.assertEqual([real_frames for _, real_frames in batches], [1, 0])
+    self.assertEqual(batches[0][0].dtype, np.float32)
+    np.testing.assert_allclose(
+        batches[0][0].reshape(-1),
+        np.frombuffer(first_frame, dtype=np.uint8).astype(np.float32)
+        * (2.0 / 255.0)
+        - 1.0,
+    )
+    np.testing.assert_array_equal(batches[1][0], -1.0)
+
+  @mock.patch("utils.video_reader.subprocess.Popen")
+  def test_incomplete_rgb_frame_is_rejected(self, popen):
+    popen.return_value = _FakeProcess(b"short")
+
+    with self.assertRaisesRegex(RuntimeError, "incomplete RGB frame"):
+      list(
+          video_reader.iter_video_batches_1p5(
+              "input.mp4",
+              video_length=1,
+              video_height=1,
+              video_width=2,
+          )
+      )
+
+
+if __name__ == "__main__":
+  unittest.main()

utils/video_reader.py

@@ -19,6 +19,7 @@
 import os
 import subprocess
 import tempfile
+from collections.abc import Iterator
 
 import numpy as np
 
@@ -157,16 +158,53 @@ def load_video_1p0(
   return video, video_resized
 
 
-def load_video_1p5(
+def _video_1p5_command(
+    filepath: str,
+    transpose: bool,
+    video_fps: int,
+    video_height: int,
+    video_width: int,
+    ffmpeg_path: str,
+) -> list[str]:
+  """Build the FFmpeg command used by the UVQ 1.5 reader."""
+  transpose_filter = "transpose=2," if transpose else ""
+  video_filter = (
+      f"fps={video_fps}:start_time=0:round=up,{transpose_filter}"
+      f"scale=w={video_width}:h={video_height}:flags=bicubic,format=rgb24"
+  )
+  return [
+      ffmpeg_path,
+      "-hide_banner",
+      "-loglevel",
+      "error",
+      "-i",
+      filepath,
+      "-vf",
+      video_filter,
+      "-fps_mode",
+      "passthrough",
+      "-f",
+      "rawvideo",
+      "-pix_fmt",
+      "rgb24",
+      "pipe:1",
+  ]
+
+
+def iter_video_batches_1p5(
     filepath: str,
     video_length: int,
     transpose: bool = False,
     video_fps: int = 1,
     video_height: int = 1080,
     video_width: int = 1920,
+    batch_size: int = 24,
     ffmpeg_path: str = "ffmpeg",
-) -> tuple[np.ndarray, int]:
-  """Load input video for UVQ 1.5.
+) -> Iterator[tuple[np.ndarray, int]]:
+  """Decode normalized RGB frames for UVQ 1.5 in batches.
+
+  A truncated decode is padded with black frames up to
+  ``video_length * video_fps`` frames.
 
   Args:
     filepath: Path to the video file.
@@ -175,74 +213,156 @@ def load_video_1p5(
     video_fps: Frames per second to sample for inference.
     video_height: Height of the video to resize to.
     video_width: Width of the video to resize to.
+    batch_size: Maximum number of frames per yielded batch.
+    ffmpeg_path: Path to the FFmpeg executable.
 
-  Returns:
-    A tuple containing the loaded video as a numpy array and the number of
-    real frames.
+  Yields:
+    A tuple containing a float32 array of shape
+    (frames, video_height, video_width, 3) with values in [-1, 1], and the
+    number of real decoded frames in that batch (0 for padding batches).
+
+  Raises:
+    ValueError: If batch_size is less than 1.
+    RuntimeError: If FFmpeg returns an incomplete RGB frame.
+    subprocess.CalledProcessError: If FFmpeg exits with an error.
+    RuntimeError: If FFmpeg does not decode a single frame.
   """
-  video_channel = 3
-  # Rotate video if requested
-  if transpose:
-    transpose_param = "transpose=2,"  # rotate 90 degrees counterclockwise
-  else:
-    transpose_param = ""
+  if batch_size < 1:
+    raise ValueError("batch_size must be at least 1")
 
-  # Sample at constant frame rate, and save as RGB24 (RGBRGB...)
-  fd, temp_filename = tempfile.mkstemp()
-  cmd = (
-      f"{ffmpeg_path} -i {filepath} -vf"
-      f" {transpose_param}scale=w={video_width}:h={video_height}:flags=bicubic,format=rgb24"
-      f" -r {video_fps} -f rawvideo -pix_fmt rgb24 -y {temp_filename}"
+  video_channels = 3
+  single_frame_size = video_width * video_height * video_channels
+  expected_frames = video_length * video_fps
+  command = _video_1p5_command(
+      filepath,
+      transpose,
+      video_fps,
+      video_height,
+      video_width,
+      ffmpeg_path,
   )
+  logging.info("Run with cmd: %s\n", subprocess.list2cmdline(command))
 
-  try:
-    logging.info("Run with cmd:% s\n", cmd)
-    subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
-  except subprocess.CalledProcessError as error:
-    logging.fatal(
-        "Run with cmd: %s \n terminated with return code %s\n%s",
-        cmd,
-        str(error.returncode),
-        error.output,
+  num_real_frames = 0
+  emitted_frames = 0
+  with tempfile.TemporaryFile() as error_file:
+    process = subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=error_file,
     )
-    raise error
+    assert process.stdout is not None
+    completed_read = False
 
-  # For video, the entire video is divided into 1s chunks in 5 fps
-  with open(temp_filename, "rb") as rgb_file:
-    single_frame_size = video_width * video_height * video_channel
-    full_decode_size = video_length * video_fps * single_frame_size
-    rgb_file.seek(0, 2)
-    rgb_file_size = rgb_file.tell()
-    rgb_file.seek(0)
-    num_real_frames = rgb_file_size // single_frame_size
-    assert rgb_file_size >= single_frame_size, (
-        f"Decoding failed to output a single frame: {rgb_file_size} <"
-        f" {single_frame_size}"
-    )
-    if rgb_file_size < full_decode_size:
-      logging.warning(
-          "Decoding may be truncated: %d bytes (%d frames) < %d bytes (%d"
-          " frames), or video length (%ds) may be too incorrect",
-          rgb_file_size,
-          rgb_file_size / single_frame_size,
-          full_decode_size,
-          full_decode_size / single_frame_size,
-          video_length,
-      )
+    try:
+      while True:
+        rgb = process.stdout.read(batch_size * single_frame_size)
+        if not rgb:
+          break
+        if len(rgb) % single_frame_size != 0:
+          raise RuntimeError(
+              "FFmpeg returned an incomplete RGB frame: "
+              f"{len(rgb)} bytes is not divisible by {single_frame_size}"
+          )
 
-    rgb = _extend_array(bytearray(rgb_file.read()), full_decode_size)
-    video = (
-        np.reshape(
-            np.frombuffer(rgb, "uint8"),
-            (video_length, int(video_fps), video_height, video_width, 3),
+        decoded_frames = len(rgb) // single_frame_size
+        num_real_frames += decoded_frames
+        frames_to_emit = min(decoded_frames, expected_frames - emitted_frames)
+        if frames_to_emit <= 0:
+          continue
+
+        rgb_array = np.frombuffer(
+            rgb, dtype=np.uint8, count=frames_to_emit * single_frame_size
         )
-        / 255.0
-        - 0.5
-    ) * 2
+        video = rgb_array.reshape(
+            frames_to_emit,
+            video_height,
+            video_width,
+            video_channels,
+        ).astype(np.float32)
+        video *= 2.0 / 255.0
+        video -= 1.0
+        emitted_frames += frames_to_emit
+        yield video, frames_to_emit
+      completed_read = True
+      return_code = process.wait()
+    finally:
+      process.stdout.close()
+      if process.poll() is None:
+        process.terminate()
+        process.wait()
 
-  # Delete temp files
-  os.close(fd)
-  os.remove(temp_filename)
-  logging.info("Load %s done successfully.", filepath)
+    if completed_read and return_code != 0:
+      error_file.seek(0)
+      error_output = error_file.read().decode(errors="replace")
+      raise subprocess.CalledProcessError(
+          return_code, command, stderr=error_output
+      )
+
+  if num_real_frames < 1:
+    raise RuntimeError(
+        f"Decoding failed to output a single frame for '{filepath}'"
+    )
+  if num_real_frames < expected_frames:
+    logging.warning(
+        "Decoding may be truncated: %d frames < %d frames, or video length"
+        " (%ds) may be too incorrect",
+        num_real_frames,
+        expected_frames,
+        video_length,
+    )
+
+  while emitted_frames < expected_frames:
+    padding_frames = min(batch_size, expected_frames - emitted_frames)
+    padding = np.full(
+        (padding_frames, video_height, video_width, video_channels),
+        -1.0,
+        dtype=np.float32,
+    )
+    emitted_frames += padding_frames
+    yield padding, 0
+
+
+def load_video_1p5(
+    filepath: str,
+    video_length: int,
+    transpose: bool = False,
+    video_fps: int = 1,
+    video_height: int = 1080,
+    video_width: int = 1920,
+    ffmpeg_path: str = "ffmpeg",
+) -> tuple[np.ndarray, int]:
+  """Load input video for UVQ 1.5.
+
+  Args:
+    filepath: Path to the video file.
+    video_length: Length of the video in seconds.
+    transpose: Whether to transpose the video.
+    video_fps: Frames per second to sample for inference.
+    video_height: Height of the video to resize to.
+    video_width: Width of the video to resize to.
+    ffmpeg_path: Path to the FFmpeg executable.
+
+  Returns:
+    A tuple containing the loaded video as a numpy array and the number of
+    real frames.
+  """
+  batches = []
+  num_real_frames = 0
+  for batch, real_frames in iter_video_batches_1p5(
+      filepath,
+      video_length,
+      transpose,
+      video_fps,
+      video_height,
+      video_width,
+      ffmpeg_path=ffmpeg_path,
+  ):
+    batches.append(batch)
+    num_real_frames += real_frames
 
+  video = np.concatenate(batches).reshape(
+      video_length, video_fps, video_height, video_width, 3
+  )
+  logging.info("Load %s done successfully.", filepath)
   return video, num_real_frames