refactor get_orig_filename to return an acceptable filename to the processor.
[mediagoblin.git] / mediagoblin / media_types / image / processing.py
1 # GNU MediaGoblin -- federated, autonomous media hosting
2 # Copyright (C) 2011, 2012 MediaGoblin contributors. See AUTHORS.
3 #
4 # This program is free software: you can redistribute it and/or modify
5 # it under the terms of the GNU Affero General Public License as published by
6 # the Free Software Foundation, either version 3 of the License, or
7 # (at your option) any later version.
8 #
9 # This program is distributed in the hope that it will be useful,
10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12 # GNU Affero General Public License for more details.
13 #
14 # You should have received a copy of the GNU Affero General Public License
15 # along with this program. If not, see <http://www.gnu.org/licenses/>.
16
17 try:
18 from PIL import Image
19 except ImportError:
20 import Image
21 import os
22 import logging
23 import argparse
24
25 from mediagoblin import mg_globals as mgg
26 from mediagoblin.processing import (
27 BadMediaFail, FilenameBuilder,
28 MediaProcessor, ProcessingManager,
29 request_from_args, get_process_filename,
30 store_public, copy_original)
31 from mediagoblin.tools.exif import exif_fix_image_orientation, \
32 extract_exif, clean_exif, get_gps_data, get_useful, \
33 exif_image_needs_rotation
34
35 _log = logging.getLogger(__name__)
36
37 PIL_FILTERS = {
38 'NEAREST': Image.NEAREST,
39 'BILINEAR': Image.BILINEAR,
40 'BICUBIC': Image.BICUBIC,
41 'ANTIALIAS': Image.ANTIALIAS}
42
43 MEDIA_TYPE = 'mediagoblin.media_types.image'
44
45
46 def resize_image(entry, resized, keyname, target_name, new_size,
47 exif_tags, workdir, quality, filter):
48 """
49 Store a resized version of an image and return its pathname.
50
51 Arguments:
52 proc_state -- the processing state for the image to resize
53 resized -- an image from Image.open() of the original image being resized
54 keyname -- Under what key to save in the db.
55 target_name -- public file path for the new resized image
56 exif_tags -- EXIF data for the original image
57 workdir -- directory path for storing converted image files
58 new_size -- 2-tuple size for the resized image
59 quality -- level of compression used when resizing images
60 filter -- One of BICUBIC, BILINEAR, NEAREST, ANTIALIAS
61 """
62 resized = exif_fix_image_orientation(resized, exif_tags) # Fix orientation
63
64 try:
65 resize_filter = PIL_FILTERS[filter.upper()]
66 except KeyError:
67 raise Exception('Filter "{0}" not found, choose one of {1}'.format(
68 unicode(filter),
69 u', '.join(PIL_FILTERS.keys())))
70
71 resized.thumbnail(new_size, resize_filter)
72
73 # Copy the new file to the conversion subdir, then remotely.
74 tmp_resized_filename = os.path.join(workdir, target_name)
75 with file(tmp_resized_filename, 'w') as resized_file:
76 resized.save(resized_file, quality=quality)
77 store_public(entry, keyname, tmp_resized_filename, target_name)
78
79
80 def resize_tool(entry,
81 force, keyname, orig_file, target_name,
82 conversions_subdir, exif_tags, quality, filter, new_size=None):
83 # Use the default size if new_size was not given
84 if not new_size:
85 max_width = mgg.global_config['media:' + keyname]['max_width']
86 max_height = mgg.global_config['media:' + keyname]['max_height']
87 new_size = (max_width, max_height)
88
89 # If the size of the original file exceeds the specified size for the desized
90 # file, a target_name file is created and later associated with the media
91 # entry.
92 # Also created if the file needs rotation, or if forced.
93 try:
94 im = Image.open(orig_file)
95 except IOError:
96 raise BadMediaFail()
97 if force \
98 or im.size[0] > new_size[0]\
99 or im.size[1] > new_size[1]\
100 or exif_image_needs_rotation(exif_tags):
101 resize_image(
102 entry, im, unicode(keyname), target_name,
103 new_size,
104 exif_tags, conversions_subdir,
105 quality, filter)
106
107
108 SUPPORTED_FILETYPES = ['png', 'gif', 'jpg', 'jpeg', 'tiff']
109
110
111 def sniff_handler(media_file, **kw):
112 _log.info('Sniffing {0}'.format(MEDIA_TYPE))
113 if kw.get('media') is not None: # That's a double negative!
114 name, ext = os.path.splitext(kw['media'].filename)
115 clean_ext = ext[1:].lower() # Strip the . from ext and make lowercase
116
117 if clean_ext in SUPPORTED_FILETYPES:
118 _log.info('Found file extension in supported filetypes')
119 return MEDIA_TYPE
120 else:
121 _log.debug('Media present, extension not found in {0}'.format(
122 SUPPORTED_FILETYPES))
123 else:
124 _log.warning('Need additional information (keyword argument \'media\')'
125 ' to be able to handle sniffing')
126
127 return None
128
129
130 class CommonImageProcessor(MediaProcessor):
131 """
132 Provides a base for various media processing steps
133 """
134 # list of acceptable file keys in order of prefrence for reprocessing
135 acceptable_files = ['original', 'medium']
136
137 def common_setup(self):
138 """
139 Set up the workbench directory and pull down the original file
140 """
141 self.image_config = mgg.global_config[
142 'media_type:mediagoblin.media_types.image']
143
144 ## @@: Should this be two functions?
145 # Conversions subdirectory to avoid collisions
146 self.conversions_subdir = os.path.join(
147 self.workbench.dir, 'convirsions')
148 os.mkdir(self.conversions_subdir)
149
150 # Pull down and set up the processing file
151 self.process_filename = get_process_filename(
152 self.entry, self.workbench, self.acceptable_files)
153 self.name_builder = FilenameBuilder(self.process_filename)
154
155 # Exif extraction
156 self.exif_tags = extract_exif(self.process_filename)
157
158 def generate_medium_if_applicable(self, size=None, quality=None,
159 filter=None):
160 if not quality:
161 quality = self.image_config['quality']
162 if not filter:
163 filter = self.image_config['resize_filter']
164
165 resize_tool(self.entry, False, 'medium', self.process_filename,
166 self.name_builder.fill('{basename}.medium{ext}'),
167 self.conversions_subdir, self.exif_tags, quality,
168 filter, size)
169
170 def generate_thumb(self, size=None, quality=None, filter=None):
171 if not quality:
172 quality = self.image_config['quality']
173 if not filter:
174 filter = self.image_config['resize_filter']
175
176 resize_tool(self.entry, True, 'thumb', self.process_filename,
177 self.name_builder.fill('{basename}.thumbnail{ext}'),
178 self.conversions_subdir, self.exif_tags, quality,
179 filter, size)
180
181 def copy_original(self):
182 copy_original(
183 self.entry, self.process_filename,
184 self.name_builder.fill('{basename}{ext}'))
185
186 def extract_metadata(self):
187 # Is there any GPS data
188 gps_data = get_gps_data(self.exif_tags)
189
190 # Insert exif data into database
191 exif_all = clean_exif(self.exif_tags)
192
193 if len(exif_all):
194 self.entry.media_data_init(exif_all=exif_all)
195
196 if len(gps_data):
197 for key in list(gps_data.keys()):
198 gps_data['gps_' + key] = gps_data.pop(key)
199 self.entry.media_data_init(**gps_data)
200
201
202 class InitialProcessor(CommonImageProcessor):
203 """
204 Initial processing step for new images
205 """
206 name = "initial"
207 description = "Initial processing"
208
209 @classmethod
210 def media_is_eligible(cls, entry=None, state=None):
211 """
212 Determine if this media type is eligible for processing
213 """
214 if not state:
215 state = entry.state
216 return state in (
217 "unprocessed", "failed")
218
219 ###############################
220 # Command line interface things
221 ###############################
222
223 @classmethod
224 def generate_parser(cls):
225 parser = argparse.ArgumentParser(
226 description=cls.description,
227 prog=cls.name)
228
229 parser.add_argument(
230 '--size',
231 nargs=2,
232 metavar=('max_width', 'max_height'),
233 type=int)
234
235 parser.add_argument(
236 '--thumb-size',
237 nargs=2,
238 metavar=('max_width', 'max_height'),
239 type=int)
240
241 parser.add_argument(
242 '--filter',
243 choices=['BICUBIC', 'BILINEAR', 'NEAREST', 'ANTIALIAS'])
244
245 parser.add_argument(
246 '--quality',
247 type=int,
248 help='level of compression used when resizing images')
249
250 return parser
251
252 @classmethod
253 def args_to_request(cls, args):
254 return request_from_args(
255 args, ['size', 'thumb_size', 'filter', 'quality'])
256
257 def process(self, size=None, thumb_size=None, quality=None, filter=None):
258 self.common_setup()
259 self.generate_medium_if_applicable(size=size, filter=filter,
260 quality=quality)
261 self.generate_thumb(size=thumb_size, filter=filter, quality=quality)
262 self.copy_original()
263 self.extract_metadata()
264 self.delete_queue_file()
265
266
267 class Resizer(CommonImageProcessor):
268 """
269 Resizing process steps for processed media
270 """
271 name = 'resize'
272 description = 'Resize image'
273 thumb_size = 'size'
274
275 @classmethod
276 def media_is_eligible(cls, entry=None, state=None):
277 """
278 Determine if this media type is eligible for processing
279 """
280 if not state:
281 state = entry.state
282 return state in 'processed'
283
284 ###############################
285 # Command line interface things
286 ###############################
287
288 @classmethod
289 def generate_parser(cls):
290 parser = argparse.ArgumentParser(
291 description=cls.description,
292 prog=cls.name)
293
294 parser.add_argument(
295 '--size',
296 nargs=2,
297 metavar=('max_width', 'max_height'),
298 type=int)
299
300 parser.add_argument(
301 '--filter',
302 choices=['BICUBIC', 'BILINEAR', 'NEAREST', 'ANTIALIAS'])
303
304 parser.add_argument(
305 '--quality',
306 type=int,
307 help='level of compression used when resizing images')
308
309 parser.add_argument(
310 'file',
311 choices=['medium', 'thumb'])
312
313 return parser
314
315 @classmethod
316 def args_to_request(cls, args):
317 return request_from_args(
318 args, ['size', 'file', 'quality', 'filter'])
319
320 def process(self, file, size=None, filter=None, quality=None):
321 self.common_setup()
322 if file == 'medium':
323 self.generate_medium_if_applicable(size=size, filter=filter,
324 quality=quality)
325 elif file == 'thumb':
326 self.generate_thumb(size=size, filter=filter, quality=quality)
327
328
329 class ImageProcessingManager(ProcessingManager):
330 def __init__(self):
331 super(self.__class__, self).__init__()
332 self.add_processor(InitialProcessor)
333 self.add_processor(Resizer)
334
335
336 if __name__ == '__main__':
337 import sys
338 import pprint
339
340 pp = pprint.PrettyPrinter()
341
342 result = extract_exif(sys.argv[1])
343 gps = get_gps_data(result)
344 clean = clean_exif(result)
345 useful = get_useful(clean)
346
347 print pp.pprint(
348 clean)