io_mesh_ply/import_ply.py

   1 # ##### BEGIN GPL LICENSE BLOCK #####
   2 #
   3 #  This program is free software; you can redistribute it and/or
   4 #  modify it under the terms of the GNU General Public License
   5 #  as published by the Free Software Foundation; either version 2
   6 #  of the License, or (at your option) any later version.
   7 #
   8 #  This program is distributed in the hope that it will be useful,
   9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 #  GNU General Public License for more details.
  12 #
  13 #  You should have received a copy of the GNU General Public License
  14 #  along with this program; if not, write to the Free Software Foundation,
  15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16 #
  17 # ##### END GPL LICENSE BLOCK #####
  18
  19 # <pep8 compliant>
  20
  21 import re
  22 import struct
  23
  24
  25 class element_spec(object):
  26     __slots__ = ("name",
  27                  "count",
  28                  "properties",
  29                  )
  30
  31     def __init__(self, name, count):
  32         self.name = name
  33         self.count = count
  34         self.properties = []
  35
  36     def load(self, format, stream):
  37         if format == b'ascii':
  38             stream = stream.readline().split()
  39         return [x.load(format, stream) for x in self.properties]
  40
  41     def index(self, name):
  42         for i, p in enumerate(self.properties):
  43             if p.name == name:
  44                 return i
  45         return -1
  46
  47
  48 class property_spec(object):
  49     __slots__ = ("name",
  50                  "list_type",
  51                  "numeric_type",
  52                  )
  53
  54     def __init__(self, name, list_type, numeric_type):
  55         self.name = name
  56         self.list_type = list_type
  57         self.numeric_type = numeric_type
  58
  59     def read_format(self, format, count, num_type, stream):
  60         if format == b'ascii':
  61             if num_type == 's':
  62                 ans = []
  63                 for i in range(count):
  64                     s = stream[i]
  65                     if len(s) < 2 or s[0] != '"' or s[-1] != '"':
  66                         print('Invalid string', s)
  67                         print('Note: ply_import.py does not handle whitespace in strings')
  68                         return None
  69                     ans.append(s[1:-1])
  70                 stream[:count] = []
  71                 return ans
  72             if num_type == 'f' or num_type == 'd':
  73                 mapper = float
  74             else:
  75                 mapper = int
  76             ans = [mapper(x) for x in stream[:count]]
  77             stream[:count] = []
  78             return ans
  79         else:
  80             if num_type == 's':
  81                 ans = []
  82                 for i in range(count):
  83                     fmt = format + 'i'
  84                     data = stream.read(struct.calcsize(fmt))
  85                     length = struct.unpack(fmt, data)[0]
  86                     fmt = '%s%is' % (format, length)
  87                     data = stream.read(struct.calcsize(fmt))
  88                     s = struct.unpack(fmt, data)[0]
  89                     ans.append(s[:-1])  # strip the NULL
  90                 return ans
  91             else:
  92                 fmt = '%s%i%s' % (format, count, num_type)
  93                 data = stream.read(struct.calcsize(fmt))
  94                 return struct.unpack(fmt, data)
  95
  96     def load(self, format, stream):
  97         if self.list_type is not None:
  98             count = int(self.read_format(format, 1, self.list_type, stream)[0])
  99             return self.read_format(format, count, self.numeric_type, stream)
 100         else:
 101             return self.read_format(format, 1, self.numeric_type, stream)[0]
 102
 103
 104 class object_spec(object):
 105     __slots__ = ("specs",
 106                 )
 107     'A list of element_specs'
 108     def __init__(self):
 109         self.specs = []
 110
 111     def load(self, format, stream):
 112         return dict([(i.name, [i.load(format, stream) for j in range(i.count)]) for i in self.specs])
 113
 114         '''
 115         # Longhand for above LC
 116         answer = {}
 117         for i in self.specs:
 118             answer[i.name] = []
 119             for j in range(i.count):
 120                 if not j % 100 and meshtools.show_progress:
 121                     Blender.Window.DrawProgressBar(float(j) / i.count, 'Loading ' + i.name)
 122                 answer[i.name].append(i.load(format, stream))
 123         return answer
 124             '''
 125
 126
 127 def read(filepath):
 128     format = b''
 129     texture = b''
 130     version = b'1.0'
 131     format_specs = {b'binary_little_endian': '<',
 132                     b'binary_big_endian': '>',
 133                     b'ascii': b'ascii'}
 134     type_specs = {b'char': 'b',
 135                   b'uchar': 'B',
 136                   b'int8': 'b',
 137                   b'uint8': 'B',
 138                   b'int16': 'h',
 139                   b'uint16': 'H',
 140                   b'short': 'h',
 141                   b'ushort': 'H',
 142                   b'int': 'i',
 143                   b'int32': 'i',
 144                   b'uint': 'I',
 145                   b'uint32': 'I',
 146                   b'float': 'f',
 147                   b'float32': 'f',
 148                   b'float64': 'd',
 149                   b'double': 'd',
 150                   b'string': 's'}
 151     obj_spec = object_spec()
 152     invalid_ply = (None, None, None)
 153
 154     with open(filepath, 'rb') as plyf:
 155         signature = plyf.readline()
 156
 157         if not signature.startswith(b'ply'):
 158             print('Signature line was invalid')
 159             return invalid_ply
 160
 161         valid_header = False
 162         for line in plyf:
 163             tokens = re.split(br'[ \r\n]+', line)
 164
 165             if len(tokens) == 0:
 166                 continue
 167             if tokens[0] == b'end_header':
 168                 valid_header = True
 169                 break
 170             elif tokens[0] == b'comment':
 171                 if len(tokens) < 2:
 172                     continue
 173                 elif tokens[1] == b'TextureFile':
 174                     if len(tokens) < 4:
 175                         print('Invalid texture line')
 176                     else:
 177                         texture = tokens[2]
 178                 continue
 179             elif tokens[0] == b'obj_info':
 180                 continue
 181             elif tokens[0] == b'format':
 182                 if len(tokens) < 3:
 183                     print('Invalid format line')
 184                     return invalid_ply
 185                 if tokens[1] not in format_specs:
 186                     print('Unknown format', tokens[1])
 187                     return invalid_ply
 188                 try:
 189                     version_test = float(tokens[2])
 190                 except Exception as ex:
 191                     print('Unknown version', ex)
 192                     version_test = None
 193                 if version_test != float(version):
 194                     print('Unknown version', tokens[2])
 195                     return invalid_ply
 196                 del version_test
 197                 format = tokens[1]
 198             elif tokens[0] == b'element':
 199                 if len(tokens) < 3:
 200                     print(b'Invalid element line')
 201                     return invalid_ply
 202                 obj_spec.specs.append(element_spec(tokens[1], int(tokens[2])))
 203             elif tokens[0] == b'property':
 204                 if not len(obj_spec.specs):
 205                     print('Property without element')
 206                     return invalid_ply
 207                 if tokens[1] == b'list':
 208                     obj_spec.specs[-1].properties.append(property_spec(tokens[4], type_specs[tokens[2]], type_specs[tokens[3]]))
 209                 else:
 210                     obj_spec.specs[-1].properties.append(property_spec(tokens[2], None, type_specs[tokens[1]]))
 211         if not valid_header:
 212             print("Invalid header ('end_header' line not found!)")
 213             return invalid_ply
 214
 215         obj = obj_spec.load(format_specs[format], plyf)
 216
 217     return obj_spec, obj, texture
 218
 219
 220 import bpy
 221
 222
 223 def load_ply_mesh(filepath, ply_name):
 224     from bpy_extras.io_utils import unpack_face_list
 225     # from bpy_extras.image_utils import load_image  # UNUSED
 226
 227     obj_spec, obj, texture = read(filepath)
 228     if obj is None:
 229         print('Invalid file')
 230         return
 231
 232     uvindices = colindices = None
 233     colmultiply = None
 234
 235     # noindices = None # Ignore normals
 236
 237     for el in obj_spec.specs:
 238         if el.name == b'vertex':
 239             vindices_x, vindices_y, vindices_z = el.index(b'x'), el.index(b'y'), el.index(b'z')
 240             # noindices = (el.index('nx'), el.index('ny'), el.index('nz'))
 241             # if -1 in noindices: noindices = None
 242             uvindices = (el.index(b's'), el.index(b't'))
 243             if -1 in uvindices:
 244                 uvindices = None
 245             colindices = el.index(b'red'), el.index(b'green'), el.index(b'blue'), el.index(b'alpha')
 246             if -1 in colindices:
 247                 colindices = None
 248             else:  # if not a float assume uchar
 249                 colmultiply = [1.0 if el.properties[i].numeric_type in {'f', 'd'} else (1.0 / 255.0) for i in colindices]
 250
 251         elif el.name == b'face':
 252             findex = el.index(b'vertex_indices')
 253         elif el.name == b'tristrips':
 254             trindex = el.index(b'vertex_indices')
 255         elif el.name == b'edge':
 256             eindex1, eindex2 = el.index(b'vertex1'), el.index(b'vertex2')
 257
 258     mesh_faces = []
 259     mesh_uvs = []
 260     mesh_colors = []
 261
 262     def add_face(vertices, indices, uvindices, colindices):
 263         mesh_faces.append(indices)
 264         if uvindices:
 265             mesh_uvs.append([(vertices[index][uvindices[0]], vertices[index][uvindices[1]]) for index in indices])
 266         if colindices:
 267             mesh_colors.append([(vertices[index][colindices[0]] * colmultiply[0],
 268                                  vertices[index][colindices[1]] * colmultiply[1],
 269                                  vertices[index][colindices[2]] * colmultiply[2],
 270                                  vertices[index][colindices[3]] * colmultiply[3],
 271                                  ) for index in indices])
 272
 273     if uvindices or colindices:
 274         # If we have Cols or UVs then we need to check the face order.
 275         add_face_simple = add_face
 276
 277         # EVIL EEKADOODLE - face order annoyance.
 278         def add_face(vertices, indices, uvindices, colindices):
 279             if len(indices) == 4:
 280                 if indices[2] == 0 or indices[3] == 0:
 281                     indices = indices[2], indices[3], indices[0], indices[1]
 282             elif len(indices) == 3:
 283                 if indices[2] == 0:
 284                     indices = indices[1], indices[2], indices[0]
 285
 286             add_face_simple(vertices, indices, uvindices, colindices)
 287
 288     verts = obj[b'vertex']
 289
 290     if b'face' in obj:
 291         for f in obj[b'face']:
 292             ind = f[findex]
 293             len_ind = len(ind)
 294             if len_ind <= 4:
 295                 add_face(verts, ind, uvindices, colindices)
 296             else:
 297                 # Fan fill the face
 298                 for j in range(len_ind - 2):
 299                     add_face(verts, (ind[0], ind[j + 1], ind[j + 2]), uvindices, colindices)
 300
 301     if b'tristrips' in obj:
 302         for t in obj[b'tristrips']:
 303             ind = t[trindex]
 304             len_ind = len(ind)
 305             for j in range(len_ind - 2):
 306                 add_face(verts, (ind[j], ind[j + 1], ind[j + 2]), uvindices, colindices)
 307
 308     mesh = bpy.data.meshes.new(name=ply_name)
 309
 310     mesh.vertices.add(len(obj[b'vertex']))
 311
 312     mesh.vertices.foreach_set("co", [a for v in obj[b'vertex'] for a in (v[vindices_x], v[vindices_y], v[vindices_z])])
 313
 314     if b'edge' in obj:
 315         mesh.edges.add(len(obj[b'edge']))
 316         mesh.edges.foreach_set("vertices", [a for e in obj[b'edge'] for a in (e[eindex1], e[eindex2])])
 317
 318     if mesh_faces:
 319         mesh.tessfaces.add(len(mesh_faces))
 320         mesh.tessfaces.foreach_set("vertices_raw", unpack_face_list(mesh_faces))
 321
 322         if uvindices or colindices:
 323             if uvindices:
 324                 uvlay = mesh.tessface_uv_textures.new()
 325             if colindices:
 326                 vcol_lay = mesh.tessface_vertex_colors.new()
 327
 328             if uvindices:
 329                 for i, f in enumerate(uvlay.data):
 330                     ply_uv = mesh_uvs[i]
 331                     for j, uv in enumerate(f.uv):
 332                         uv[0], uv[1] = ply_uv[j]
 333
 334             if colindices:
 335                 for i, f in enumerate(vcol_lay.data):
 336                     # XXX, colors dont come in right, needs further investigation.
 337                     ply_col = mesh_colors[i]
 338                     if len(ply_col) == 4:
 339                         f_col = f.color1, f.color2, f.color3, f.color4
 340                     else:
 341                         f_col = f.color1, f.color2, f.color3
 342
 343                     for j, col in enumerate(f_col):
 344                         col[0] = ply_col[j][0]
 345                         col[1] = ply_col[j][1]
 346                         col[2] = ply_col[j][2]
 347                         col[3] = ply_col[j][3]
 348
 349     mesh.validate()
 350     mesh.update()
 351
 352     if texture and uvindices:
 353
 354         import os
 355         import sys
 356         from bpy_extras.image_utils import load_image
 357
 358         encoding = sys.getfilesystemencoding()
 359         encoded_texture = texture.decode(encoding=encoding)
 360         name = bpy.path.display_name_from_filepath(texture)
 361         image = load_image(encoded_texture, os.path.dirname(filepath), recursive=True, place_holder=True)
 362
 363         if image:
 364             texture = bpy.data.textures.new(name=name, type='IMAGE')
 365             texture.image = image
 366
 367             material = bpy.data.materials.new(name=name)
 368             material.use_shadeless = True
 369
 370             mtex = material.texture_slots.add()
 371             mtex.texture = texture
 372             mtex.texture_coords = 'UV'
 373             mtex.use_map_color_diffuse = True
 374
 375             mesh.materials.append(material)
 376             for face in mesh.uv_textures[0].data:
 377                 face.image = image
 378
 379     return mesh
 380
 381
 382 def load_ply(filepath):
 383     import time
 384
 385     t = time.time()
 386     ply_name = bpy.path.display_name_from_filepath(filepath)
 387
 388     mesh = load_ply_mesh(filepath, ply_name)
 389     if not mesh:
 390         return {'CANCELLED'}
 391
 392     scn = bpy.context.scene
 393
 394     obj = bpy.data.objects.new(ply_name, mesh)
 395     scn.objects.link(obj)
 396     scn.objects.active = obj
 397     obj.select = True
 398
 399     print('\nSuccessfully imported %r in %.3f sec' % (filepath, time.time() - t))
 400     return {'FINISHED'}
 401
 402
 403 def load(operator, context, filepath=""):
 404     return load_ply(filepath)