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 os
  18 import json
  19 import logging
  20 import subprocess
  21 import pkg_resources
  22
  23 from mediagoblin import mg_globals as mgg
  24 from mediagoblin.processing import create_pub_filepath, \
  25     FilenameBuilder
  26
  27 from mediagoblin.media_types.stl import model_loader
  28
  29
  30 _log = logging.getLogger(__name__)
  31 SUPPORTED_FILETYPES = ['stl', 'obj']
  32 MEDIA_TYPE = 'mediagoblin.media_types.stl'
  33
  34 BLEND_FILE = pkg_resources.resource_filename(
  35     'mediagoblin.media_types.stl',
  36     os.path.join(
  37         'assets',
  38         'blender_render.blend'))
  39 BLEND_SCRIPT = pkg_resources.resource_filename(
  40     'mediagoblin.media_types.stl',
  41     os.path.join(
  42         'assets',
  43         'blender_render.py'))
  44
  45
  46 def sniff_handler(media_file, **kw):
  47     _log.info('Sniffing {0}'.format(MEDIA_TYPE))
  48     if kw.get('media') is not None:
  49         name, ext = os.path.splitext(kw['media'].filename)
  50         clean_ext = ext[1:].lower()
  51
  52         if clean_ext in SUPPORTED_FILETYPES:
  53             _log.info('Found file extension in supported filetypes')
  54             return MEDIA_TYPE
  55         else:
  56             _log.debug('Media present, extension not found in {0}'.format(
  57                     SUPPORTED_FILETYPES))
  58     else:
  59         _log.warning('Need additional information (keyword argument \'media\')'
  60                      ' to be able to handle sniffing')
  61
  62     return None
  63
  64
  65 def blender_render(config):
  66     """
  67     Called to prerender a model.
  68     """
  69     env = {"RENDER_SETUP" : json.dumps(config), "DISPLAY":":0"}
  70     subprocess.call(
  71         ["blender",
  72          "-b", BLEND_FILE,
  73          "-F", "JPEG",
  74          "-P", BLEND_SCRIPT],
  75         env=env)
  76
  77
  78 def process_stl(proc_state):
  79     """Code to process an stl or obj model. Will be run by celery.
  80
  81     A Workbench() represents a local tempory dir. It is automatically
  82     cleaned up when this function exits.
  83     """
  84     entry = proc_state.entry
  85     workbench = proc_state.workbench
  86
  87     queued_filepath = entry.queued_media_file
  88     queued_filename = workbench.localized_file(
  89         mgg.queue_store, queued_filepath, 'source')
  90     name_builder = FilenameBuilder(queued_filename)
  91
  92     ext = queued_filename.lower().strip()[-4:]
  93     if ext.startswith("."):
  94         ext = ext[1:]
  95     else:
  96         ext = None
  97
  98     # Attempt to parse the model file and divine some useful
  99     # information about it.
 100     with open(queued_filename, 'rb') as model_file:
 101         model = model_loader.auto_detect(model_file, ext)
 102
 103     # generate preview images
 104     greatest = [model.width, model.height, model.depth]
 105     greatest.sort()
 106     greatest = greatest[-1]
 107
 108     def snap(name, camera, width=640, height=640, project="ORTHO"):
 109         filename = name_builder.fill(name)
 110         workbench_path = workbench.joinpath(filename)
 111         shot = {
 112             "model_path": queued_filename,
 113             "model_ext": ext,
 114             "camera_coord": camera,
 115             "camera_focus": model.average,
 116             "camera_clip": greatest*10,
 117             "greatest": greatest,
 118             "projection": project,
 119             "width": width,
 120             "height": height,
 121             "out_file": workbench_path,
 122             }
 123         blender_render(shot)
 124
 125         # make sure the image rendered to the workbench path
 126         assert os.path.exists(workbench_path)
 127
 128         # copy it up!
 129         with open(workbench_path, 'rb') as rendered_file:
 130             public_path = create_pub_filepath(entry, filename)
 131
 132             with mgg.public_store.get_file(public_path, "wb") as public_file:
 133                 public_file.write(rendered_file.read())
 134
 135         return public_path
 136
 137     thumb_path = snap(
 138         "{basename}.thumb.jpg",
 139         [0, greatest*-1.5, greatest],
 140         mgg.global_config['media:thumb']['max_width'],
 141         mgg.global_config['media:thumb']['max_height'],
 142         project="PERSP")
 143
 144     perspective_path = snap(
 145         "{basename}.perspective.jpg",
 146         [0, greatest*-1.5, greatest], project="PERSP")
 147
 148     topview_path = snap(
 149         "{basename}.top.jpg",
 150         [model.average[0], model.average[1], greatest*2])
 151
 152     frontview_path = snap(
 153         "{basename}.front.jpg",
 154         [model.average[0], greatest*-2, model.average[2]])
 155
 156     sideview_path = snap(
 157         "{basename}.side.jpg",
 158         [greatest*-2, model.average[1], model.average[2]])
 159
 160     ## Save the public file stuffs
 161     model_filepath = create_pub_filepath(
 162         entry, name_builder.fill('{basename}{ext}'))
 163
 164     with mgg.public_store.get_file(model_filepath, 'wb') as model_file:
 165         with open(queued_filename, 'rb') as queued_file:
 166             model_file.write(queued_file.read())
 167
 168     # Remove queued media file from storage and database.
 169     # queued_filepath is in the task_id directory which should
 170     # be removed too, but fail if the directory is not empty to be on
 171     # the super-safe side.
 172     mgg.queue_store.delete_file(queued_filepath)      # rm file
 173     mgg.queue_store.delete_dir(queued_filepath[:-1])  # rm dir
 174     entry.queued_media_file = []
 175
 176     # Insert media file information into database
 177     media_files_dict = entry.setdefault('media_files', {})
 178     media_files_dict[u'original'] = model_filepath
 179     media_files_dict[u'thumb'] = thumb_path
 180     media_files_dict[u'perspective'] = perspective_path
 181     media_files_dict[u'top'] = topview_path
 182     media_files_dict[u'side'] = sideview_path
 183     media_files_dict[u'front'] = frontview_path
 184
 185     # Put model dimensions into the database
 186     dimensions = {
 187         "center_x" : model.average[0],
 188         "center_y" : model.average[1],
 189         "center_z" : model.average[2],
 190         "width" : model.width,
 191         "height" : model.height,
 192         "depth" : model.depth,
 193         "file_type" : ext,
 194         }
 195     entry.media_data_init(**dimensions)