From afc45c197982fb6f05ae8be726745fd0819c08cb Mon Sep 17 00:00:00 2001 From: Karl Hobley Date: Sun, 6 Mar 2022 17:46:20 +0000 Subject: [PATCH] WIP: FFMpeg plugin --- runtests.py | 1 + tests/test_ffmpeg.py | 101 ++++++++++++++++++++++++++++++++ willow/__init__.py | 9 ++- willow/image.py | 16 ++++- willow/plugins/ffmpeg.py | 122 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 246 insertions(+), 3 deletions(-) create mode 100644 tests/test_ffmpeg.py create mode 100644 willow/plugins/ffmpeg.py diff --git a/runtests.py b/runtests.py index 6de0bab1..aea355e1 100755 --- a/runtests.py +++ b/runtests.py @@ -6,6 +6,7 @@ from tests.test_registry import * from tests.test_pillow import * from tests.test_wand import * +from tests.test_ffmpeg import * from tests.test_image import * diff --git a/tests/test_ffmpeg.py b/tests/test_ffmpeg.py new file mode 100644 index 00000000..ddae48f4 --- /dev/null +++ b/tests/test_ffmpeg.py @@ -0,0 +1,101 @@ +import unittest +import io +import imghdr + +from PIL import Image as PILImage + +from willow.image import ( + GIFImageFile, BadImageOperationError, WebMVP9ImageFile, OggTheoraImageFile, MP4H264ImageFile +) +from willow.plugins.ffmpeg import FFMpegLazyVideo, probe + + +class TestFFMpegOperations(unittest.TestCase): + def setUp(self): + self.f = open('tests/images/newtons_cradle.gif', 'rb') + self.image = FFMpegLazyVideo.open(GIFImageFile(self.f)) + + def tearDown(self): + self.f.close() + + def test_get_size(self): + width, height = self.image.get_size() + self.assertEqual(width, 480) + self.assertEqual(height, 360) + + def test_get_frame_count(self): + frames = self.image.get_frame_count() + self.assertEqual(frames, 34) + + def test_resize(self): + resized_image = self.image.resize((100, 75)) + self.assertEqual(resized_image.get_size(), (100, 75)) + + def test_crop(self): + cropped_image = self.image.crop((10, 10, 100, 100)) + + # Cropping not supported, but image will be resized + self.assertEqual(cropped_image.get_size(), (90, 90)) + + def test_rotate(self): + rotated_image = self.image.rotate(90) + width, height = rotated_image.get_size() + + # Not supported, image will not be rotated + self.assertEqual((width, height), (480, 360)) + + def test_set_background_color_rgb(self): + # Not supported, would do nothing + red_background_image = self.image.set_background_color_rgb((255, 0, 0)) + self.assertFalse(red_background_image.has_alpha()) + + def test_save_as_webm_vp9(self): + output = io.BytesIO() + return_value = self.image.save_as_webm_vp9(output) + output.seek(0) + + probe_data = probe(output) + + self.assertEqual(probe_data['format']['format_name'], 'matroska,webm') + self.assertEqual(probe_data['streams'][0]['codec_name'], 'vp9') + self.assertIsInstance(return_value, WebMVP9ImageFile) + self.assertEqual(return_value.f, output) + + def test_save_as_ogg_theora(self): + output = io.BytesIO() + return_value = self.image.save_as_ogg_theora(output) + output.seek(0) + + probe_data = probe(output) + + self.assertEqual(probe_data['format']['format_name'], 'ogg') + self.assertEqual(probe_data['streams'][0]['codec_name'], 'theora') + self.assertIsInstance(return_value, OggTheoraImageFile) + self.assertEqual(return_value.f, output) + + def test_save_as_mp4_h264(self): + output = io.BytesIO() + return_value = self. image.save_as_mp4_h264(output) + output.seek(0) + + probe_data = probe(output) + + self.assertEqual(probe_data['format']['format_name'], 'mov,mp4,m4a,3gp,3g2,mj2') + self.assertEqual(probe_data['streams'][0]['codec_name'], 'h264') + self.assertIsInstance(return_value, MP4H264ImageFile) + self.assertEqual(return_value.f, output) + + def test_has_alpha(self): + has_alpha = self.image.has_alpha() + self.assertFalse(has_alpha) + + def test_has_animation(self): + has_animation = self.image.has_animation() + self.assertTrue(has_animation) + + def test_transparent_gif(self): + with open('tests/images/transparent.gif', 'rb') as f: + image = FFMpegLazyVideo.open(GIFImageFile(f)) + + # Transparency not supported + self.assertFalse(image.has_alpha()) diff --git a/willow/__init__.py b/willow/__init__.py index cfc6ca49..e7787cbb 100644 --- a/willow/__init__.py +++ b/willow/__init__.py @@ -12,8 +12,11 @@ def setup(): RGBAImageBuffer, TIFFImageFile, WebPImageFile, + WebMVP9ImageFile, + OggTheoraImageFile, + MP4H264ImageFile, ) - from willow.plugins import pillow, wand, opencv + from willow.plugins import pillow, wand, opencv, ffmpeg registry.register_image_class(JPEGImageFile) registry.register_image_class(PNGImageFile) @@ -21,12 +24,16 @@ def setup(): registry.register_image_class(BMPImageFile) registry.register_image_class(TIFFImageFile) registry.register_image_class(WebPImageFile) + registry.register_image_class(WebMVP9ImageFile) + registry.register_image_class(OggTheoraImageFile) + registry.register_image_class(MP4H264ImageFile) registry.register_image_class(RGBImageBuffer) registry.register_image_class(RGBAImageBuffer) registry.register_plugin(pillow) registry.register_plugin(wand) registry.register_plugin(opencv) + registry.register_plugin(ffmpeg) setup() diff --git a/willow/image.py b/willow/image.py index 096e0d36..eb769f0e 100644 --- a/willow/image.py +++ b/willow/image.py @@ -100,10 +100,10 @@ def open(cls, f): def save(self, image_format, output): # Get operation name - if image_format not in ['jpeg', 'png', 'gif', 'bmp', 'tiff', 'webp']: + if image_format not in ['jpeg', 'png', 'gif', 'bmp', 'tiff', 'webp', 'webm/vp9', 'ogg/theora', 'mp4/h264']: raise ValueError("Unknown image format: %s" % image_format) - operation_name = 'save_as_' + image_format + operation_name = 'save_as_' + image_format.replace('/', '_') return getattr(self, operation_name)(output) @@ -180,6 +180,18 @@ class WebPImageFile(ImageFile): format_name = 'webp' +class WebMVP9ImageFile(ImageFile): + format_name = 'webm/vp9' + + +class OggTheoraImageFile(ImageFile): + format_name = 'ogg/theora' + + +class MP4H264ImageFile(ImageFile): + format_name = 'mp4/h264' + + INITIAL_IMAGE_CLASSES = { # A mapping of image formats to their initial class 'jpeg': JPEGImageFile, diff --git a/willow/plugins/ffmpeg.py b/willow/plugins/ffmpeg.py new file mode 100644 index 00000000..3d18e2da --- /dev/null +++ b/willow/plugins/ffmpeg.py @@ -0,0 +1,122 @@ +import subprocess +import os.path +from itertools import product +import json +from tempfile import NamedTemporaryFile, TemporaryDirectory + +from willow.image import Image + +from willow.image import ( + GIFImageFile, + WebPImageFile, + WebMVP9ImageFile, + OggTheoraImageFile, + MP4H264ImageFile, +) + + +def probe(file): + with NamedTemporaryFile() as src: + src.write(file.read()) + result = subprocess.run(["ffprobe", "-show_format", "-show_streams", "-loglevel", "quiet", "-print_format", "json", src.name], capture_output=True) + return json.loads(result.stdout) + + +def transcode(source_file, output_file, output_resolution, format, codec): + with NamedTemporaryFile() as src, TemporaryDirectory() as outdir: + src.write(source_file.read()) + + args = ["ffmpeg", "-i", src.name, "-f", format, "-codec:v", codec] + + if output_resolution: + args += ["-s", f"{output_resolution[0]}x{output_resolution[1]}"] + + args.append(os.path.join(outdir, 'out')) + + subprocess.run(args) + + with open(os.path.join(outdir, 'out'), 'rb') as out: + output_file.write(out.read()) + + +class FFMpegLazyVideo(Image): + def __init__(self, source_file, output_resolution=None): + self.source_file = source_file + self.output_resolution = output_resolution + + @Image.operation + def get_size(self): + if self.output_resolution: + return self.output_resolution + + # Find the size from the source file + data = probe(self.source_file.f) + for stream in data['streams']: + if stream['codec_type'] == 'video': + return stream['width'], stream['height'] + + @Image.operation + def get_frame_count(self): + # Find the frame count from the source file + data = probe(self.source_file.f) + for stream in data['streams']: + if stream['codec_type'] == 'video': + return int(stream['nb_frames']) + + @Image.operation + def has_alpha(self): + # Alpha not supported + return False + + @Image.operation + def has_animation(self): + return True + + @Image.operation + def resize(self, size): + return FFMpegLazyVideo(self.source_file, size) + + @Image.operation + def crop(self, rect): + # Not supported, but resize the image to match the crop rect size + left, top, right, bottom = rect + width = right - left + height = bottom - top + return FFMpegLazyVideo(self.source_file, (width, height)) + + @Image.operation + def rotate(self, angle): + # Not supported + return self + + @Image.operation + def set_background_color_rgb(self, color): + # Alpha not supported + return self + + @classmethod + @Image.converter_from(GIFImageFile) + @Image.converter_from(WebPImageFile) + @Image.converter_from(WebMVP9ImageFile) + @Image.converter_from(OggTheoraImageFile) + @Image.converter_from(MP4H264ImageFile) + def open(cls, file): + return cls(file) + + @Image.operation + def save_as_webm_vp9(self, f): + transcode(self.source_file.f, f, self.output_resolution, 'webm', 'libvpx-vp9') + return WebMVP9ImageFile(f) + + @Image.operation + def save_as_ogg_theora(self, f): + transcode(self.source_file.f, f, self.output_resolution, 'ogg', 'libtheora') + return OggTheoraImageFile(f) + + @Image.operation + def save_as_mp4_h264(self, f): + transcode(self.source_file.f, f, self.output_resolution, 'mp4', 'libx264') + return MP4H264ImageFile(f) + + +willow_image_classes = [FFMpegLazyVideo]