test_vision.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. import subprocess
  2. import pytest
  3. from facefusion.download import conditional_download
  4. from facefusion.vision import count_video_frame_total, create_image_resolutions, create_video_resolutions, detect_image_resolution, detect_video_fps, detect_video_resolution, get_video_frame, normalize_resolution, pack_resolution, restrict_image_resolution, restrict_video_fps, restrict_video_resolution, unpack_resolution
  5. from .helper import get_test_example_file, get_test_examples_directory
  6. @pytest.fixture(scope = 'module', autouse = True)
  7. def before_all() -> None:
  8. conditional_download(get_test_examples_directory(),
  9. [
  10. 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg',
  11. 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4',
  12. 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-1080p.mp4'
  13. ])
  14. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ])
  15. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-1080p.mp4'), '-vframes', '1', get_test_example_file('target-1080p.jpg') ])
  16. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', '-vf', 'transpose=0', get_test_example_file('target-240p-90deg.jpg') ])
  17. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-1080p.mp4'), '-vframes', '1', '-vf', 'transpose=0', get_test_example_file('target-1080p-90deg.jpg') ])
  18. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=25', get_test_example_file('target-240p-25fps.mp4') ])
  19. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=30', get_test_example_file('target-240p-30fps.mp4') ])
  20. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=60', get_test_example_file('target-240p-60fps.mp4') ])
  21. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'transpose=0', get_test_example_file('target-240p-90deg.mp4') ])
  22. subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-1080p.mp4'), '-vf', 'transpose=0', get_test_example_file('target-1080p-90deg.mp4') ])
  23. def test_detect_image_resolution() -> None:
  24. assert detect_image_resolution(get_test_example_file('target-240p.jpg')) == (426, 226)
  25. assert detect_image_resolution(get_test_example_file('target-240p-90deg.jpg')) == (226, 426)
  26. assert detect_image_resolution(get_test_example_file('target-1080p.jpg')) == (2048, 1080)
  27. assert detect_image_resolution(get_test_example_file('target-1080p-90deg.jpg')) == (1080, 2048)
  28. assert detect_image_resolution('invalid') is None
  29. def test_restrict_image_resolution() -> None:
  30. assert restrict_image_resolution(get_test_example_file('target-1080p.jpg'), (426, 226)) == (426, 226)
  31. assert restrict_image_resolution(get_test_example_file('target-1080p.jpg'), (2048, 1080)) == (2048, 1080)
  32. assert restrict_image_resolution(get_test_example_file('target-1080p.jpg'), (4096, 2160)) == (2048, 1080)
  33. def test_create_image_resolutions() -> None:
  34. assert create_image_resolutions((426, 226)) == [ '106x56', '212x112', '320x170', '426x226', '640x340', '852x452', '1064x564', '1278x678', '1492x792', '1704x904' ]
  35. assert create_image_resolutions((226, 426)) == [ '56x106', '112x212', '170x320', '226x426', '340x640', '452x852', '564x1064', '678x1278', '792x1492', '904x1704' ]
  36. assert create_image_resolutions((2048, 1080)) == [ '512x270', '1024x540', '1536x810', '2048x1080', '3072x1620', '4096x2160', '5120x2700', '6144x3240', '7168x3780', '8192x4320' ]
  37. assert create_image_resolutions((1080, 2048)) == [ '270x512', '540x1024', '810x1536', '1080x2048', '1620x3072', '2160x4096', '2700x5120', '3240x6144', '3780x7168', '4320x8192' ]
  38. assert create_image_resolutions(None) == []
  39. def test_get_video_frame() -> None:
  40. assert get_video_frame(get_test_example_file('target-240p-25fps.mp4')) is not None
  41. assert get_video_frame('invalid') is None
  42. def test_count_video_frame_total() -> None:
  43. assert count_video_frame_total(get_test_example_file('target-240p-25fps.mp4')) == 270
  44. assert count_video_frame_total(get_test_example_file('target-240p-30fps.mp4')) == 324
  45. assert count_video_frame_total(get_test_example_file('target-240p-60fps.mp4')) == 648
  46. assert count_video_frame_total('invalid') == 0
  47. def test_detect_video_fps() -> None:
  48. assert detect_video_fps(get_test_example_file('target-240p-25fps.mp4')) == 25.0
  49. assert detect_video_fps(get_test_example_file('target-240p-30fps.mp4')) == 30.0
  50. assert detect_video_fps(get_test_example_file('target-240p-60fps.mp4')) == 60.0
  51. assert detect_video_fps('invalid') is None
  52. def test_restrict_video_fps() -> None:
  53. assert restrict_video_fps(get_test_example_file('target-1080p.mp4'), 20.0) == 20.0
  54. assert restrict_video_fps(get_test_example_file('target-1080p.mp4'), 25.0) == 25.0
  55. assert restrict_video_fps(get_test_example_file('target-1080p.mp4'), 60.0) == 25.0
  56. def test_detect_video_resolution() -> None:
  57. assert detect_video_resolution(get_test_example_file('target-240p.mp4')) == (426, 226)
  58. assert detect_video_resolution(get_test_example_file('target-240p-90deg.mp4')) == (226, 426)
  59. assert detect_video_resolution(get_test_example_file('target-1080p.mp4')) == (2048, 1080)
  60. assert detect_video_resolution(get_test_example_file('target-1080p-90deg.mp4')) == (1080, 2048)
  61. assert detect_video_resolution('invalid') is None
  62. def test_restrict_video_resolution() -> None:
  63. assert restrict_video_resolution(get_test_example_file('target-1080p.mp4'), (426, 226)) == (426, 226)
  64. assert restrict_video_resolution(get_test_example_file('target-1080p.mp4'), (2048, 1080)) == (2048, 1080)
  65. assert restrict_video_resolution(get_test_example_file('target-1080p.mp4'), (4096, 2160)) == (2048, 1080)
  66. def test_create_video_resolutions() -> None:
  67. assert create_video_resolutions((426, 226)) == [ '426x226', '452x240', '678x360', '904x480', '1018x540', '1358x720', '2036x1080', '2714x1440', '4072x2160', '8144x4320' ]
  68. assert create_video_resolutions((226, 426)) == [ '226x426', '240x452', '360x678', '480x904', '540x1018', '720x1358', '1080x2036', '1440x2714', '2160x4072', '4320x8144' ]
  69. assert create_video_resolutions((2048, 1080)) == [ '456x240', '682x360', '910x480', '1024x540', '1366x720', '2048x1080', '2730x1440', '4096x2160', '8192x4320' ]
  70. assert create_video_resolutions((1080, 2048)) == [ '240x456', '360x682', '480x910', '540x1024', '720x1366', '1080x2048', '1440x2730', '2160x4096', '4320x8192' ]
  71. assert create_video_resolutions(None) == []
  72. def test_normalize_resolution() -> None:
  73. assert normalize_resolution((2.5, 2.5)) == (2, 2)
  74. assert normalize_resolution((3.0, 3.0)) == (4, 4)
  75. assert normalize_resolution((6.5, 6.5)) == (6, 6)
  76. def test_pack_resolution() -> None:
  77. assert pack_resolution((1, 1)) == '0x0'
  78. assert pack_resolution((2, 2)) == '2x2'
  79. def test_unpack_resolution() -> None:
  80. assert unpack_resolution('0x0') == (0, 0)
  81. assert unpack_resolution('2x2') == (2, 2)