Cleanup image common 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_orig_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):
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 """
60 config = mgg.global_config['media_type:mediagoblin.media_types.image']
61
62 resized = exif_fix_image_orientation(resized, exif_tags) # Fix orientation
63
64 filter_config = config['resize_filter']
65 try:
66 resize_filter = PIL_FILTERS[filter_config.upper()]
67 except KeyError:
68 raise Exception('Filter "{0}" not found, choose one of {1}'.format(
69 unicode(filter_config),
70 u', '.join(PIL_FILTERS.keys())))
71
72 resized.thumbnail(new_size, resize_filter)
73
74 # Copy the new file to the conversion subdir, then remotely.
75 tmp_resized_filename = os.path.join(workdir, target_name)
76 with file(tmp_resized_filename, 'w') as resized_file:
77 resized.save(resized_file, quality=config['quality'])
78 store_public(entry, keyname, tmp_resized_filename, target_name)
79
80
81 def resize_tool(entry,
82 force, keyname, orig_file, target_name,
83 conversions_subdir, exif_tags, new_size=None):
84 # Use the default size if new_size was not given
85 if not new_size:
86 max_width = mgg.global_config['media:' + keyname]['max_width']
87 max_height = mgg.global_config['media:' + keyname]['max_height']
88 new_size = (max_width, max_height)
89
90 # If the size of the original file exceeds the specified size for the desized
91 # file, a target_name file is created and later associated with the media
92 # entry.
93 # Also created if the file needs rotation, or if forced.
94 try:
95 im = Image.open(orig_file)
96 except IOError:
97 raise BadMediaFail()
98 if force \
99 or im.size[0] > new_size[0]\
100 or im.size[1] > new_size[1]\
101 or exif_image_needs_rotation(exif_tags):
102 resize_image(
103 entry, im, unicode(keyname), target_name,
104 new_size,
105 exif_tags, conversions_subdir)
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 def common_setup(self):
135 """
136 Set up the workbench directory and pull down the original file
137 """
138 ## @@: Should this be two functions?
139 # Conversions subdirectory to avoid collisions
140 self.conversions_subdir = os.path.join(
141 self.workbench.dir, 'convirsions')
142 os.mkdir(self.conversions_subdir)
143
144 # Pull down and set up the original file
145 self.orig_filename = get_orig_filename(
146 self.entry, self.workbench)
147 self.name_builder = FilenameBuilder(self.orig_filename)
148
149 # Exif extraction
150 self.exif_tags = extract_exif(self.orig_filename)
151
152 def generate_medium_if_applicable(self, size=None):
153 resize_tool(self.entry, False, 'medium', self.orig_filename,
154 self.name_builder.fill('{basename}.medium{ext}'),
155 self.conversions_subdir, self.exif_tags, size)
156
157 def generate_thumb(self, size=None):
158 resize_tool(self.entry, True, 'thumb', self.orig_filename,
159 self.name_builder.fill('{basename}.thumbnail{ext}'),
160 self.conversions_subdir, self.exif_tags, size)
161
162 def copy_original(self):
163 copy_original(
164 self.entry, self.orig_filename,
165 self.name_builder.fill('{basename}{ext}'))
166
167 def extract_metadata(self):
168 # Is there any GPS data
169 gps_data = get_gps_data(self.exif_tags)
170
171 # Insert exif data into database
172 exif_all = clean_exif(self.exif_tags)
173
174 if len(exif_all):
175 self.entry.media_data_init(exif_all=exif_all)
176
177 if len(gps_data):
178 for key in list(gps_data.keys()):
179 gps_data['gps_' + key] = gps_data.pop(key)
180 self.entry.media_data_init(**gps_data)
181
182
183 class InitialProcessor(CommonImageProcessor):
184 """
185 Initial processing step for new images
186 """
187 name = "initial"
188 description = "Initial processing"
189
190 @classmethod
191 def media_is_eligible(cls, entry=None, state=None):
192 """
193 Determine if this media type is eligible for processing
194 """
195 if not state:
196 state = entry.state
197 return state in (
198 "unprocessed", "failed")
199
200 ###############################
201 # Command line interface things
202 ###############################
203
204 @classmethod
205 def generate_parser(cls):
206 parser = argparse.ArgumentParser(
207 description=cls.description,
208 prog=cls.name)
209
210 parser.add_argument(
211 '--size',
212 nargs=2,
213 metavar=('max_width', 'max_height'),
214 type=int)
215
216 parser.add_argument(
217 '--thumb-size',
218 nargs=2,
219 metavar=('max_width', 'max_height'),
220 type=int)
221
222 return parser
223
224 @classmethod
225 def args_to_request(cls, args):
226 return request_from_args(
227 args, ['size', 'thumb_size'])
228
229 def process(self, size=None, thumb_size=None):
230 self.common_setup()
231 self.generate_medium_if_applicable(size=size)
232 self.generate_thumb(size=thumb_size)
233 self.copy_original()
234 self.extract_metadata()
235 self.delete_queue_file()
236
237
238 class Resizer(CommonImageProcessor):
239 """
240 Resizing process steps for processed media
241 """
242 name = 'resize'
243 description = 'Resize image'
244
245 @classmethod
246 def media_is_eligible(cls, entry=None, state=None):
247 """
248 Determine if this media type is eligible for processing
249 """
250 if not state:
251 state = entry.state
252 return state in 'processed'
253
254 ###############################
255 # Command line interface things
256 ###############################
257
258 @classmethod
259 def generate_parser(cls):
260 parser = argparse.ArgumentParser(
261 description=cls.description,
262 prog=cls.name)
263
264 parser.add_argument(
265 '--size',
266 nargs=2,
267 metavar=('max_width', 'max_height'),
268 type=int)
269
270 parser.add_argument(
271 'file',
272 choices=['medium', 'thumb'])
273
274 return parser
275
276 @classmethod
277 def args_to_request(cls, args):
278 return request_from_args(
279 args, ['size', 'file'])
280
281 def process(self, file, size=None):
282 self.common_setup()
283 if file == 'medium':
284 self.generate_medium_if_applicable(size=size)
285 elif file == 'thumb':
286 self.generate_thumb(size=size)
287
288
289 class ImageProcessingManager(ProcessingManager):
290 def __init__(self):
291 super(self.__class__, self).__init__()
292 self.add_processor(InitialProcessor)
293 self.add_processor(Resizer)
294
295
296 if __name__ == '__main__':
297 import sys
298 import pprint
299
300 pp = pprint.PrettyPrinter()
301
302 result = extract_exif(sys.argv[1])
303 gps = get_gps_data(result)
304 clean = clean_exif(result)
305 useful = get_useful(clean)
306
307 print pp.pprint(
308 clean)