Fixed a bug that had binary stls be handled by the ascii stl parser.
[mediagoblin.git] / mediagoblin / media_types / stl / processing.py
index 0a5a24c1a8884fd7993c71ef784cf2650c6bdf4f..77744ac52b58ec889945c8e47ea5717dfce5b6f5 100644 (file)
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import os
+import json
 import logging
+import subprocess
+import pkg_resources
 
 from mediagoblin import mg_globals as mgg
 from mediagoblin.processing import create_pub_filepath, \
@@ -27,6 +30,17 @@ from mediagoblin.media_types.stl import model_loader
 _log = logging.getLogger(__name__)
 SUPPORTED_FILETYPES = ['stl', 'obj']
 
+BLEND_FILE = pkg_resources.resource_filename(
+    'mediagoblin.media_types.stl',
+    os.path.join(
+        'assets',
+        'blender_render.blend'))
+BLEND_SCRIPT = pkg_resources.resource_filename(
+    'mediagoblin.media_types.stl',
+    os.path.join(
+        'assets',
+        'blender_render.py'))
+
 
 def sniff_handler(media_file, **kw):
     if kw.get('media') is not None:
@@ -46,16 +60,30 @@ def sniff_handler(media_file, **kw):
     return False
 
 
-def process_stl(entry):
+def blender_render(config):
+    """
+    Called to prerender a model.
     """
-    Code to process an stl or obj model.
+    arg_string = "blender -b blender_render.blend -F "
+    arg_string +="JPEG -P blender_render.py"
+    env = {"RENDER_SETUP" : json.dumps(config), "DISPLAY":":0"}
+    subprocess.call(
+        ["blender",
+         "-b", BLEND_FILE,
+         "-F", "JPEG",
+         "-P", BLEND_SCRIPT],
+        env=env)
+
+
+def process_stl(proc_state):
+    """Code to process an stl or obj model. Will be run by celery.
+
+    A Workbench() represents a local tempory dir. It is automatically
+    cleaned up when this function exits.
     """
+    entry = proc_state.entry
+    workbench = proc_state.workbench
 
-    workbench = mgg.workbench_manager.create_workbench()
-    # Conversions subdirectory to avoid collisions
-    conversions_subdir = os.path.join(
-        workbench.dir, 'conversions')
-    os.mkdir(conversions_subdir)
     queued_filepath = entry.queued_media_file
     queued_filename = workbench.localized_file(
         mgg.queue_store, queued_filepath, 'source')
@@ -72,9 +100,64 @@ def process_stl(entry):
     with open(queued_filename, 'rb') as model_file:
         model = model_loader.auto_detect(model_file, ext)
 
-    # TODO: generate blender previews
-
-    # Save the public file stuffs
+    # generate preview images
+    greatest = [model.width, model.height, model.depth]
+    greatest.sort()
+    greatest = greatest[-1]
+
+    def snap(name, camera, width=640, height=640, project="ORTHO"):
+        filename = name_builder.fill(name)
+        workbench_path = workbench.joinpath(filename)
+        shot = {
+            "model_path": queued_filename,
+            "model_ext": ext,
+            "camera_coord": camera,
+            "camera_focus": model.average,
+            "camera_clip": greatest*10,
+            "greatest": greatest,
+            "projection": project,
+            "width": width,
+            "height": height,
+            "out_file": workbench_path,
+            }
+        blender_render(shot)
+
+        # make sure the image rendered to the workbench path
+        assert os.path.exists(workbench_path)
+
+        # copy it up!
+        with open(workbench_path, 'rb') as rendered_file:
+            public_path = create_pub_filepath(entry, filename)
+
+            with mgg.public_store.get_file(public_path, "wb") as public_file:
+                public_file.write(rendered_file.read())
+
+        return public_path
+
+    thumb_path = snap(
+        "{basename}.thumb.jpg",
+        [0, greatest*-1.5, greatest],
+        mgg.global_config['media:thumb']['max_width'],
+        mgg.global_config['media:thumb']['max_height'],
+        project="PERSP")
+
+    perspective_path = snap(
+        "{basename}.perspective.jpg",
+        [0, greatest*-1.5, greatest], project="PERSP")
+
+    topview_path = snap(
+        "{basename}.top.jpg",
+        [model.average[0], model.average[1], greatest*2])
+
+    frontview_path = snap(
+        "{basename}.front.jpg",
+        [model.average[0], greatest*-2, model.average[2]])
+
+    sideview_path = snap(
+        "{basename}.side.jpg",
+        [greatest*-2, model.average[1], model.average[2]])
+
+    ## Save the public file stuffs
     model_filepath = create_pub_filepath(
         entry, name_builder.fill('{basename}{ext}'))
 
@@ -82,15 +165,18 @@ def process_stl(entry):
         with open(queued_filename, 'rb') as queued_file:
             model_file.write(queued_file.read())
 
-
     # Remove queued media file from storage and database
     mgg.queue_store.delete_file(queued_filepath)
     entry.queued_media_file = []
-        
+
     # Insert media file information into database
     media_files_dict = entry.setdefault('media_files', {})
     media_files_dict[u'original'] = model_filepath
-    media_files_dict[u'thumb'] = ["mgoblin_static/images/404.png"]
+    media_files_dict[u'thumb'] = thumb_path
+    media_files_dict[u'perspective'] = perspective_path
+    media_files_dict[u'top'] = topview_path
+    media_files_dict[u'side'] = sideview_path
+    media_files_dict[u'front'] = frontview_path
 
     # Put model dimensions into the database
     dimensions = {
@@ -100,8 +186,6 @@ def process_stl(entry):
         "width" : model.width,
         "height" : model.height,
         "depth" : model.depth,
+        "file_type" : ext,
         }
     entry.media_data_init(**dimensions)
-
-    # clean up workbench
-    workbench.destroy_self()