100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""CPU-only contracts for pixel-space frame resampling."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from fastvideo.mlx_runtime.frame_upsample import (
|
|
PIXEL_UPSAMPLE_MODES,
|
|
unsharp,
|
|
upsample_frame,
|
|
upsample_frames,
|
|
)
|
|
from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count
|
|
|
|
|
|
def _frame(height: int = 24, width: int = 40, seed: int = 0) -> np.ndarray:
|
|
rng = np.random.default_rng(seed)
|
|
return rng.integers(0, 256, (height, width, 3), dtype=np.uint8)
|
|
|
|
|
|
@pytest.mark.parametrize("mode", PIXEL_UPSAMPLE_MODES)
|
|
def test_upsample_frame_hits_target_size_in_every_mode(mode: str) -> None:
|
|
out = upsample_frame(_frame(), width=80, height=48, mode=mode)
|
|
assert out.shape == (48, 80, 3)
|
|
assert out.dtype == np.uint8
|
|
|
|
|
|
def test_upsample_frame_rejects_unknown_mode() -> None:
|
|
with pytest.raises(ValueError, match="Unsupported pixel upsample mode"):
|
|
upsample_frame(_frame(), width=80, height=48, mode="mitchell")
|
|
|
|
|
|
def test_upsample_frame_rejects_bad_shape() -> None:
|
|
with pytest.raises(ValueError, match="HxWx3"):
|
|
upsample_frame(np.zeros((24, 40), dtype=np.uint8), width=80, height=48)
|
|
|
|
|
|
def test_upsample_frame_rejects_nonpositive_size() -> None:
|
|
with pytest.raises(ValueError, match="must be positive"):
|
|
upsample_frame(_frame(), width=0, height=48)
|
|
|
|
|
|
def test_upsample_frame_at_target_size_is_a_passthrough() -> None:
|
|
frame = _frame()
|
|
np.testing.assert_array_equal(upsample_frame(frame, width=40, height=24), frame)
|
|
|
|
|
|
def test_nearest_upsample_is_exact_block_replication() -> None:
|
|
"""A 2x nearest resize must duplicate pixels, not blend them."""
|
|
frame = _frame(height=4, width=6)
|
|
out = upsample_frame(frame, width=12, height=8, mode="nearest")
|
|
np.testing.assert_array_equal(out, np.repeat(np.repeat(frame, 2, axis=0), 2, axis=1))
|
|
|
|
|
|
def test_unsharp_zero_amount_is_identity() -> None:
|
|
frame = _frame()
|
|
assert unsharp(frame, 0.0) is frame
|
|
|
|
|
|
def test_unsharp_steepens_a_soft_edge() -> None:
|
|
# A blurred step edge: unsharp masking should raise its peak gradient.
|
|
import cv2
|
|
|
|
frame = np.zeros((24, 40, 3), dtype=np.uint8)
|
|
frame[:, 20:] = 200
|
|
frame = cv2.GaussianBlur(frame, (0, 0), 2.0)
|
|
|
|
def peak_gradient(image: np.ndarray) -> float:
|
|
return float(np.abs(np.diff(image[:, :, 0].astype(np.int32), axis=1)).max())
|
|
|
|
assert peak_gradient(unsharp(frame, 0.8)) > peak_gradient(frame)
|
|
|
|
|
|
def test_upsample_frames_preserves_order_and_leaves_input_alone() -> None:
|
|
frames = [_frame(seed=index) for index in range(3)]
|
|
originals = [frame.copy() for frame in frames]
|
|
out = upsample_frames(frames, width=80, height=48)
|
|
|
|
assert len(out) == 3
|
|
assert all(frame.shape == (48, 80, 3) for frame in out)
|
|
# Distinct inputs stay distinct and in order.
|
|
assert not np.array_equal(out[0], out[1])
|
|
for frame, original in zip(frames, originals, strict=True):
|
|
np.testing.assert_array_equal(frame, original)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("target_frames", "factor", "keyframes"),
|
|
[(1, 999, 1), (81, 2, 41), (81, 3, 29), (82, 2, 45), (121, 4, 33)],
|
|
)
|
|
def test_aligned_keyframes_expand_to_exact_target(target_frames: int, factor: int, keyframes: int) -> None:
|
|
assert aligned_keyframe_count(target_frames, factor) == keyframes
|
|
interpolated_frames = (keyframes - 1) * factor + 1
|
|
assert interpolated_frames >= target_frames
|
|
assert keyframes % 4 == 1
|
|
if keyframes > 1:
|
|
assert (keyframes - 5) * factor + 1 < target_frames
|