diff --git a/facefusion/apis/stream_video.py b/facefusion/apis/stream_video.py index 4e93f3c0..278f9694 100644 --- a/facefusion/apis/stream_video.py +++ b/facefusion/apis/stream_video.py @@ -12,7 +12,7 @@ import numpy from facefusion import rtc, state_manager, streamer from facefusion.apis.stream_event import create_receive_event from facefusion.audio import create_empty_audio_frame -from facefusion.codecs import aom_decoder, aom_encoder, vp9_decoder, vp9_encoder, vpx_decoder, vpx_encoder +from facefusion.codecs import aom_decoder, aom_encoder, vpx_decoder, vpx_encoder from facefusion.types import AomDecoder, AomEncoder, AomPointer, BitRate, Resolution, RtcPeer, RtcPeerVideo, VideoCodec, VideoPack, VisionFrame, VpxDecoder, VpxEncoder, VpxPointer @@ -118,18 +118,12 @@ def decode_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | Ao if aom_pointer: return normalize_vision_frame(aom_pointer) - if video_codec == 'vp8': + if video_codec in [ 'vp8', 'vp9' ]: vpx_pointer = vpx_decoder.decode(video_decoder, input_buffer) if vpx_pointer: return normalize_vision_frame(vpx_pointer) - if video_codec == 'vp9': - vpx_pointer = vp9_decoder.decode(video_decoder, input_buffer) - - if vpx_pointer: - return normalize_vision_frame(vpx_pointer) - return None @@ -137,12 +131,9 @@ def encode_video_frame(video_codec : VideoCodec, video_encoder : VpxEncoder | Ao if video_codec == 'av1': return aom_encoder.encode(video_encoder, input_buffer, frame_resolution, frame_index) - if video_codec == 'vp8': + if video_codec in [ 'vp8', 'vp9' ]: return vpx_encoder.encode(video_encoder, input_buffer, frame_resolution, frame_index) - if video_codec == 'vp9': - return vp9_encoder.encode(video_encoder, input_buffer, frame_resolution, frame_index) - return bytes() @@ -156,11 +147,8 @@ def create_video_decoder(video_codec : VideoCodec) -> Optional[VpxDecoder | AomD if video_codec == 'av1': return aom_decoder.create(8) - if video_codec == 'vp8': - return vpx_decoder.create(8) - - if video_codec == 'vp9': - return vp9_decoder.create(8) + if video_codec in [ 'vp8', 'vp9' ]: + return vpx_decoder.create(video_codec, 8) return None @@ -169,11 +157,8 @@ def create_video_encoder(video_codec : VideoCodec, frame_resolution : Resolution if video_codec == 'av1': return aom_encoder.create(frame_resolution, bitrate, 8, 10) - if video_codec == 'vp8': - return vpx_encoder.create(frame_resolution, bitrate, 8, 10) - - if video_codec == 'vp9': - return vp9_encoder.create(frame_resolution, bitrate, 8, 10) + if video_codec in [ 'vp8', 'vp9' ]: + return vpx_encoder.create(video_codec, frame_resolution, bitrate, 8, 10) return None @@ -182,34 +167,25 @@ def destroy_video_decoder(video_codec : VideoCodec, video_decoder : VpxDecoder | if video_codec == 'av1': aom_decoder.destroy(video_decoder) - if video_codec == 'vp8': + if video_codec in [ 'vp8', 'vp9' ]: vpx_decoder.destroy(video_decoder) - if video_codec == 'vp9': - vp9_decoder.destroy(video_decoder) - def destroy_video_encoder(video_codec : VideoCodec, video_encoder : VpxEncoder | AomEncoder) -> None: if video_codec == 'av1': aom_encoder.destroy(video_encoder) - if video_codec == 'vp8': + if video_codec in [ 'vp8', 'vp9' ]: vpx_encoder.destroy(video_encoder) - if video_codec == 'vp9': - vp9_encoder.destroy(video_encoder) - def update_video_encoder_bitrate(video_codec : VideoCodec, video_encoder : VpxEncoder | AomEncoder, bitrate : BitRate) -> bool: if video_codec == 'av1': return aom_encoder.update_bitrate(video_encoder, bitrate) - if video_codec == 'vp8': + if video_codec in [ 'vp8', 'vp9' ]: return vpx_encoder.update_bitrate(video_encoder, bitrate) - if video_codec == 'vp9': - return vp9_encoder.update_bitrate(video_encoder, bitrate) - return False diff --git a/facefusion/codecs/vpx_decoder.py b/facefusion/codecs/vpx_decoder.py index 894a1e01..c6305f71 100644 --- a/facefusion/codecs/vpx_decoder.py +++ b/facefusion/codecs/vpx_decoder.py @@ -3,10 +3,10 @@ import struct from typing import Optional from facefusion.libraries import vpx as vpx_module -from facefusion.types import VideoCodec, VpxDecoder, VpxPointer +from facefusion.types import VpxDecoder, VpxPointer, VxpVideoCodec -def create(video_codec : VideoCodec, thread_count : int) -> Optional[VpxDecoder]: +def create(video_codec : VxpVideoCodec, thread_count : int) -> Optional[VpxDecoder]: vpx_library = vpx_module.create_static_library() if vpx_library: diff --git a/facefusion/codecs/vpx_encoder.py b/facefusion/codecs/vpx_encoder.py index ac7ccaf0..8f61c560 100644 --- a/facefusion/codecs/vpx_encoder.py +++ b/facefusion/codecs/vpx_encoder.py @@ -3,10 +3,10 @@ import struct from typing import Optional from facefusion.libraries import vpx as vpx_module -from facefusion.types import BitRate, Resolution, VideoCodec, VpxEncoder +from facefusion.types import BitRate, Resolution, VideoCodec, VpxEncoder, VxpVideoCodec -def create(video_codec : VideoCodec, frame_resolution : Resolution, bitrate : BitRate, thread_count : int, cpu_count : int) -> Optional[VpxEncoder]: +def create(video_codec : VxpVideoCodec, frame_resolution : Resolution, bitrate : BitRate, thread_count : int, cpu_count : int) -> Optional[VpxEncoder]: vpx_library = vpx_module.create_static_library() if vpx_library: diff --git a/facefusion/types.py b/facefusion/types.py index d6a70e5d..5f942a38 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -99,6 +99,9 @@ Resolution : TypeAlias = Tuple[int, int] AudioCodec : TypeAlias = Literal['opus'] VideoCodec : TypeAlias = Literal['av1', 'vp8', 'vp9'] +AomVideoCodec : TypeAlias = Literal['av1'] +VxpVideoCodec : TypeAlias = Literal['vp8', 'vp9'] + FrameHandler : TypeAlias = Callable[..., None] AomEncoder : TypeAlias = ctypes.Array[ctypes.c_char] diff --git a/tests/test_api_stream_video.py b/tests/test_api_stream_video.py index 486ba673..f3e61bc8 100644 --- a/tests/test_api_stream_video.py +++ b/tests/test_api_stream_video.py @@ -11,7 +11,7 @@ import pytest from facefusion import rtc, rtc_store, state_manager from facefusion.apis.stream_video import create_video_decoder, create_video_encoder, decode_video_frame, destroy_video_decoder, destroy_video_encoder, encode_video_frame, handle_video_frame, receive_video_frames, run_video_encode_loop, update_video_encoder_bitrate -from facefusion.codecs import aom_encoder, vp9_encoder, vpx_encoder +from facefusion.codecs import aom_encoder, vpx_encoder from facefusion.common_helper import is_linux, is_macos, is_windows from facefusion.download import conditional_download from facefusion.hash_helper import create_hash @@ -165,12 +165,9 @@ def test_create_and_destroy_video_decoder(video_codec : VideoCodec) -> None: if video_codec == 'av1': video_encoder = aom_encoder.create((426, 226), 1000, 1, 0) encode_buffer = aom_encoder.encode(video_encoder, input_buffer, (426, 226), 0) - if video_codec == 'vp8': - video_encoder = vpx_encoder.create((426, 226), 1000, 1, 0) + if video_codec in [ 'vp8', 'vp9' ]: + video_encoder = vpx_encoder.create(video_codec, (426, 226), 1000, 1, 0) encode_buffer = vpx_encoder.encode(video_encoder, input_buffer, (426, 226), 0) - if video_codec == 'vp9': - video_encoder = vp9_encoder.create((426, 226), 1000, 1, 0) - encode_buffer = vp9_encoder.encode(video_encoder, input_buffer, (426, 226), 0) video_decoder = create_video_decoder(video_codec) @@ -189,19 +186,15 @@ def test_create_and_destroy_video_encoder(video_codec : VideoCodec) -> None: if video_codec == 'av1': assert aom_encoder.encode(video_encoder, input_buffer, (426, 226), 0) - if video_codec == 'vp8': + if video_codec in [ 'vp8', 'vp9' ]: assert vpx_encoder.encode(video_encoder, input_buffer, (426, 226), 0) - if video_codec == 'vp9': - assert vp9_encoder.encode(video_encoder, input_buffer, (426, 226), 0) destroy_video_encoder(video_codec, video_encoder) if video_codec == 'av1': assert aom_encoder.encode(video_encoder, input_buffer, (426, 226), 1) == bytes() - if video_codec == 'vp8': + if video_codec in [ 'vp8', 'vp9' ]: assert vpx_encoder.encode(video_encoder, input_buffer, (426, 226), 1) == bytes() - if video_codec == 'vp9': - assert vp9_encoder.encode(video_encoder, input_buffer, (426, 226), 1) == bytes() @pytest.mark.parametrize('video_codec', [ 'av1', 'vp8', 'vp9' ]) diff --git a/tests/test_codec_vpx_decoder.py b/tests/test_codec_vpx_decoder.py index e3f27390..bb436143 100644 --- a/tests/test_codec_vpx_decoder.py +++ b/tests/test_codec_vpx_decoder.py @@ -11,6 +11,7 @@ from facefusion.common_helper import is_linux, is_macos, is_windows from facefusion.download import conditional_download from facefusion.hash_helper import create_hash from facefusion.libraries import vpx as vpx_module +from facefusion.types import VxpVideoCodec from facefusion.vision import read_video_frame @@ -23,30 +24,41 @@ def before_all() -> None: vpx_module.pre_check() -def test_create() -> None: - assert create(1) +@pytest.mark.parametrize('video_codec', [ 'vp8', 'vp9' ]) +def test_create(video_codec : VxpVideoCodec) -> None: + assert create(video_codec, 1) with patch('facefusion.libraries.vpx.create_static_library', return_value = None): - assert create(1) is None + assert create(video_codec, 1) is None -def test_decode() -> None: +@pytest.mark.parametrize('video_codec', [ 'vp8', 'vp9' ]) +def test_decode(video_codec : VxpVideoCodec) -> None: video_frame = read_video_frame(get_test_example_file('target-240p.mp4')) video_buffer = cv2.cvtColor(video_frame, cv2.COLOR_BGR2YUV_I420).tobytes() video_resolution = (video_frame.shape[1], video_frame.shape[0]) - vpx_encoder = create_encoder(video_resolution, 1000, 1, 0) + vpx_encoder = create_encoder(video_codec, video_resolution, 1000, 1, 0) encode_buffer = encode(vpx_encoder, video_buffer, video_resolution, 0) - vpx_decoder = create(1) + vpx_decoder = create(video_codec, 1) if is_linux() or is_windows(): - assert create_hash(decode(vpx_decoder, encode_buffer).get('buffer')) == 'dc9c8864' + if video_codec == 'vp8': + assert create_hash(decode(vpx_decoder, encode_buffer).get('buffer')) == 'dc9c8864' + + if video_codec == 'vp9': + assert create_hash(decode(vpx_decoder, encode_buffer).get('buffer')) == '50f4c1fa' if is_macos(): - assert create_hash(decode(vpx_decoder, encode_buffer).get('buffer')) == '87450f70' + if video_codec == 'vp8': + assert create_hash(decode(vpx_decoder, encode_buffer).get('buffer')) == '87450f70' + + if video_codec == 'vp9': + pytest.skip() -def test_destroy() -> None: - vpx_decoder = create(1) +@pytest.mark.parametrize('video_codec', [ 'vp8', 'vp9' ]) +def test_destroy(video_codec : VxpVideoCodec) -> None: + vpx_decoder = create(video_codec, 1) with patch.object(vpx_module.create_static_library(), 'vpx_codec_destroy') as mock: destroy(vpx_decoder) diff --git a/tests/test_codec_vpx_encoder.py b/tests/test_codec_vpx_encoder.py index 289adedd..b70c4c07 100644 --- a/tests/test_codec_vpx_encoder.py +++ b/tests/test_codec_vpx_encoder.py @@ -2,6 +2,8 @@ from unittest.mock import patch import cv2 import pytest + +from facefusion.types import VxpVideoCodec from tests.assert_helper import get_test_example_file, get_test_examples_directory from facefusion import state_manager @@ -22,28 +24,35 @@ def before_all() -> None: vpx_module.pre_check() -def test_create() -> None: - assert create((320, 240), 1000, 8, 16) +@pytest.mark.parametrize('video_codec', [ 'vp8', 'vp9' ]) +def test_create(video_codec : VxpVideoCodec) -> None: + assert create(video_codec, (320, 240), 1000, 8, 16) with patch('facefusion.libraries.vpx.create_static_library', return_value = None): - assert create((320, 240), 1000, 8, 16) is None + assert create(video_codec, (320, 240), 1000, 8, 16) is None -def test_encode() -> None: +@pytest.mark.parametrize('video_codec', [ 'vp8', 'vp9' ]) +def test_encode(video_codec : VxpVideoCodec) -> None: video_frame = read_video_frame(get_test_example_file('target-240p.mp4')) video_buffer = cv2.cvtColor(video_frame, cv2.COLOR_BGR2YUV_I420).tobytes() video_resolution = (video_frame.shape[1], video_frame.shape[0]) - vpx_encoder = create(video_resolution, 1000, 1, 0) + vpx_encoder = create(video_codec, video_resolution, 1000, 1, 0) if is_linux() or is_windows(): - assert create_hash(encode(vpx_encoder, video_buffer, video_resolution, 3)) == 'ce133a1f' + if video_codec == 'vp8': + assert create_hash(encode(vpx_encoder, video_buffer, video_resolution, 3)) == 'ce133a1f' + + if video_codec == 'vp9': + assert create_hash(encode(vpx_encoder, video_buffer, video_resolution, 3)) == 'b0760d5e' if is_macos(): pytest.skip() -def test_destroy() -> None: - vpx_encoder = create((320, 240), 1000, 8, 16) +@pytest.mark.parametrize('video_codec', [ 'vp8', 'vp9' ]) +def test_destroy(video_codec : VxpVideoCodec) -> None: + vpx_encoder = create(video_codec, (320, 240), 1000, 8, 16) with patch.object(vpx_module.create_static_library(), 'vpx_codec_destroy') as mock: destroy(vpx_encoder)