glTF exporter: fix / enhancement of texture images export

1 files changed, 63 insertions, 76 deletions

diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_gather_image.py b/io_scene_gltf2/blender/exp/gltf2_blender_gather_image.py
index b570e616..71913034 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_image.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_image.py
@@ -24,8 +24,11 @@ from io_scene_gltf2.blender.exp import gltf2_blender_search_node_tree
 from io_scene_gltf2.io.exp import gltf2_io_binary_data
 from io_scene_gltf2.io.exp import gltf2_io_image_data
 from io_scene_gltf2.io.com import gltf2_io_debug
+from io_scene_gltf2.blender.exp import gltf2_blender_image
+from io_scene_gltf2.blender.exp.gltf2_blender_gather_cache import cached
 
 
+@cached
 def gather_image(
         blender_shader_sockets_or_texture_slots: typing.Union[typing.Tuple[bpy.types.NodeSocket],
                                                               typing.Tuple[bpy.types.Texture]],
@@ -39,13 +42,11 @@ def gather_image(
         # The blender image has no data
         return None
 
-    mime_type = __gather_mime_type(uri.filepath if uri is not None else "")
-
     image = gltf2_io.Image(
         buffer_view=buffer_view,
         extensions=__gather_extensions(blender_shader_sockets_or_texture_slots, export_settings),
         extras=__gather_extras(blender_shader_sockets_or_texture_slots, export_settings),
-        mime_type=mime_type,
+        mime_type=__gather_mime_type(blender_shader_sockets_or_texture_slots, export_settings),
         name=__gather_name(blender_shader_sockets_or_texture_slots, export_settings),
         uri=uri
     )
@@ -64,7 +65,7 @@ def __gather_buffer_view(sockets_or_slots, export_settings):
         if image is None:
             return None
         return gltf2_io_binary_data.BinaryData(
-            data=image.to_image_data(__gather_mime_type()))
+            data=image.encode(__gather_mime_type(sockets_or_slots, export_settings)))
     return None
 
 
@@ -76,29 +77,43 @@ def __gather_extras(sockets_or_slots, export_settings):
     return None
 
 
-def __gather_mime_type(filepath=""):
-    extension_types = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg'}
-    default_extension = extension_types['.png']
-
-    matches = re.findall(r'\.\w+$', filepath)
-    extension = matches[0] if len(matches) > 0 else default_extension
-    return extension_types[extension] if extension.lower() in extension_types.keys() else default_extension
+def __gather_mime_type(sockets_or_slots, export_settings):
+    if export_settings["gltf_image_format"] == "NAME":
+        image_name = __get_texname_from_slot(sockets_or_slots, export_settings)
+        _, extension = os.path.splitext(image_name)
+        if extension in [".jpeg", ".jpg", ".png"]:
+            return {
+                ".jpeg": "image/jpeg",
+                ".jpg": "image/jpeg",
+                ".png": "image/png",
+            }[extension]
+        return "image/png"
+
+    elif export_settings["gltf_image_format"] == "JPEG":
+        return "image/jpeg"
+    else:
+        return "image/png"
 
 
 def __gather_name(sockets_or_slots, export_settings):
-    if __is_socket(sockets_or_slots):
-        node = __get_tex_from_socket(sockets_or_slots[0])
-        if node is not None:
-            return node.shader_node.image.name
-    elif isinstance(sockets_or_slots[0], bpy.types.MaterialTextureSlot):
-        return sockets_or_slots[0].name
-    return None
+    image_name = __get_texname_from_slot(sockets_or_slots, export_settings)
+
+    name, extension = os.path.splitext(image_name)
+    if extension in [".jpeg", ".jpg", ".png"]:
+        return name
+    return image_name
 
 
 def __gather_uri(sockets_or_slots, export_settings):
     if export_settings[gltf2_blender_export_keys.FORMAT] == 'GLTF_SEPARATE':
         # as usual we just store the data in place instead of already resolving the references
-        return __get_image_data(sockets_or_slots, export_settings)
+        mime_type = __gather_mime_type(sockets_or_slots, export_settings)
+        return gltf2_io_image_data.ImageData(
+            data=__get_image_data(sockets_or_slots, export_settings).encode(mime_type=mime_type),
+            mime_type=mime_type,
+            name=__gather_name(sockets_or_slots, export_settings)
+        )
+
     return None
 
 
@@ -110,7 +125,7 @@ def __is_slot(sockets_or_slots):
     return isinstance(sockets_or_slots[0], bpy.types.MaterialTextureSlot)
 
 
-def __get_image_data(sockets_or_slots, export_settings):
+def __get_image_data(sockets_or_slots, export_settings) -> gltf2_blender_image.ExportImage:
     # For shared ressources, such as images, we just store the portion of data that is needed in the glTF property
     # in a helper class. During generation of the glTF in the exporter these will then be combined to actual binary
     # ressources.
@@ -128,8 +143,8 @@ def __get_image_data(sockets_or_slots, export_settings):
         return channels
 
     if __is_socket(sockets_or_slots):
-        results = [__get_tex_from_socket(socket) for socket in sockets_or_slots]
-        image = None
+        results = [__get_tex_from_socket(socket, export_settings) for socket in sockets_or_slots]
+        composed_image = None
         for result, socket in zip(results, sockets_or_slots):
             if result.shader_node.image.channels == 0:
                 gltf2_io_debug.print_console("WARNING",
@@ -138,9 +153,7 @@ def __get_image_data(sockets_or_slots, export_settings):
                 continue
 
             # rudimentarily try follow the node tree to find the correct image data.
-            source_channel = None
-            target_channel = None
-            source_channels_length = None
+            source_channel = 0
             for elem in result.path:
                 if isinstance(elem.from_node, bpy.types.ShaderNodeSeparateRGB):
                     source_channel = {
@@ -149,68 +162,30 @@ def __get_image_data(sockets_or_slots, export_settings):
                         'B': 2
                     }[elem.from_socket.name]
 
-            if source_channel is not None:
-                pixels = [split_pixels_by_channels(result.shader_node.image, export_settings)[source_channel]]
-                target_channel = source_channel
-                source_channel = 0
-                source_channels_length = 1
-            else:
-                pixels = split_pixels_by_channels(result.shader_node.image, export_settings)
-                target_channel = 0
-                source_channel = 0
-                source_channels_length = len(pixels)
+            image = gltf2_blender_image.ExportImage.from_blender_image(result.shader_node.image)
+
+            if composed_image is None:
+                composed_image = gltf2_blender_image.ExportImage.white_image(image.width, image.height)
 
             # Change target channel for metallic and roughness.
             if elem.to_socket.name == 'Metallic':
-                target_channel = 2
-                source_channels_length = 1
+                composed_image[2] = image[source_channel]
             elif elem.to_socket.name == 'Roughness':
-                target_channel = 1
-                source_channels_length = 1
-
-            file_name = os.path.splitext(result.shader_node.image.name)[0]
-            if result.shader_node.image.packed_file is None:
-                file_path = result.shader_node.image.filepath
+                composed_image[1] = image[source_channel]
             else:
-                # empty path for packed textures, because they are converted to png anyway
-                file_path = ""
-
-            image_data = gltf2_io_image_data.ImageData(
-                file_name,
-                file_path,
-                result.shader_node.image.size[0],
-                result.shader_node.image.size[1],
-                source_channel,
-                target_channel,
-                source_channels_length,
-                pixels)
-
-            if image is None:
-                image = image_data
-            else:
-                image.add_to_image(target_channel, image_data)
+                composed_image.update(image)
+
+        return composed_image
 
-        return image
     elif __is_slot(sockets_or_slots):
         texture = __get_tex_from_slot(sockets_or_slots[0])
-        pixels = split_pixels_by_channels(texture.image, export_settings)
-
-        image_data = gltf2_io_image_data.ImageData(
-            texture.name,
-            texture.image.filepath,
-            texture.image.size[0],
-            texture.image.size[1],
-            0,
-            0,
-            len(pixels),
-            pixels)
-        return image_data
+        image = gltf2_blender_image.ExportImage.from_blender_image(texture.image)
+        return image
     else:
-        # Texture slots
         raise NotImplementedError()
 
-
-def __get_tex_from_socket(blender_shader_socket: bpy.types.NodeSocket):
+@cached
+def __get_tex_from_socket(blender_shader_socket: bpy.types.NodeSocket, export_settings):
     result = gltf2_blender_search_node_tree.from_socket(
         blender_shader_socket,
         gltf2_blender_search_node_tree.FilterByType(bpy.types.ShaderNodeTexImage))
@@ -222,3 +197,15 @@ def __get_tex_from_socket(blender_shader_socket: bpy.types.NodeSocket):
 def __get_tex_from_slot(blender_texture_slot):
     return blender_texture_slot.texture
 
+
+@cached
+def __get_texname_from_slot(sockets_or_slots, export_settings):
+    if __is_socket(sockets_or_slots):
+        node = __get_tex_from_socket(sockets_or_slots[0], export_settings)
+        if node is None:
+            return None
+        return node.shader_node.image.name
+
+    elif isinstance(sockets_or_slots[0], bpy.types.MaterialTextureSlot):
+        return sockets_or_slots[0].name
+