glTF importer/exporter: Draco decoder + encoder fixes

We can now read Draco compressed files.
This also fix exporting vertex color Draco compressed files. Fix #T75550

2 files changed, 149 insertions, 0 deletions

diff --git a/io_scene_gltf2/blender/imp/gltf2_blender_mesh.py b/io_scene_gltf2/blender/imp/gltf2_blender_mesh.py
index a6896491..328bc1a8 100755
--- a/io_scene_gltf2/blender/imp/gltf2_blender_mesh.py
+++ b/io_scene_gltf2/blender/imp/gltf2_blender_mesh.py
@@ -19,6 +19,8 @@ import numpy as np
 from ...io.imp.gltf2_io_binary import BinaryData
 from ..com.gltf2_blender_extras import set_extras
 from .gltf2_blender_material import BlenderMaterial
+from ...io.com.gltf2_io_debug import print_console
+from .gltf2_io_draco_compression_extension import decode_primitive
 
 
 class BlenderMesh():
@@ -134,6 +136,10 @@ def do_primitives(gltf, mesh_idx, skin_idx, mesh, ob):
 
         vert_index_base = len(vert_locs)
 
+        if prim.extensions is not None and 'KHR_draco_mesh_compression' in prim.extensions:
+            print_console('INFO', 'Draco Decoder: Decode primitive {}'.format(pymesh.name or '[unnamed]'))
+            decode_primitive(gltf, prim)
+
         if prim.indices is not None:
             indices = BinaryData.decode_accessor(gltf, prim.indices)
             indices = indices.reshape(len(indices))
diff --git a/io_scene_gltf2/blender/imp/gltf2_io_draco_compression_extension.py b/io_scene_gltf2/blender/imp/gltf2_io_draco_compression_extension.py
new file mode 100644
index 00000000..5a93673f
--- /dev/null
+++ b/io_scene_gltf2/blender/imp/gltf2_io_draco_compression_extension.py
@@ -0,0 +1,143 @@
+# Copyright 2018-2019 The glTF-Blender-IO authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ctypes import *
+
+from io_scene_gltf2.io.com.gltf2_io import BufferView
+from io_scene_gltf2.io.imp.gltf2_io_binary import BinaryData
+from ...io.com.gltf2_io_debug import print_console
+from io_scene_gltf2.io.com.gltf2_io_draco_compression_extension import dll_path
+
+
+def decode_primitive(gltf, prim):
+    """
+    Handles draco compression.
+    Moves decoded data into new buffers and buffer views held by the accessors of the given primitive.
+    """
+
+    # Load DLL and setup function signatures.
+    dll = cdll.LoadLibrary(str(dll_path().resolve()))
+
+    dll.decoderCreate.restype = c_void_p
+    dll.decoderCreate.argtypes = []
+
+    dll.decoderRelease.restype = None
+    dll.decoderRelease.argtypes = [c_void_p]
+
+    dll.decoderDecode.restype = c_bool
+    dll.decoderDecode.argtypes = [c_void_p, c_void_p, c_size_t]
+
+    dll.decoderReadAttribute.restype = c_bool
+    dll.decoderReadAttribute.argtypes = [c_void_p, c_uint32, c_size_t, c_char_p]
+
+    dll.decoderGetVertexCount.restype = c_uint32
+    dll.decoderGetVertexCount.argtypes = [c_void_p]
+
+    dll.decoderGetIndexCount.restype = c_uint32
+    dll.decoderGetIndexCount.argtypes = [c_void_p]
+
+    dll.decoderAttributeIsNormalized.restype = c_bool
+    dll.decoderAttributeIsNormalized.argtypes = [c_void_p, c_uint32]
+
+    dll.decoderGetAttributeByteLength.restype = c_size_t
+    dll.decoderGetAttributeByteLength.argtypes = [c_void_p, c_uint32]
+
+    dll.decoderCopyAttribute.restype = None
+    dll.decoderCopyAttribute.argtypes = [c_void_p, c_uint32, c_void_p]
+
+    dll.decoderReadIndices.restype = c_bool
+    dll.decoderReadIndices.argtypes = [c_void_p, c_size_t]
+
+    dll.decoderGetIndicesByteLength.restype = c_size_t
+    dll.decoderGetIndicesByteLength.argtypes = [c_void_p]
+
+    dll.decoderCopyIndices.restype = None
+    dll.decoderCopyIndices.argtypes = [c_void_p, c_void_p]
+    
+    decoder = dll.decoderCreate()
+    extension = prim.extensions['KHR_draco_mesh_compression']
+
+    name = prim.name if hasattr(prim, 'name') else '[unnamed]'
+
+    # Create Draco decoder.
+    draco_buffer = bytes(BinaryData.get_buffer_view(gltf, extension['bufferView']))
+    if not dll.decoderDecode(decoder, draco_buffer, len(draco_buffer)):
+        print_console('ERROR', 'Draco Decoder: Unable to decode. Skipping primitive {}.'.format(name))
+        return
+
+    # Choose a buffer index which does not yet exist, skipping over existing glTF buffers yet to be loaded
+    # and buffers which were generated and did not exist in the initial glTF file, like this decoder does.
+    base_buffer_idx = len(gltf.data.buffers)
+    for existing_buffer_idx in gltf.buffers:
+        if base_buffer_idx <= existing_buffer_idx:
+            base_buffer_idx = existing_buffer_idx + 1
+    
+    # Read indices.
+    index_accessor = gltf.data.accessors[prim.indices]
+    if dll.decoderGetIndexCount(decoder) != index_accessor.count:
+        print_console('WARNING', 'Draco Decoder: Index count of accessor and decoded index count does not match. Updating accessor.')
+        index_accessor.count = dll.decoderGetIndexCount(decoder)
+    if not dll.decoderReadIndices(decoder, index_accessor.component_type):
+        print_console('ERROR', 'Draco Decoder: Unable to decode indices. Skipping primitive {}.'.format(name))
+        return
+    
+    indices_byte_length = dll.decoderGetIndicesByteLength(decoder)
+    decoded_data = bytes(indices_byte_length)
+    dll.decoderCopyIndices(decoder, decoded_data)
+
+    # Generate a new buffer holding the decoded indices.
+    gltf.buffers[base_buffer_idx] = decoded_data
+
+    # Create a buffer view referencing the new buffer.
+    gltf.data.buffer_views.append(BufferView.from_dict({
+        'buffer': base_buffer_idx,
+        'byteLength': indices_byte_length
+    }))
+
+    # Update accessor to point to the new buffer view.
+    index_accessor.buffer_view = len(gltf.data.buffer_views) - 1
+    
+    # Read each attribute.
+    for attr_idx, attr in enumerate(extension['attributes']):
+        dracoId = extension['attributes'][attr]
+        if attr not in prim.attributes:
+            print_console('ERROR', 'Draco Decoder: Draco attribute {} not in primitive attributes. Skipping primitive {}.'.format(attr, name))
+            return
+        
+        accessor = gltf.data.accessors[prim.attributes[attr]]
+        if dll.decoderGetVertexCount(decoder) != accessor.count:
+            print_console('WARNING', 'Draco Decoder: Vertex count of accessor and decoded vertex count does not match for attribute {}. Updating accessor.'.format(attr, name))
+            accessor.count = dll.decoderGetVertexCount(decoder)
+        if not dll.decoderReadAttribute(decoder, dracoId, accessor.component_type, accessor.type.encode()):
+            print_console('ERROR', 'Draco Decoder: Could not decode attribute {}. Skipping primitive {}.'.format(attr, name))
+            return
+        
+        byte_length = dll.decoderGetAttributeByteLength(decoder, dracoId)
+        decoded_data = bytes(byte_length)
+        dll.decoderCopyAttribute(decoder, dracoId, decoded_data)
+
+        # Generate a new buffer holding the decoded vertex data.
+        buffer_idx = base_buffer_idx + 1 + attr_idx
+        gltf.buffers[buffer_idx] = decoded_data
+        
+        # Create a buffer view referencing the new buffer.
+        gltf.data.buffer_views.append(BufferView.from_dict({
+            'buffer': buffer_idx,
+            'byteLength': byte_length
+        }))
+
+        # Update accessor to point to the new buffer view.
+        accessor.buffer_view = len(gltf.data.buffer_views) - 1
+
+    dll.decoderRelease(decoder)