diff --git a/README.md b/README.md index f4e73bb..473b3d1 100644 --- a/README.md +++ b/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`). diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..c360d12 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for UVQ inference utilities.""" diff --git a/tests/test_video_reader.py b/tests/test_video_reader.py new file mode 100644 index 0000000..e217856 --- /dev/null +++ b/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() diff --git a/utils/video_reader.py b/utils/video_reader.py index 50ed956..f2511d4 100644 --- a/utils/video_reader.py +++ b/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 diff --git a/uvq1p5_pytorch/utils/uvq1p5.py b/uvq1p5_pytorch/utils/uvq1p5.py index a398fc1..519ca62 100644 --- a/uvq1p5_pytorch/utils/uvq1p5.py +++ b/uvq1p5_pytorch/utils/uvq1p5.py @@ -110,6 +110,7 @@ def infer( orig_fps: float | None = None, ffmpeg_path: str = "ffmpeg", device: str = "cpu", + batch_size: int | None = None, ) -> dict[str, Any]: """Runs UVQ 1.5 inference on a video file. @@ -122,35 +123,35 @@ def infer( calculation. ffmpeg_path: Path to ffmpeg executable. device: Device to run inference on (e.g., 'cpu' or 'cuda'). + batch_size: Maximum number of sampled frames per inference batch. + Defaults to 1 on CPU and 24 on other devices. Returns: A dictionary containing the overall UVQ 1.5 score, per-frame scores, and frame indices. """ - video_1080p, _ = self.load_video( - video_filename, - video_length, - transpose, - fps=fps, - ffmpeg_path=ffmpeg_path, - ) - num_seconds, read_fps, c, h, w = video_1080p.shape - # reshape to (num_seconds * fps, 1, 3, h, w) to process all frames - num_frames = num_seconds * read_fps - video_1080p = video_1080p.reshape(num_frames, 1, c, h, w) - - batch_size = 24 - if num_frames > batch_size: # if video is longer than batch size, run inference in batches to avoid OOM - predictions = [] - with torch.inference_mode(): - for i in range(0, num_frames, batch_size): - batch = video_1080p[i : i + batch_size].to(device) - prediction_batch = self.uvq1p5_core(batch) - predictions.append(prediction_batch) - prediction = torch.cat(predictions, dim=0) - else: - with torch.inference_mode(): - prediction = self.uvq1p5_core(video_1080p.to(device)) + if batch_size is None: + batch_size = 1 if device == "cpu" else 24 + + # Keep convolution weights in the layout of the decoded RGB input on CPU. + if device == "cpu": + self.uvq1p5_core.to(memory_format=torch.channels_last) + + predictions = [] + with torch.inference_mode(): + for video_batch, _ in video_reader.iter_video_batches_1p5( + video_filename, + video_length, + transpose, + video_fps=fps, + batch_size=batch_size, + ffmpeg_path=ffmpeg_path, + ): + batch = torch.from_numpy(video_batch).permute(0, 3, 1, 2) + batch = batch.unsqueeze(1).to(device) + predictions.append(self.uvq1p5_core(batch)) + + prediction = torch.cat(predictions, dim=0) video_score = torch.mean(prediction).item() frame_scores = prediction.cpu().numpy().flatten().tolist() diff --git a/uvq_inference.py b/uvq_inference.py index 0031223..25544d8 100644 --- a/uvq_inference.py +++ b/uvq_inference.py @@ -98,6 +98,7 @@ def run_batch_inference(args): orig_fps=orig_fps, ffmpeg_path=args.ffmpeg_path, device=args.device, + batch_size=args.batch_size, ) score = results["uvq1p5_score"] elif args.model_version == "1.0": @@ -170,6 +171,7 @@ def run_single_inference(args): orig_fps=orig_fps, ffmpeg_path=args.ffmpeg_path, device=args.device, + batch_size=args.batch_size, ) elif args.model_version == "1.0": uvq_inference = uvq1p0.UVQ1p0() @@ -206,6 +208,22 @@ def main(): parser = setup_parser() args = parser.parse_args() + if args.batch_size is None: + args.batch_size = 1 if args.device == "cpu" else 24 + if args.batch_size < 1: + parser.error("--batch_size must be at least 1") + if args.threads is None and args.device == "cpu": + try: + available_cpus = len(os.sched_getaffinity(0)) + except AttributeError: + available_cpus = os.cpu_count() or 1 + args.threads = min(4, available_cpus) + if args.threads is not None and args.threads < 1: + parser.error("--threads must be at least 1") + if args.threads is not None: + torch.set_num_threads(args.threads) + torch.set_num_interop_threads(1) + if args.device == "cuda" and not torch.cuda.is_available(): print("Error: CUDA is not available, please use --device cpu") return @@ -276,6 +294,24 @@ def setup_parser(): help="Frames per second to sample for UVQ1.5. -1 to sample all frames." " Ignored for UVQ1.0.", ) + parser.add_argument( + "--batch_size", + type=int, + default=None, + help=( + "Maximum sampled frames per UVQ1.5 inference batch. Defaults to 1" + " on CPU and 24 on CUDA." + ), + ) + parser.add_argument( + "--threads", + type=int, + default=None, + help=( + "PyTorch CPU inference threads. Defaults to at most 4 available" + " CPUs on CPU; PyTorch chooses the thread count on CUDA." + ), + ) parser.add_argument( "--output_all_stats", action="store_true",