Claude
a4eaabc20e
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BzJLJWVu7RZnbRqm3nYAnX
104 lines
4.3 KiB
Python
104 lines
4.3 KiB
Python
"""Video preparation: the loader choices (rate, skip, cut, every Nth, cap, crop, size), on made-up frames and a real file."""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from modules.media.load import prepare as prep
|
|
|
|
|
|
def _clip(n=48, h=40, w=80):
|
|
"""Frame i is filled with the value i/255, so which frames survived can be read off the pixels."""
|
|
return torch.stack([torch.full((h, w, 3), i / 255.0) for i in range(n)])
|
|
|
|
|
|
def _ids(frames):
|
|
return [round(float(f[0, 0, 0]) * 255) for f in frames]
|
|
|
|
|
|
def test_frames_are_picked_in_loader_order():
|
|
frames, fps, start, span = prep.prepare(_clip(), 24, {"skip_first": 4, "cut_last": 4, "every_nth": 2, "frame_cap": 10})
|
|
assert _ids(frames) == [4, 6, 8, 10, 12, 14, 16, 18, 20, 22]
|
|
assert fps == 12 and start == pytest.approx(4 / 24) and span == pytest.approx(20 / 24)
|
|
|
|
|
|
def test_target_rate_picks_the_frame_showing_at_each_tick():
|
|
frames, fps, *_ = prep.prepare(_clip(48), 48, {"fps": 24})
|
|
assert fps == 24 and _ids(frames) == list(range(0, 48, 2))
|
|
|
|
|
|
def test_nothing_left_is_refused():
|
|
with pytest.raises(ValueError, match="nothing is left"):
|
|
prep.prepare(_clip(10), 24, {"skip_first": 6, "cut_last": 4})
|
|
|
|
|
|
def test_crop_fit_fills_the_target_and_loses_the_overflow():
|
|
out = prep.shape(_clip(2, 40, 80), 32, 32, "crop")
|
|
assert out.shape == (2, 32, 32, 3)
|
|
|
|
|
|
def test_pad_fit_keeps_everything_with_black_bars():
|
|
out = prep.shape(_clip(1, 40, 80) + 0.5, 64, 64, "pad")
|
|
assert out.shape == (1, 64, 64, 3)
|
|
assert float(out[0, 0, 0, 0]) == 0.0 and float(out[0, 32, 32, 0]) > 0.4
|
|
|
|
|
|
def test_stretch_and_one_sided_target_and_multiple():
|
|
assert prep.shape(_clip(1, 40, 80), 64, 64, "stretch").shape[1:3] == (64, 64)
|
|
assert prep.shape(_clip(1, 40, 80), 64, 0).shape[1:3] == (32, 64) # height follows the shape
|
|
assert prep.shape(_clip(1, 50, 90), 0, 0, multiple_of=32).shape[1:3] == (32, 64)
|
|
|
|
|
|
def test_margin_crop_and_a_crop_that_removes_everything():
|
|
assert prep.shape(_clip(1, 40, 80), crop_left=25, crop_right=25).shape[1:3] == (40, 40)
|
|
with pytest.raises(ValueError, match="no picture"):
|
|
prep.shape(_clip(1, 4, 4), crop_left=50, crop_right=50)
|
|
|
|
|
|
def test_sound_follows_the_picked_stretch():
|
|
wave = torch.arange(48000, dtype=torch.float32).reshape(1, 1, -1)
|
|
cut = prep.cut_audio({"waveform": wave, "sample_rate": 48000}, 0.25, 0.5)
|
|
assert cut["waveform"].shape[-1] == 24000 and float(cut["waveform"][0, 0, 0]) == 12000
|
|
|
|
|
|
# --- the node, on a real encoded file --------------------------------------------------------
|
|
|
|
@pytest.fixture
|
|
def stored(comfyui, tmp_path, monkeypatch):
|
|
from core import config, media
|
|
monkeypatch.setattr(config, "MEDIA_DIR", tmp_path / "media")
|
|
import av
|
|
path = tmp_path / "clip.mp4"
|
|
with av.open(str(path), "w") as out:
|
|
stream = out.add_stream("libx264", rate=24)
|
|
stream.width, stream.height, stream.pix_fmt = 64, 32, "yuv420p"
|
|
for i in range(48):
|
|
frame = av.VideoFrame.from_ndarray(torch.full((32, 64, 3), i * 5, dtype=torch.uint8).numpy(), format="rgb24")
|
|
for packet in stream.encode(frame):
|
|
out.mux(packet)
|
|
for packet in stream.encode():
|
|
out.mux(packet)
|
|
return media.save_upload("clip.mp4", path.read_bytes())
|
|
|
|
|
|
def test_the_node_prepares_a_real_video(stored):
|
|
from modules.media.load.nodes import FunUILoadMedia
|
|
options = {"skip_first": 8, "cut_last": 8, "every_nth": 2, "width": 32, "height": 32, "fit": "crop", "multiple_of": 16}
|
|
first, frames, _audio = FunUILoadMedia.execute(stored["id"], full=True, prepare=json.dumps(options)).result
|
|
assert frames.shape == (16, 32, 32, 3) and first.shape[0] == 1
|
|
|
|
|
|
def test_no_options_keeps_the_old_behaviour(stored):
|
|
from modules.media.load.nodes import FunUILoadMedia
|
|
_first, frames, _audio = FunUILoadMedia.execute(stored["id"], full=True).result
|
|
assert frames.shape[0] == 48 and frames.shape[1:3] == (32, 64)
|
|
|
|
|
|
def test_a_cut_that_leaves_nothing_is_said_in_plain_words(stored):
|
|
from modules.media.load.nodes import FunUILoadMedia
|
|
with pytest.raises(RuntimeError, match="cannot be prepared"):
|
|
FunUILoadMedia.execute(stored["id"], full=True, prepare=json.dumps({"skip_first": 40, "cut_last": 20})).result
|
|
with pytest.raises(RuntimeError, match="cannot be prepared"):
|
|
FunUILoadMedia.execute(stored["id"], full=True, prepare="{not json").result
|