WIP: FFMpeg plugin

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 *
 
 

tests/test_ffmpeg.py

@@ -0,0 +1,164 @@
+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))
+        self.assertEqual(cropped_image.get_size(), (90, 90))
+
+    def test_crop_out_of_bounds(self):
+        # crop rectangle should be clamped to the image boundaries
+        bottom_right_cropped_image = self.image.crop((150, 100, 250, 200))
+        self.assertEqual(bottom_right_cropped_image.get_size(), (50, 50))
+
+        top_left_cropped_image = self.image.crop((-50, -50, 50, 50))
+        self.assertEqual(top_left_cropped_image.get_size(), (50, 50))
+
+        # fail if the crop rectangle is entirely to the left of the image
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((-100, 50, -50, 100))
+        # right edge of crop rectangle is exclusive, so 0 is also invalid
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((-50, 50, 0, 100))
+
+        # fail if the crop rectangle is entirely above the image
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((50, -100, 100, -50))
+        # bottom edge of crop rectangle is exclusive, so 0 is also invalid
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((50, -50, 100, 0))
+
+        # fail if the crop rectangle is entirely to the right of the image
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((250, 50, 300, 100))
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((200, 50, 250, 100))
+
+        # fail if the crop rectangle is entirely below the image
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((50, 200, 100, 250))
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((50, 150, 100, 200))
+
+        # fail if left edge >= right edge
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((125, 25, 25, 125))
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((100, 25, 100, 125))
+
+        # fail if bottom edge >= top edge
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((25, 125, 125, 25))
+        with self.assertRaises(BadImageOperationError):
+            self.image.crop((25, 100, 125, 100))
+
+    def test_rotate(self):
+        rotated_image = self.image.rotate(90)
+        width, height = rotated_image.get_size()
+        self.assertEqual((width, height), (150, 200))
+
+    def test_rotate_without_multiple_of_90(self):
+        with self.assertRaises(UnsupportedRotation) as e:
+            rotated_image = self.image.rotate(45)
+
+    def test_rotate_greater_than_360(self):
+        # 450 should end up the same as a 90 rotation
+        rotated_image = self.image.rotate(450)
+        width, height = rotated_image.get_size()
+        self.assertEqual((width, height), (150, 200))
+
+    def test_rotate_multiple_of_360(self):
+        rotated_image = self.image.rotate(720)
+        width, height = rotated_image.get_size()
+        self.assertEqual((width, height), (200, 150))
+
+    def test_set_background_color_rgb(self):
+        red_background_image = self.image.set_background_color_rgb((255, 0, 0))
+        self.assertFalse(red_background_image.has_alpha())
+        self.assertEqual(red_background_image.image.getpixel((10, 10)), (255, 0, 0))
+
+    def test_set_background_color_rgb_color_argument_check(self):
+        with self.assertRaises(TypeError) as e:
+            self.image.set_background_color_rgb('rgb(255, 0, 0)')
+
+            self.assertEqual(str(e.exception), "the 'color' argument must be a 3-element tuple or list")
+
+    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())

willow/__init__.py

@@ -12,21 +12,28 @@ 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)
     registry.register_image_class(GIFImageFile)
     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()
 

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,

willow/plugins/ffmpeg.py

@@ -0,0 +1,118 @@
+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, crop_rect, 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 crop_rect:
+            pass
+
+        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, crop_rect=None, output_resolution=None):
+        self.source_file = source_file
+        self.crop_rect = crop_rect
+        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, self.crop_rect, size)
+
+    @Image.operation
+    def crop(self, rect):
+        # TODO: Combine with existing rect
+        return FFMpegLazyVideo(self.source_file, rect, self.output_resolution)
+
+    @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.crop_rect, 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.crop_rect, self.output_resolution, 'ogg', 'libtheora')
+        return OggTheoraImageFile(f)
+
+    @Image.operation
+    def save_as_mp4_h264(self, f):
+        transcode(self.source_file.f, f, self.crop_rect, self.output_resolution, 'mp4', 'libx264')
+        return MP4H264ImageFile(f)
+
+
+willow_image_classes = [FFMpegLazyVideo]