test_ffmpeg.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. import glob
  2. import subprocess
  3. import pytest
  4. from facefusion import process_manager, state_manager
  5. from facefusion.download import conditional_download
  6. from facefusion.ffmpeg import concat_video, extract_frames, read_audio_buffer
  7. from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_directory_path
  8. from .helper import get_test_example_file, get_test_examples_directory, get_test_output_file, prepare_test_output_directory
  9. @pytest.fixture(scope = 'module', autouse = True)
  10. def before_all() -> None:
  11. process_manager.start()
  12. conditional_download(get_test_examples_directory(),
  13. [
  14. 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg',
  15. 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.mp3',
  16. 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4'
  17. ])
  18. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), get_test_example_file('source.wav') ])
  19. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=25', get_test_example_file('target-240p-25fps.mp4') ])
  20. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=30', get_test_example_file('target-240p-30fps.mp4') ])
  21. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=60', get_test_example_file('target-240p-60fps.mp4') ])
  22. state_manager.init_item('temp_frame_format', 'jpg')
  23. state_manager.init_item('output_audio_encoder', 'aac')
  24. @pytest.fixture(scope = 'function', autouse = True)
  25. def before_each() -> None:
  26. state_manager.clear_item('trim_frame_start')
  27. state_manager.clear_item('trim_frame_end')
  28. prepare_test_output_directory()
  29. def test_extract_frames() -> None:
  30. target_paths =\
  31. [
  32. get_test_example_file('target-240p-25fps.mp4'),
  33. get_test_example_file('target-240p-30fps.mp4'),
  34. get_test_example_file('target-240p-60fps.mp4')
  35. ]
  36. for target_path in target_paths:
  37. temp_directory_path = get_temp_directory_path(target_path)
  38. create_temp_directory(target_path)
  39. assert extract_frames(target_path, '452x240', 30.0) is True
  40. assert len(glob.glob1(temp_directory_path, '*.jpg')) == 324
  41. clear_temp_directory(target_path)
  42. def test_extract_frames_with_trim_start() -> None:
  43. state_manager.init_item('trim_frame_start', 224)
  44. providers =\
  45. [
  46. (get_test_example_file('target-240p-25fps.mp4'), 55),
  47. (get_test_example_file('target-240p-30fps.mp4'), 100),
  48. (get_test_example_file('target-240p-60fps.mp4'), 212)
  49. ]
  50. for target_path, frame_total in providers:
  51. temp_directory_path = get_temp_directory_path(target_path)
  52. create_temp_directory(target_path)
  53. assert extract_frames(target_path, '452x240', 30.0) is True
  54. assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total
  55. clear_temp_directory(target_path)
  56. def test_extract_frames_with_trim_start_and_trim_end() -> None:
  57. state_manager.init_item('trim_frame_start', 124)
  58. state_manager.init_item('trim_frame_end', 224)
  59. providers =\
  60. [
  61. (get_test_example_file('target-240p-25fps.mp4'), 120),
  62. (get_test_example_file('target-240p-30fps.mp4'), 100),
  63. (get_test_example_file('target-240p-60fps.mp4'), 50)
  64. ]
  65. for target_path, frame_total in providers:
  66. temp_directory_path = get_temp_directory_path(target_path)
  67. create_temp_directory(target_path)
  68. assert extract_frames(target_path, '452x240', 30.0) is True
  69. assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total
  70. clear_temp_directory(target_path)
  71. def test_extract_frames_with_trim_end() -> None:
  72. state_manager.init_item('trim_frame_end', 100)
  73. providers =\
  74. [
  75. (get_test_example_file('target-240p-25fps.mp4'), 120),
  76. (get_test_example_file('target-240p-30fps.mp4'), 100),
  77. (get_test_example_file('target-240p-60fps.mp4'), 50)
  78. ]
  79. for target_path, frame_total in providers:
  80. temp_directory_path = get_temp_directory_path(target_path)
  81. create_temp_directory(target_path)
  82. assert extract_frames(target_path, '426x240', 30.0) is True
  83. assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total
  84. clear_temp_directory(target_path)
  85. def test_concat_video() -> None:
  86. output_path = get_test_output_file('test-concat-video.mp4')
  87. temp_output_paths =\
  88. [
  89. get_test_example_file('target-240p.mp4'),
  90. get_test_example_file('target-240p.mp4')
  91. ]
  92. assert concat_video(output_path, temp_output_paths) is True
  93. def test_read_audio_buffer() -> None:
  94. assert isinstance(read_audio_buffer(get_test_example_file('source.mp3'), 1, 1), bytes)
  95. assert isinstance(read_audio_buffer(get_test_example_file('source.wav'), 1, 1), bytes)
  96. assert read_audio_buffer(get_test_example_file('invalid.mp3'), 1, 1) is None