glTF exporter: roundtrip all texture wrap modes

4 files changed, 126 insertions, 60 deletions

diff --git a/io_scene_gltf2/__init__.py b/io_scene_gltf2/__init__.py
index d8bee58c..127ea6e5 100755
--- a/io_scene_gltf2/__init__.py
+++ b/io_scene_gltf2/__init__.py
@@ -15,7 +15,7 @@
 bl_info = {
     'name': 'glTF 2.0 format',
     'author': 'Julien Duroure, Scurest, Norbert Nopper, Urs Hanselmann, Moritz Becher, Benjamin Schmithüsen, Jim Eckerlein, and many external contributors',
-    "version": (1, 5, 9),
+    "version": (1, 5, 10),
     'blender': (2, 91, 0),
     'location': 'File > Import-Export',
     'description': 'Import-Export as glTF 2.0',
diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_gather_sampler.py b/io_scene_gltf2/blender/exp/gltf2_blender_gather_sampler.py
index abbd7e94..40dfdb16 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_sampler.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_sampler.py
@@ -17,18 +17,25 @@ from io_scene_gltf2.io.com import gltf2_io
 from io_scene_gltf2.blender.exp.gltf2_blender_gather_cache import cached
 from io_scene_gltf2.io.exp.gltf2_io_user_extensions import export_user_extensions
 from io_scene_gltf2.io.com.gltf2_io_constants import TextureFilter, TextureWrap
+from io_scene_gltf2.blender.exp.gltf2_blender_get import (
+    previous_node,
+    previous_socket,
+    get_const_from_socket,
+)
 
 
 @cached
 def gather_sampler(blender_shader_node: bpy.types.Node, export_settings):
+    wrap_s, wrap_t = __gather_wrap(blender_shader_node, export_settings)
+
     sampler = gltf2_io.Sampler(
         extensions=__gather_extensions(blender_shader_node, export_settings),
         extras=__gather_extras(blender_shader_node, export_settings),
         mag_filter=__gather_mag_filter(blender_shader_node, export_settings),
         min_filter=__gather_min_filter(blender_shader_node, export_settings),
         name=__gather_name(blender_shader_node, export_settings),
-        wrap_s=__gather_wrap_s(blender_shader_node, export_settings),
-        wrap_t=__gather_wrap_t(blender_shader_node, export_settings)
+        wrap_s=wrap_s,
+        wrap_t=wrap_t,
     )
 
     export_user_extensions('gather_sampler_hook', export_settings, sampler, blender_shader_node)
@@ -83,13 +90,83 @@ def __gather_name(blender_shader_node, export_settings):
     return None
 
 
-def __gather_wrap_s(blender_shader_node, export_settings):
-    if blender_shader_node.extension == 'EXTEND':
-        return TextureWrap.ClampToEdge
-    return None
-
-
-def __gather_wrap_t(blender_shader_node, export_settings):
+def __gather_wrap(blender_shader_node, export_settings):
+    # First gather from the Texture node
     if blender_shader_node.extension == 'EXTEND':
-        return TextureWrap.ClampToEdge
-    return None
+        wrap_s = TextureWrap.ClampToEdge
+    elif blender_shader_node.extension == 'CLIP':
+        # Not possible in glTF, but ClampToEdge is closest
+        wrap_s = TextureWrap.ClampToEdge
+    else:
+        wrap_s = TextureWrap.Repeat
+    wrap_t = wrap_s
+
+    # Take manual wrapping into account
+    result = detect_manual_uv_wrapping(blender_shader_node)
+    if result:
+        if result['wrap_s'] is not None: wrap_s = result['wrap_s']
+        if result['wrap_t'] is not None: wrap_t = result['wrap_t']
+
+    # Omit if both are repeat
+    if (wrap_s, wrap_t) == (TextureWrap.Repeat, TextureWrap.Repeat):
+        wrap_s, wrap_t = None, None
+
+    return wrap_s, wrap_t
+
+
+def detect_manual_uv_wrapping(blender_shader_node):
+    # Detects UV wrapping done using math nodes. This is for emulating wrap
+    # modes Blender doesn't support. It looks like
+    #
+    #     next_socket => [Sep XYZ] => [Wrap S] => [Comb XYZ] => blender_shader_node
+    #                              => [Wrap T] =>
+    #
+    # The [Wrap _] blocks are either math nodes (eg. PINGPONG for mirrored
+    # repeat), or can be omitted.
+    #
+    # Returns None if not detected. Otherwise a dict containing the wrap
+    # mode in each direction (or None), and next_socket.
+    result = {}
+
+    comb = previous_node(blender_shader_node.inputs['Vector'])
+    if comb is None or comb.type != 'COMBXYZ': return None
+
+    for soc in ['X', 'Y']:
+        node = previous_node(comb.inputs[soc])
+        if node is None: return None
+
+        if node.type == 'SEPXYZ':
+            # Passed through without change
+            wrap = None
+            prev_socket = previous_socket(comb.inputs[soc])
+        elif node.type == 'MATH':
+            # Math node applies a manual wrap
+            if (node.operation == 'PINGPONG' and
+                    get_const_from_socket(node.inputs[1], kind='VALUE') == 1.0):  # scale = 1
+                wrap = TextureWrap.MirroredRepeat
+            elif (node.operation == 'WRAP' and
+                    get_const_from_socket(node.inputs[1], kind='VALUE') == 0.0 and  # min = 0
+                    get_const_from_socket(node.inputs[2], kind='VALUE') == 1.0):    # max = 1
+                wrap = TextureWrap.Repeat
+            else:
+                return None
+
+            prev_socket = previous_socket(node.inputs[0])
+        else:
+            return None
+
+        if prev_socket is None: return None
+        prev_node = prev_socket.node
+        if prev_node.type != 'SEPXYZ': return None
+        # Make sure X goes to X, etc.
+        if prev_socket.name != soc: return None
+        # Make sure both attach to the same SeparateXYZ node
+        if soc == 'X':
+            sep = prev_node
+        else:
+            if sep != prev_node: return None
+
+        result['wrap_s' if soc == 'X' else 'wrap_t'] = wrap
+
+    result['next_socket'] = sep.inputs[0]
+    return result
diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_gather_texture_info.py b/io_scene_gltf2/blender/exp/gltf2_blender_gather_texture_info.py
index 30975a3f..6c4acb82 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_texture_info.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_texture_info.py
@@ -19,6 +19,8 @@ from io_scene_gltf2.io.com import gltf2_io
 from io_scene_gltf2.blender.exp import gltf2_blender_gather_texture
 from io_scene_gltf2.blender.exp import gltf2_blender_search_node_tree
 from io_scene_gltf2.blender.exp import gltf2_blender_get
+from io_scene_gltf2.blender.exp.gltf2_blender_get import previous_node
+from io_scene_gltf2.blender.exp.gltf2_blender_gather_sampler import detect_manual_uv_wrapping
 from io_scene_gltf2.io.com.gltf2_io_extensions import Extension
 from io_scene_gltf2.io.exp.gltf2_io_user_extensions import export_user_extensions
 
@@ -47,11 +49,13 @@ def __gather_texture_info_helper(
     if not __filter_texture_info(primary_socket, blender_shader_sockets, export_settings):
         return None
 
+    tex_transform, tex_coord = __gather_texture_transform_and_tex_coord(primary_socket, export_settings)
+
     fields = {
-        'extensions': __gather_extensions(primary_socket, export_settings),
+        'extensions': __gather_extensions(tex_transform, export_settings),
         'extras': __gather_extras(blender_shader_sockets, export_settings),
         'index': __gather_index(blender_shader_sockets, export_settings),
-        'tex_coord': __gather_tex_coord(primary_socket, export_settings),
+        'tex_coord': tex_coord,
     }
 
     if kind == 'DEFAULT':
@@ -89,17 +93,9 @@ def __filter_texture_info(primary_socket, blender_shader_sockets, export_setting
     return True
 
 
-def __gather_extensions(primary_socket, export_settings):
-    if not hasattr(primary_socket, 'links'):
-        return None
-
-    texture_node = __get_tex_from_socket(primary_socket).shader_node
-    if texture_node is None:
-        return None
-    texture_transform = gltf2_blender_get.get_texture_transform_from_texture_node(texture_node)
+def __gather_extensions(texture_transform, export_settings):
     if texture_transform is None:
         return None
-
     extension = Extension("KHR_texture_transform", texture_transform)
     return {"KHR_texture_transform": extension}
 
@@ -144,33 +140,37 @@ def __gather_index(blender_shader_sockets, export_settings):
     return gltf2_blender_gather_texture.gather_texture(blender_shader_sockets, export_settings)
 
 
-def __gather_tex_coord(primary_socket, export_settings):
+def __gather_texture_transform_and_tex_coord(primary_socket, export_settings):
+    # We're expecting
+    #
+    #     [UV Map] => [Mapping] => [UV Wrapping] => [Texture Node] => ... => primary_socket
+    #
+    # The [UV Wrapping] is for wrap modes like MIRROR that use nodes,
+    # [Mapping] is for KHR_texture_transform, and [UV Map] is for texCoord.
     blender_shader_node = __get_tex_from_socket(primary_socket).shader_node
-    if len(blender_shader_node.inputs['Vector'].links) == 0:
-        return 0
-
-    input_node = blender_shader_node.inputs['Vector'].links[0].from_node
-
-    if isinstance(input_node, bpy.types.ShaderNodeMapping):
-
-        if len(input_node.inputs['Vector'].links) == 0:
-            return 0
-
-        input_node = input_node.inputs['Vector'].links[0].from_node
-
-    if not isinstance(input_node, bpy.types.ShaderNodeUVMap):
-        return 0
-
-    if input_node.uv_map == '':
-        return 0
-
-    # Try to gather map index.
-    for blender_mesh in bpy.data.meshes:
-        texCoordIndex = blender_mesh.uv_layers.find(input_node.uv_map)
-        if texCoordIndex >= 0:
-            return texCoordIndex
 
-    return 0
+    # Skip over UV wrapping stuff (it goes in the sampler)
+    result = detect_manual_uv_wrapping(blender_shader_node)
+    if result:
+        node = previous_node(result['next_socket'])
+    else:
+        node = previous_node(blender_shader_node.inputs['Vector'])
+
+    texture_transform = None
+    if node and node.type == 'MAPPING':
+        texture_transform = gltf2_blender_get.get_texture_transform_from_mapping_node(node)
+        node = previous_node(node.inputs['Vector'])
+
+    texcoord_idx = 0
+    if node and node.type == 'UVMAP' and node.uv_map:
+        # Try to gather map index.
+        for blender_mesh in bpy.data.meshes:
+            i = blender_mesh.uv_layers.find(node.uv_map)
+            if i >= 0:
+                texcoord_idx = i
+                break
+
+    return texture_transform, texcoord_idx or None
 
 
 def __get_tex_from_socket(socket):
diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_get.py b/io_scene_gltf2/blender/exp/gltf2_blender_get.py
index ee63aa7e..b7a52114 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_get.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_get.py
@@ -132,18 +132,7 @@ def find_shader_image_from_shader_socket(shader_socket, max_hops=10):
     return None
 
 
-def get_texture_transform_from_texture_node(texture_node):
-    if not isinstance(texture_node, bpy.types.ShaderNodeTexImage):
-        return None
-
-    mapping_socket = texture_node.inputs["Vector"]
-    if len(mapping_socket.links) == 0:
-        return None
-
-    mapping_node = mapping_socket.links[0].from_node
-    if not isinstance(mapping_node, bpy.types.ShaderNodeMapping):
-        return None
-
+def get_texture_transform_from_mapping_node(mapping_node):
     if mapping_node.vector_type not in ["TEXTURE", "POINT", "VECTOR"]:
         gltf2_io_debug.print_console("WARNING",
             "Skipping exporting texture transform because it had type " +