mediagoblin/media_types/stl/model_loader.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
  18 import struct
  19
  20
  21 class ThreeDeeParseError(Exception):
  22     pass
  23
  24
  25 class ThreeDee():
  26     """
  27     3D model parser base class.  Derrived classes are used for basic
  28     analysis of 3D models, and are not intended to be used for 3D
  29     rendering.
  30     """
  31
  32     def __init__(self, fileob):
  33         self.verts = []
  34         self.average = [0, 0, 0]
  35         self.min = [None, None, None]
  36         self.max = [None, None, None]
  37         self.width = 0  # x axis
  38         self.depth = 0  # y axis
  39         self.height = 0 # z axis
  40
  41         self.load(fileob)
  42         if not len(self.verts):
  43             raise ThreeDeeParseError("Empty model.")
  44
  45         for vector in self.verts:
  46             for i in range(3):
  47                 num = vector[i]
  48                 self.average[i] += num
  49                 if not self.min[i]:
  50                     self.min[i] = num
  51                     self.max[i] = num
  52                 else:
  53                     if self.min[i] > num:
  54                         self.min[i] = num
  55                     if self.max[i] < num:
  56                         self.max[i] = num
  57
  58         for i in range(3):
  59             self.average[i]/=len(self.verts)
  60
  61         self.width = abs(self.min[0] - self.max[0])
  62         self.depth = abs(self.min[1] - self.max[1])
  63         self.height = abs(self.min[2] - self.max[2])
  64
  65
  66     def load(self, fileob):
  67         """Override this method in your subclass."""
  68         pass
  69
  70
  71 class ObjModel(ThreeDee):
  72     """
  73     Parser for textureless wavefront obj files.  File format
  74     reference: http://en.wikipedia.org/wiki/Wavefront_.obj_file
  75     """
  76
  77     def __vector(self, line, expected=3):
  78         nums = map(float, line.strip().split(" ")[1:])
  79         return tuple(nums[:expected])
  80
  81     def load(self, fileob):
  82         for line in fileob:
  83             line = line.strip()
  84             if line[0] == "v":
  85                 self.verts.append(self.__vector(line))
  86
  87
  88 class BinaryStlModel(ThreeDee):
  89     """
  90     Parser for ascii-encoded stl files.  File format reference:
  91     http://en.wikipedia.org/wiki/STL_%28file_format%29#Binary_STL
  92     """
  93
  94     def load(self, fileob):
  95         fileob.seek(80) # skip the header
  96         count = struct.unpack("<I", fileob.read(4))[0]
  97         for i in range(count):
  98             fileob.read(12) # skip the normal vector
  99             for v in range(3):
 100                 self.verts.append(struct.unpack("<3f", fileob.read(12)))
 101             fileob.read(2) # skip the attribute bytes
 102
 103
 104 def auto_detect(fileob, hint):
 105     """
 106     Attempt to divine which parser to use to divine information about
 107     the model / verify the file."""
 108
 109     if hint == "obj" or not hint:
 110         try:
 111             return ObjModel(fileob)
 112         except ThreeDeeParseError:
 113             pass
 114
 115     if hint == "stl" or not hint:
 116         try:
 117             # HACK Ascii formatted stls are similar enough to obj
 118             # files that we can just use the same parser for both.
 119             # Isn't that something?
 120             return ObjModel(fileob)
 121         except ThreeDeeParseError:
 122             pass
 123         except ValueError:
 124             pass
 125         except IndexError:
 126             pass
 127         try:
 128             # It is pretty important that the binary stl model loader
 129             # is tried second, because its possible for it to parse
 130             # total garbage from plaintext =)
 131             return BinaryStlModel(fileob)
 132         except ThreeDeeParseError:
 133             pass
 134         except MemoryError:
 135             pass
 136
 137     raise ThreeDeeParseError("Could not successfully parse the model :(")