mediagoblin/media_types/stl/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 import argparse
  18 import os
  19 import json
  20 import logging
  21 import subprocess
  22 import pkg_resources
  23
  24 from mediagoblin import mg_globals as mgg
  25 from mediagoblin.processing import (
  26     FilenameBuilder, MediaProcessor,
  27     ProcessingManager, request_from_args,
  28     get_process_filename, store_public,
  29     copy_original)
  30
  31 from mediagoblin.media_types.stl import model_loader
  32
  33
  34 _log = logging.getLogger(__name__)
  35 SUPPORTED_FILETYPES = ['stl', 'obj']
  36 MEDIA_TYPE = 'mediagoblin.media_types.stl'
  37
  38 BLEND_FILE = pkg_resources.resource_filename(
  39     'mediagoblin.media_types.stl',
  40     os.path.join(
  41         'assets',
  42         'blender_render.blend'))
  43 BLEND_SCRIPT = pkg_resources.resource_filename(
  44     'mediagoblin.media_types.stl',
  45     os.path.join(
  46         'assets',
  47         'blender_render.py'))
  48
  49
  50 def sniff_handler(media_file, filename):
  51     _log.info('Sniffing {0}'.format(MEDIA_TYPE))
  52
  53     name, ext = os.path.splitext(filename)
  54     clean_ext = ext[1:].lower()
  55
  56     if clean_ext in SUPPORTED_FILETYPES:
  57         _log.info('Found file extension in supported filetypes')
  58         return MEDIA_TYPE
  59     else:
  60         _log.debug('Media present, extension not found in {0}'.format(
  61                 SUPPORTED_FILETYPES))
  62
  63     return None
  64
  65
  66 def blender_render(config):
  67     """
  68     Called to prerender a model.
  69     """
  70     env = {"RENDER_SETUP" : json.dumps(config), "DISPLAY":":0"}
  71     subprocess.call(
  72         ["blender",
  73          "-b", BLEND_FILE,
  74          "-F", "JPEG",
  75          "-P", BLEND_SCRIPT],
  76         env=env)
  77
  78
  79 class CommonStlProcessor(MediaProcessor):
  80     """
  81     Provides a common base for various stl processing steps
  82     """
  83     acceptable_files = ['original']
  84
  85     def common_setup(self):
  86         # Pull down and set up the processing file
  87         self.process_filename = get_process_filename(
  88             self.entry, self.workbench, self.acceptable_files)
  89         self.name_builder = FilenameBuilder(self.process_filename)
  90
  91         self._set_ext()
  92         self._set_model()
  93         self._set_greatest()
  94
  95     def _set_ext(self):
  96         ext = self.name_builder.ext[1:]
  97
  98         if not ext:
  99             ext = None
 100
 101         self.ext = ext
 102
 103     def _set_model(self):
 104         """
 105         Attempt to parse the model file and divine some useful
 106         information about it.
 107         """
 108         with open(self.process_filename, 'rb') as model_file:
 109             self.model = model_loader.auto_detect(model_file, self.ext)
 110
 111     def _set_greatest(self):
 112         greatest = [self.model.width, self.model.height, self.model.depth]
 113         greatest.sort()
 114         self.greatest = greatest[-1]
 115
 116     def copy_original(self):
 117         copy_original(
 118             self.entry, self.process_filename,
 119             self.name_builder.fill('{basename}{ext}'))
 120
 121     def _snap(self, keyname, name, camera, size, project="ORTHO"):
 122         filename = self.name_builder.fill(name)
 123         workbench_path = self.workbench.joinpath(filename)
 124         shot = {
 125             "model_path": self.process_filename,
 126             "model_ext": self.ext,
 127             "camera_coord": camera,
 128             "camera_focus": self.model.average,
 129             "camera_clip": self.greatest*10,
 130             "greatest": self.greatest,
 131             "projection": project,
 132             "width": size[0],
 133             "height": size[1],
 134             "out_file": workbench_path,
 135             }
 136         blender_render(shot)
 137
 138         # make sure the image rendered to the workbench path
 139         assert os.path.exists(workbench_path)
 140
 141         # copy it up!
 142         store_public(self.entry, keyname, workbench_path, filename)
 143
 144     def _skip_processing(self, keyname, **kwargs):
 145         file_metadata = self.entry.get_file_metadata(keyname)
 146
 147         if not file_metadata:
 148             return False
 149         skip = True
 150
 151         if keyname == 'thumb':
 152             if kwargs.get('thumb_size') != file_metadata.get('thumb_size'):
 153                 skip = False
 154         else:
 155             if kwargs.get('size') != file_metadata.get('size'):
 156                 skip = False
 157
 158         return skip
 159
 160     def generate_thumb(self, thumb_size=None):
 161         if not thumb_size:
 162             thumb_size = (mgg.global_config['media:thumb']['max_width'],
 163                           mgg.global_config['media:thumb']['max_height'])
 164
 165         if self._skip_processing('thumb', thumb_size=thumb_size):
 166             return
 167
 168         self._snap(
 169             "thumb",
 170             "{basename}.thumb.jpg",
 171             [0, self.greatest*-1.5, self.greatest],
 172             thumb_size,
 173             project="PERSP")
 174
 175         self.entry.set_file_metadata('thumb', thumb_size=thumb_size)
 176
 177     def generate_perspective(self, size=None):
 178         if not size:
 179             size = (mgg.global_config['media:medium']['max_width'],
 180                     mgg.global_config['media:medium']['max_height'])
 181
 182         if self._skip_processing('perspective', size=size):
 183             return
 184
 185         self._snap(
 186             "perspective",
 187             "{basename}.perspective.jpg",
 188             [0, self.greatest*-1.5, self.greatest],
 189             size,
 190             project="PERSP")
 191
 192         self.entry.set_file_metadata('perspective', size=size)
 193
 194     def generate_topview(self, size=None):
 195         if not size:
 196             size = (mgg.global_config['media:medium']['max_width'],
 197                     mgg.global_config['media:medium']['max_height'])
 198
 199         if self._skip_processing('top', size=size):
 200             return
 201
 202         self._snap(
 203             "top",
 204             "{basename}.top.jpg",
 205             [self.model.average[0], self.model.average[1],
 206              self.greatest*2],
 207             size)
 208
 209         self.entry.set_file_metadata('top', size=size)
 210
 211     def generate_frontview(self, size=None):
 212         if not size:
 213             size = (mgg.global_config['media:medium']['max_width'],
 214                     mgg.global_config['media:medium']['max_height'])
 215
 216         if self._skip_processing('front', size=size):
 217             return
 218
 219         self._snap(
 220             "front",
 221             "{basename}.front.jpg",
 222             [self.model.average[0], self.greatest*-2,
 223              self.model.average[2]],
 224             size)
 225
 226         self.entry.set_file_metadata('front', size=size)
 227
 228     def generate_sideview(self, size=None):
 229         if not size:
 230             size = (mgg.global_config['media:medium']['max_width'],
 231                     mgg.global_config['media:medium']['max_height'])
 232
 233         if self._skip_processing('side', size=size):
 234             return
 235
 236         self._snap(
 237             "side",
 238             "{basename}.side.jpg",
 239             [self.greatest*-2, self.model.average[1],
 240              self.model.average[2]],
 241             size)
 242
 243         self.entry.set_file_metadata('side', size=size)
 244
 245     def store_dimensions(self):
 246         """
 247         Put model dimensions into the database
 248         """
 249         dimensions = {
 250             "center_x": self.model.average[0],
 251             "center_y": self.model.average[1],
 252             "center_z": self.model.average[2],
 253             "width": self.model.width,
 254             "height": self.model.height,
 255             "depth": self.model.depth,
 256             "file_type": self.ext,
 257             }
 258         self.entry.media_data_init(**dimensions)
 259
 260
 261 class InitialProcessor(CommonStlProcessor):
 262     """
 263     Initial processing step for new stls
 264     """
 265     name = "initial"
 266     description = "Initial processing"
 267
 268     @classmethod
 269     def media_is_eligible(cls, entry=None, state=None):
 270         """
 271         Determine if this media type is eligible for processing
 272         """
 273         if not state:
 274             state = entry.state
 275         return state in (
 276             "unprocessed", "failed")
 277
 278     @classmethod
 279     def generate_parser(cls):
 280         parser = argparse.ArgumentParser(
 281             description=cls.description,
 282             prog=cls.name)
 283
 284         parser.add_argument(
 285             '--size',
 286             nargs=2,
 287             metavar=('max_width', 'max_height'),
 288             type=int)
 289
 290         parser.add_argument(
 291             '--thumb_size',
 292             nargs=2,
 293             metavar=('max_width', 'max_height'),
 294             type=int)
 295
 296         return parser
 297
 298     @classmethod
 299     def args_to_request(cls, args):
 300         return request_from_args(
 301             args, ['size', 'thumb_size'])
 302
 303     def process(self, size=None, thumb_size=None):
 304         self.common_setup()
 305         self.generate_thumb(thumb_size=thumb_size)
 306         self.generate_perspective(size=size)
 307         self.generate_topview(size=size)
 308         self.generate_frontview(size=size)
 309         self.generate_sideview(size=size)
 310         self.store_dimensions()
 311         self.copy_original()
 312         self.delete_queue_file()
 313
 314
 315 class Resizer(CommonStlProcessor):
 316     """
 317     Resizing process steps for processed stls
 318     """
 319     name = 'resize'
 320     description = 'Resize thumbnail and mediums'
 321     thumb_size = 'size'
 322
 323     @classmethod
 324     def media_is_eligible(cls, entry=None, state=None):
 325         """
 326         Determine if this media type is eligible for processing
 327         """
 328         if not state:
 329             state = entry.state
 330         return state in 'processed'
 331
 332     @classmethod
 333     def generate_parser(cls):
 334         parser = argparse.ArgumentParser(
 335             description=cls.description,
 336             prog=cls.name)
 337
 338         parser.add_argument(
 339             '--size',
 340             nargs=2,
 341             metavar=('max_width', 'max_height'),
 342             type=int)
 343
 344         parser.add_argument(
 345             'file',
 346             choices=['medium', 'thumb'])
 347
 348         return parser
 349
 350     @classmethod
 351     def args_to_request(cls, args):
 352         return request_from_args(
 353             args, ['size', 'file'])
 354
 355     def process(self, file, size=None):
 356         self.common_setup()
 357         if file == 'medium':
 358             self.generate_perspective(size=size)
 359             self.generate_topview(size=size)
 360             self.generate_frontview(size=size)
 361             self.generate_sideview(size=size)
 362         elif file == 'thumb':
 363             self.generate_thumb(thumb_size=size)
 364
 365
 366 class StlProcessingManager(ProcessingManager):
 367     def __init__(self):
 368         super(StlProcessingManager, self).__init__()
 369         self.add_processor(InitialProcessor)
 370         self.add_processor(Resizer)