Skip to content

Commit b0097cc

Browse files
authored
Improve VideoEncoder test against FFmpeg CLI (#1058)
1 parent 22bcf4d commit b0097cc

File tree

1 file changed

+49
-14
lines changed

1 file changed

+49
-14
lines changed

test/test_encoders.py

Lines changed: 49 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -919,7 +919,10 @@ def test_against_to_file(self, tmp_path, format, method):
919919
{"pixel_format": "yuv420p", "crf": None, "preset": None},
920920
],
921921
)
922-
def test_video_encoder_against_ffmpeg_cli(self, tmp_path, format, encode_params):
922+
@pytest.mark.parametrize("method", ("to_file", "to_tensor", "to_file_like"))
923+
def test_video_encoder_against_ffmpeg_cli(
924+
self, tmp_path, format, encode_params, method
925+
):
923926
ffmpeg_version = get_ffmpeg_major_version()
924927
if format == "webm" and (
925928
ffmpeg_version == 4 or (IS_WINDOWS and ffmpeg_version in (6, 7))
@@ -967,26 +970,45 @@ def test_video_encoder_against_ffmpeg_cli(self, tmp_path, format, encode_params)
967970
# Output path must be last
968971
ffmpeg_cmd.append(ffmpeg_encoded_path)
969972
subprocess.run(ffmpeg_cmd, check=True)
973+
ffmpeg_frames = self.decode(ffmpeg_encoded_path).data
970974

971975
# Encode with our video encoder
972-
encoder_output_path = str(tmp_path / f"encoder_output.{format}")
973976
encoder = VideoEncoder(frames=source_frames, frame_rate=frame_rate)
974-
encoder.to_file(
975-
dest=encoder_output_path,
976-
pixel_format=pixel_format,
977-
crf=crf,
978-
preset=preset,
979-
)
977+
encoder_output_path = str(tmp_path / f"encoder_output.{format}")
980978

981-
ffmpeg_frames = self.decode(ffmpeg_encoded_path).data
982-
encoder_frames = self.decode(encoder_output_path).data
979+
if method == "to_file":
980+
encoder.to_file(
981+
dest=encoder_output_path,
982+
pixel_format=pixel_format,
983+
crf=crf,
984+
preset=preset,
985+
)
986+
encoder_frames = self.decode(encoder_output_path).data
987+
elif method == "to_tensor":
988+
encoded_output = encoder.to_tensor(
989+
format=format,
990+
pixel_format=pixel_format,
991+
crf=crf,
992+
preset=preset,
993+
)
994+
encoder_frames = self.decode(encoded_output).data
995+
elif method == "to_file_like":
996+
file_like = io.BytesIO()
997+
encoder.to_file_like(
998+
file_like=file_like,
999+
format=format,
1000+
pixel_format=pixel_format,
1001+
crf=crf,
1002+
preset=preset,
1003+
)
1004+
encoder_frames = self.decode(file_like.getvalue()).data
1005+
else:
1006+
raise ValueError(f"Unknown method: {method}")
9831007

9841008
assert ffmpeg_frames.shape[0] == encoder_frames.shape[0]
9851009

986-
# If FFmpeg selects a codec or pixel format that uses qscale (not crf),
987-
# the VideoEncoder outputs *slightly* different frames.
988-
# There may be additional subtle differences in the encoder.
989-
percentage = 94 if ffmpeg_version == 6 or format == "avi" else 99
1010+
# MPEG codec used for avi format does not accept CRF
1011+
percentage = 94 if format == "avi" else 99
9901012

9911013
# Check that PSNR between both encoded versions is high
9921014
for ff_frame, enc_frame in zip(ffmpeg_frames, encoder_frames):
@@ -996,6 +1018,19 @@ def test_video_encoder_against_ffmpeg_cli(self, tmp_path, format, encode_params)
9961018
ff_frame, enc_frame, percentage=percentage, atol=2
9971019
)
9981020

1021+
# Check that video metadata is the same
1022+
if method == "to_file":
1023+
fields = ["duration", "duration_ts", "r_frame_rate", "nb_frames"]
1024+
ffmpeg_metadata = self._get_video_metadata(
1025+
ffmpeg_encoded_path,
1026+
fields=fields,
1027+
)
1028+
encoder_metadata = self._get_video_metadata(
1029+
encoder_output_path,
1030+
fields=fields,
1031+
)
1032+
assert ffmpeg_metadata == encoder_metadata
1033+
9991034
def test_to_file_like_custom_file_object(self):
10001035
"""Test to_file_like with a custom file-like object that implements write and seek."""
10011036

0 commit comments

Comments
 (0)