Skip to content

Commit 75c40f2

Browse files
kaedrohokaedroho
committed
WIP: FFMpeg plugin
1 parent ed6c532 commit 75c40f2

File tree

5 files changed

+308
-3
lines changed

5 files changed

+308
-3
lines changed

‎runtests.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from tests.test_registry import *
77
from tests.test_pillow import *
88
from tests.test_wand import *
9+
from tests.test_ffmpeg import *
910
from tests.test_image import *
1011

1112

‎tests/test_ffmpeg.py‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
import unittest
2+
import io
3+
import imghdr
4+
5+
from PIL import Image as PILImage
6+
7+
from willow.image import (
8+
GIFImageFile, BadImageOperationError, WebMVP9ImageFile, OggTheoraImageFile, MP4H264ImageFile
9+
)
10+
from willow.plugins.ffmpeg import FFMpegLazyVideo, probe
11+
12+
13+
class TestFFMpegOperations(unittest.TestCase):
14+
def setUp(self):
15+
self.f = open('tests/images/newtons_cradle.gif', 'rb')
16+
self.image = FFMpegLazyVideo.open(GIFImageFile(self.f))
17+
18+
def tearDown(self):
19+
self.f.close()
20+
21+
def test_get_size(self):
22+
width, height = self.image.get_size()
23+
self.assertEqual(width, 480)
24+
self.assertEqual(height, 360)
25+
26+
def test_get_frame_count(self):
27+
frames = self.image.get_frame_count()
28+
self.assertEqual(frames, 34)
29+
30+
def test_resize(self):
31+
resized_image = self.image.resize((100, 75))
32+
self.assertEqual(resized_image.get_size(), (100, 75))
33+
34+
def test_crop(self):
35+
cropped_image = self.image.crop((10, 10, 100, 100))
36+
self.assertEqual(cropped_image.get_size(), (90, 90))
37+
38+
def test_crop_out_of_bounds(self):
39+
# crop rectangle should be clamped to the image boundaries
40+
bottom_right_cropped_image = self.image.crop((150, 100, 250, 200))
41+
self.assertEqual(bottom_right_cropped_image.get_size(), (50, 50))
42+
43+
top_left_cropped_image = self.image.crop((-50, -50, 50, 50))
44+
self.assertEqual(top_left_cropped_image.get_size(), (50, 50))
45+
46+
# fail if the crop rectangle is entirely to the left of the image
47+
with self.assertRaises(BadImageOperationError):
48+
self.image.crop((-100, 50, -50, 100))
49+
# right edge of crop rectangle is exclusive, so 0 is also invalid
50+
with self.assertRaises(BadImageOperationError):
51+
self.image.crop((-50, 50, 0, 100))
52+
53+
# fail if the crop rectangle is entirely above the image
54+
with self.assertRaises(BadImageOperationError):
55+
self.image.crop((50, -100, 100, -50))
56+
# bottom edge of crop rectangle is exclusive, so 0 is also invalid
57+
with self.assertRaises(BadImageOperationError):
58+
self.image.crop((50, -50, 100, 0))
59+
60+
# fail if the crop rectangle is entirely to the right of the image
61+
with self.assertRaises(BadImageOperationError):
62+
self.image.crop((250, 50, 300, 100))
63+
with self.assertRaises(BadImageOperationError):
64+
self.image.crop((200, 50, 250, 100))
65+
66+
# fail if the crop rectangle is entirely below the image
67+
with self.assertRaises(BadImageOperationError):
68+
self.image.crop((50, 200, 100, 250))
69+
with self.assertRaises(BadImageOperationError):
70+
self.image.crop((50, 150, 100, 200))
71+
72+
# fail if left edge >= right edge
73+
with self.assertRaises(BadImageOperationError):
74+
self.image.crop((125, 25, 25, 125))
75+
with self.assertRaises(BadImageOperationError):
76+
self.image.crop((100, 25, 100, 125))
77+
78+
# fail if bottom edge >= top edge
79+
with self.assertRaises(BadImageOperationError):
80+
self.image.crop((25, 125, 125, 25))
81+
with self.assertRaises(BadImageOperationError):
82+
self.image.crop((25, 100, 125, 100))
83+
84+
def test_rotate(self):
85+
rotated_image = self.image.rotate(90)
86+
width, height = rotated_image.get_size()
87+
self.assertEqual((width, height), (150, 200))
88+
89+
def test_rotate_without_multiple_of_90(self):
90+
with self.assertRaises(UnsupportedRotation) as e:
91+
rotated_image = self.image.rotate(45)
92+
93+
def test_rotate_greater_than_360(self):
94+
# 450 should end up the same as a 90 rotation
95+
rotated_image = self.image.rotate(450)
96+
width, height = rotated_image.get_size()
97+
self.assertEqual((width, height), (150, 200))
98+
99+
def test_rotate_multiple_of_360(self):
100+
rotated_image = self.image.rotate(720)
101+
width, height = rotated_image.get_size()
102+
self.assertEqual((width, height), (200, 150))
103+
104+
def test_set_background_color_rgb(self):
105+
red_background_image = self.image.set_background_color_rgb((255, 0, 0))
106+
self.assertFalse(red_background_image.has_alpha())
107+
self.assertEqual(red_background_image.image.getpixel((10, 10)), (255, 0, 0))
108+
109+
def test_set_background_color_rgb_color_argument_check(self):
110+
with self.assertRaises(TypeError) as e:
111+
self.image.set_background_color_rgb('rgb(255, 0, 0)')
112+
113+
self.assertEqual(str(e.exception), "the 'color' argument must be a 3-element tuple or list")
114+
115+
def test_save_as_webm_vp9(self):
116+
output = io.BytesIO()
117+
return_value = self.image.save_as_webm_vp9(output)
118+
output.seek(0)
119+
120+
probe_data = probe(output)
121+
122+
self.assertEqual(probe_data['format']['format_name'], 'matroska,webm')
123+
self.assertEqual(probe_data['streams'][0]['codec_name'], 'vp9')
124+
self.assertIsInstance(return_value, WebMVP9ImageFile)
125+
self.assertEqual(return_value.f, output)
126+
127+
def test_save_as_ogg_theora(self):
128+
output = io.BytesIO()
129+
return_value = self.image.save_as_ogg_theora(output)
130+
output.seek(0)
131+
132+
probe_data = probe(output)
133+
134+
self.assertEqual(probe_data['format']['format_name'], 'ogg')
135+
self.assertEqual(probe_data['streams'][0]['codec_name'], 'theora')
136+
self.assertIsInstance(return_value, OggTheoraImageFile)
137+
self.assertEqual(return_value.f, output)
138+
139+
def test_save_as_mp4_h264(self):
140+
output = io.BytesIO()
141+
return_value = self. image.save_as_mp4_h264(output)
142+
output.seek(0)
143+
144+
probe_data = probe(output)
145+
146+
self.assertEqual(probe_data['format']['format_name'], 'mov,mp4,m4a,3gp,3g2,mj2')
147+
self.assertEqual(probe_data['streams'][0]['codec_name'], 'h264')
148+
self.assertIsInstance(return_value, MP4H264ImageFile)
149+
self.assertEqual(return_value.f, output)
150+
151+
def test_has_alpha(self):
152+
has_alpha = self.image.has_alpha()
153+
self.assertFalse(has_alpha)
154+
155+
def test_has_animation(self):
156+
has_animation = self.image.has_animation()
157+
self.assertTrue(has_animation)
158+
159+
def test_transparent_gif(self):
160+
with open('tests/images/transparent.gif', 'rb') as f:
161+
image = FFMpegLazyVideo.open(GIFImageFile(f))
162+
163+
# Transparency not supported
164+
self.assertFalse(image.has_alpha())

‎willow/__init__.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,21 +12,28 @@ def setup():
1212
RGBAImageBuffer,
1313
TIFFImageFile,
1414
WebPImageFile,
15+
WebMVP9ImageFile,
16+
OggTheoraImageFile,
17+
MP4H264ImageFile,
1518
)
16-
from willow.plugins import pillow, wand, opencv
19+
from willow.plugins import pillow, wand, opencv, ffmpeg
1720

1821
registry.register_image_class(JPEGImageFile)
1922
registry.register_image_class(PNGImageFile)
2023
registry.register_image_class(GIFImageFile)
2124
registry.register_image_class(BMPImageFile)
2225
registry.register_image_class(TIFFImageFile)
2326
registry.register_image_class(WebPImageFile)
27+
registry.register_image_class(WebMVP9ImageFile)
28+
registry.register_image_class(OggTheoraImageFile)
29+
registry.register_image_class(MP4H264ImageFile)
2430
registry.register_image_class(RGBImageBuffer)
2531
registry.register_image_class(RGBAImageBuffer)
2632

2733
registry.register_plugin(pillow)
2834
registry.register_plugin(wand)
2935
registry.register_plugin(opencv)
36+
registry.register_plugin(ffmpeg)
3037

3138
setup()
3239

‎willow/image.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,10 +100,10 @@ def open(cls, f):
100100

101101
def save(self, image_format, output):
102102
# Get operation name
103-
if image_format not in ['jpeg', 'png', 'gif', 'bmp', 'tiff', 'webp']:
103+
if image_format not in ['jpeg', 'png', 'gif', 'bmp', 'tiff', 'webp', 'webm/vp9', 'ogg/theora', 'mp4/h264']:
104104
raise ValueError("Unknown image format: %s" % image_format)
105105

106-
operation_name = 'save_as_' + image_format
106+
operation_name = 'save_as_' + image_format.replace('/', '_')
107107
return getattr(self, operation_name)(output)
108108

109109

@@ -180,6 +180,18 @@ class WebPImageFile(ImageFile):
180180
format_name = 'webp'
181181

182182

183+
class WebMVP9ImageFile(ImageFile):
184+
format_name = 'webm/vp9'
185+
186+
187+
class OggTheoraImageFile(ImageFile):
188+
format_name = 'ogg/theora'
189+
190+
191+
class MP4H264ImageFile(ImageFile):
192+
format_name = 'mp4/h264'
193+
194+
183195
INITIAL_IMAGE_CLASSES = {
184196
# A mapping of image formats to their initial class
185197
'jpeg': JPEGImageFile,
@@ -188,6 +200,9 @@ class WebPImageFile(ImageFile):
188200
'bmp': BMPImageFile,
189201
'tiff': TIFFImageFile,
190202
'webp': WebPImageFile,
203+
'webm/vp9': WebMVP9ImageFile,
204+
'ogg/theora': OggTheoraImageFile,
205+
'mp4/h264': MP4H264ImageFile,
191206
}
192207

193208

‎willow/plugins/ffmpeg.py‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
import subprocess
2+
import os.path
3+
from itertools import product
4+
import json
5+
from tempfile import NamedTemporaryFile, TemporaryDirectory
6+
7+
from willow.image import Image
8+
9+
from willow.image import (
10+
GIFImageFile,
11+
WebPImageFile,
12+
WebMVP9ImageFile,
13+
OggTheoraImageFile,
14+
MP4H264ImageFile,
15+
)
16+
17+
18+
def probe(file):
19+
with NamedTemporaryFile() as src:
20+
src.write(file.read())
21+
result = subprocess.run(["ffprobe", "-show_format", "-show_streams", "-loglevel", "quiet", "-print_format", "json", src.name], capture_output=True)
22+
return json.loads(result.stdout)
23+
24+
25+
def transcode(source_file, output_file, crop_rect, output_resolution, format, codec):
26+
with NamedTemporaryFile() as src, TemporaryDirectory() as outdir:
27+
src.write(source_file.read())
28+
29+
args = ["ffmpeg", "-i", src.name, "-f", format, "-codec:v", codec]
30+
31+
if crop_rect:
32+
pass
33+
34+
if output_resolution:
35+
args += ["-s", f"{output_resolution[0]}x{output_resolution[1]}"]
36+
37+
args.append(os.path.join(outdir, 'out'))
38+
39+
subprocess.run(args)
40+
41+
with open(os.path.join(outdir, 'out'), 'rb') as out:
42+
output_file.write(out.read())
43+
44+
45+
class FFMpegLazyVideo(Image):
46+
def __init__(self, source_file, crop_rect=None, output_resolution=None):
47+
self.source_file = source_file
48+
self.crop_rect = crop_rect
49+
self.output_resolution = output_resolution
50+
51+
@Image.operation
52+
def get_size(self):
53+
if self.output_resolution:
54+
return self.output_resolution
55+
56+
# Find the size from the source file
57+
data = probe(self.source_file.f)
58+
for stream in data['streams']:
59+
if stream['codec_type'] == 'video':
60+
return stream['width'], stream['height']
61+
62+
@Image.operation
63+
def get_frame_count(self):
64+
# Find the frame count from the source file
65+
data = probe(self.source_file.f)
66+
for stream in data['streams']:
67+
if stream['codec_type'] == 'video':
68+
return int(stream['nb_frames'])
69+
70+
@Image.operation
71+
def has_alpha(self):
72+
# Alpha not supported
73+
return False
74+
75+
@Image.operation
76+
def has_animation(self):
77+
return True
78+
79+
@Image.operation
80+
def resize(self, size):
81+
return FFMpegLazyVideo(self.source_file, self.crop_rect, size)
82+
83+
@Image.operation
84+
def crop(self, rect):
85+
# TODO: Combine with existing rect
86+
return FFMpegLazyVideo(self.source_file, rect, self.output_resolution)
87+
88+
@Image.operation
89+
def set_background_color_rgb(self, color):
90+
# Alpha not supported
91+
return self
92+
93+
@classmethod
94+
@Image.converter_from(GIFImageFile)
95+
@Image.converter_from(WebPImageFile)
96+
@Image.converter_from(WebMVP9ImageFile)
97+
@Image.converter_from(OggTheoraImageFile)
98+
@Image.converter_from(MP4H264ImageFile)
99+
def open(cls, file):
100+
return cls(file)
101+
102+
@Image.operation
103+
def save_as_webm_vp9(self, f):
104+
transcode(self.source_file.f, f, self.crop_rect, self.output_resolution, 'webm', 'libvpx-vp9')
105+
return WebMVP9ImageFile(f)
106+
107+
@Image.operation
108+
def save_as_ogg_theora(self, f):
109+
transcode(self.source_file.f, f, self.crop_rect, self.output_resolution, 'ogg', 'libtheora')
110+
return OggTheoraImageFile(f)
111+
112+
@Image.operation
113+
def save_as_mp4_h264(self, f):
114+
transcode(self.source_file.f, f, self.crop_rect, self.output_resolution, 'mp4', 'libx264')
115+
return MP4H264ImageFile(f)
116+
117+
118+
willow_image_classes = [FFMpegLazyVideo]

0 commit comments

Comments
 (0)