Skip to content

Commit ea5a25b

Browse files
NicolasHugNicolasHug
committed
Refac
1 parent 5a97daf commit ea5a25b

1 file changed

Lines changed: 93 additions & 87 deletions

File tree

‎test/smoke_test.py‎

Lines changed: 93 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,46 @@ def _get_devices():
4343
)
4444

4545

46+
# On CPU we encode with yuv444p (near-lossless with crf=0) so we can compare
47+
# decoded frames directly against the source.
48+
# On CUDA we encode with yuv420p instead, because yuv444p would trigger a CPU
49+
# fallback. Since yuv420p is lossy on random data we can't compare against
50+
# source, so we compare CUDA-decoded output against CPU-decoded output instead.
51+
def _make_decoder_and_ref(tmp_path, device):
52+
"""Returns (decoder, ref_decoder_or_none, source_frames_or_none).
53+
54+
On CPU: returns (cpu_decoder, None, source_frames)
55+
On CUDA: returns (cuda_decoder, cpu_decoder, None)
56+
"""
57+
if device == "cpu":
58+
path, source_frames = _make_video_file(tmp_path, pixel_format="yuv444p")
59+
return VideoDecoder(path, device="cpu"), None, source_frames
60+
else:
61+
path, _ = _make_video_file(tmp_path, pixel_format="yuv420p")
62+
return (
63+
VideoDecoder(path, device="cuda"),
64+
VideoDecoder(path, device="cpu"),
65+
None,
66+
)
67+
68+
69+
def _assert_frames_close(decoded, *, ref_decoded=None, source=None, device):
70+
"""Assert decoded frames are close to reference.
71+
72+
On CPU, compares against source frames (near-lossless yuv444p roundtrip).
73+
On CUDA, compares against CPU-decoded frames (both from same yuv420p file).
74+
"""
75+
actual = decoded.cpu() if device != "cpu" else decoded
76+
if device == "cpu":
77+
assert source is not None
78+
torch.testing.assert_close(actual, source, atol=2, rtol=0)
79+
else:
80+
assert ref_decoded is not None
81+
assert_tensor_close_on_at_least(
82+
actual, ref_decoded.cpu(), percentage=95, atol=3
83+
)
84+
85+
4686
class TestVideoDecoder:
4787
@pytest.mark.parametrize("device", _get_devices())
4888
def test_basics(self, tmp_path, device):
@@ -55,51 +95,34 @@ def test_basics(self, tmp_path, device):
5595

5696
@pytest.mark.parametrize("device", _get_devices())
5797
def test_get_frame_at(self, tmp_path, device):
58-
if device == "cpu":
59-
path, source_frames = _make_video_file(tmp_path, pixel_format="yuv444p")
60-
decoder = VideoDecoder(path, device=device)
61-
frame = decoder.get_frame_at(0)
62-
assert isinstance(frame, Frame)
63-
assert frame.data.shape == (3, HEIGHT, WIDTH)
64-
assert frame.data.dtype == torch.uint8
65-
torch.testing.assert_close(frame.data, source_frames[0], atol=2, rtol=0)
66-
else:
67-
path, _ = _make_video_file(tmp_path, pixel_format="yuv420p")
68-
cpu_decoder = VideoDecoder(path, device="cpu")
69-
cuda_decoder = VideoDecoder(path, device="cuda")
70-
cpu_frame = cpu_decoder.get_frame_at(0).data
71-
cuda_frame = cuda_decoder.get_frame_at(0).data.cpu()
72-
assert isinstance(cuda_decoder.get_frame_at(0), Frame)
73-
assert cuda_frame.shape == (3, HEIGHT, WIDTH)
74-
assert cuda_frame.dtype == torch.uint8
75-
assert_tensor_close_on_at_least(
76-
cuda_frame, cpu_frame, percentage=95, atol=3
77-
)
98+
decoder, ref_decoder, source_frames = _make_decoder_and_ref(tmp_path, device)
99+
frame = decoder.get_frame_at(0)
100+
assert isinstance(frame, Frame)
101+
assert frame.data.shape == (3, HEIGHT, WIDTH)
102+
assert frame.data.dtype == torch.uint8
103+
_assert_frames_close(
104+
frame.data,
105+
ref_decoded=ref_decoder.get_frame_at(0).data if ref_decoder else None,
106+
source=source_frames[0] if source_frames is not None else None,
107+
device=device,
108+
)
78109

79110
@pytest.mark.parametrize("device", _get_devices())
80111
def test_get_frames_in_range(self, tmp_path, device):
81-
if device == "cpu":
82-
path, source_frames = _make_video_file(tmp_path, pixel_format="yuv444p")
83-
decoder = VideoDecoder(path, device=device)
84-
batch = decoder.get_frames_in_range(start=0, stop=5)
85-
assert isinstance(batch, FrameBatch)
86-
assert batch.data.shape == (5, 3, HEIGHT, WIDTH)
87-
torch.testing.assert_close(
88-
batch.data, source_frames[:5], atol=2, rtol=0
89-
)
90-
else:
91-
path, _ = _make_video_file(tmp_path, pixel_format="yuv420p")
92-
cpu_decoder = VideoDecoder(path, device="cpu")
93-
cuda_decoder = VideoDecoder(path, device="cuda")
94-
cpu_batch = cpu_decoder.get_frames_in_range(start=0, stop=5).data
95-
cuda_batch = cuda_decoder.get_frames_in_range(start=0, stop=5).data.cpu()
96-
assert isinstance(
97-
cuda_decoder.get_frames_in_range(start=0, stop=5), FrameBatch
98-
)
99-
assert cuda_batch.shape == (5, 3, HEIGHT, WIDTH)
100-
assert_tensor_close_on_at_least(
101-
cuda_batch, cpu_batch, percentage=95, atol=3
102-
)
112+
decoder, ref_decoder, source_frames = _make_decoder_and_ref(tmp_path, device)
113+
batch = decoder.get_frames_in_range(start=0, stop=5)
114+
assert isinstance(batch, FrameBatch)
115+
assert batch.data.shape == (5, 3, HEIGHT, WIDTH)
116+
_assert_frames_close(
117+
batch.data,
118+
ref_decoded=(
119+
ref_decoder.get_frames_in_range(start=0, stop=5).data
120+
if ref_decoder
121+
else None
122+
),
123+
source=source_frames[:5] if source_frames is not None else None,
124+
device=device,
125+
)
103126

104127
@pytest.mark.parametrize("device", _get_devices())
105128
def test_get_frame_played_at(self, tmp_path, device):
@@ -112,54 +135,37 @@ def test_get_frame_played_at(self, tmp_path, device):
112135

113136
@pytest.mark.parametrize("device", _get_devices())
114137
def test_getitem(self, tmp_path, device):
115-
if device == "cpu":
116-
path, source_frames = _make_video_file(tmp_path, pixel_format="yuv444p")
117-
decoder = VideoDecoder(path, device=device)
118-
tensor = decoder[0]
119-
assert tensor.shape == (3, HEIGHT, WIDTH)
120-
torch.testing.assert_close(tensor, source_frames[0], atol=2, rtol=0)
121-
tensors = decoder[2:5]
122-
assert tensors.shape == (3, 3, HEIGHT, WIDTH)
123-
torch.testing.assert_close(
124-
tensors, source_frames[2:5], atol=2, rtol=0
125-
)
126-
else:
127-
path, _ = _make_video_file(tmp_path, pixel_format="yuv420p")
128-
cpu_decoder = VideoDecoder(path, device="cpu")
129-
cuda_decoder = VideoDecoder(path, device="cuda")
130-
cpu_tensor = cpu_decoder[0]
131-
cuda_tensor = cuda_decoder[0].cpu()
132-
assert cuda_tensor.shape == (3, HEIGHT, WIDTH)
133-
assert_tensor_close_on_at_least(
134-
cuda_tensor, cpu_tensor, percentage=95, atol=3
135-
)
136-
cpu_tensors = cpu_decoder[2:5]
137-
cuda_tensors = cuda_decoder[2:5].cpu()
138-
assert cuda_tensors.shape == (3, 3, HEIGHT, WIDTH)
139-
assert_tensor_close_on_at_least(
140-
cuda_tensors, cpu_tensors, percentage=95, atol=3
141-
)
138+
decoder, ref_decoder, source_frames = _make_decoder_and_ref(tmp_path, device)
139+
140+
tensor = decoder[0]
141+
assert tensor.shape == (3, HEIGHT, WIDTH)
142+
_assert_frames_close(
143+
tensor,
144+
ref_decoded=ref_decoder[0] if ref_decoder else None,
145+
source=source_frames[0] if source_frames is not None else None,
146+
device=device,
147+
)
148+
149+
tensors = decoder[2:5]
150+
assert tensors.shape == (3, 3, HEIGHT, WIDTH)
151+
_assert_frames_close(
152+
tensors,
153+
ref_decoded=ref_decoder[2:5] if ref_decoder else None,
154+
source=source_frames[2:5] if source_frames is not None else None,
155+
device=device,
156+
)
142157

143158
@pytest.mark.parametrize("device", _get_devices())
144159
def test_get_all_frames(self, tmp_path, device):
145-
if device == "cpu":
146-
path, source_frames = _make_video_file(tmp_path, pixel_format="yuv444p")
147-
decoder = VideoDecoder(path, device=device)
148-
all_frames = decoder.get_all_frames()
149-
assert all_frames.data.shape == (NUM_FRAMES, 3, HEIGHT, WIDTH)
150-
torch.testing.assert_close(
151-
all_frames.data, source_frames, atol=2, rtol=0
152-
)
153-
else:
154-
path, _ = _make_video_file(tmp_path, pixel_format="yuv420p")
155-
cpu_decoder = VideoDecoder(path, device="cpu")
156-
cuda_decoder = VideoDecoder(path, device="cuda")
157-
cpu_all = cpu_decoder.get_all_frames().data
158-
cuda_all = cuda_decoder.get_all_frames().data.cpu()
159-
assert cuda_all.shape == (NUM_FRAMES, 3, HEIGHT, WIDTH)
160-
assert_tensor_close_on_at_least(
161-
cuda_all, cpu_all, percentage=95, atol=3
162-
)
160+
decoder, ref_decoder, source_frames = _make_decoder_and_ref(tmp_path, device)
161+
all_frames = decoder.get_all_frames()
162+
assert all_frames.data.shape == (NUM_FRAMES, 3, HEIGHT, WIDTH)
163+
_assert_frames_close(
164+
all_frames.data,
165+
ref_decoded=(ref_decoder.get_all_frames().data if ref_decoder else None),
166+
source=source_frames,
167+
device=device,
168+
)
163169

164170
@pytest.mark.parametrize("device", _get_devices())
165171
def test_iteration(self, tmp_path, device):

0 commit comments

Comments
 (0)