glTF: better occlusion map management

At import, create a node group  where you can link an image texture
At export, use this node group to create correctly the material
TODO: doc update

4 files changed, 65 insertions, 20 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 dbee622e..3e905a7d 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_image.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_image.py
@@ -170,10 +170,12 @@ def __get_image_data(sockets_or_slots, export_settings) -> gltf2_blender_image.E
                 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':
+            if socket.name == 'Metallic':
                 composed_image[2] = image[source_channel]
-            elif elem.to_socket.name == 'Roughness':
+            elif socket.name == 'Roughness':
                 composed_image[1] = image[source_channel]
+            elif socket.name == 'Occlusion' and len(sockets_or_slots) > 2:
+                composed_image[0] = image[source_channel]
             else:
                 composed_image.update(image)
 
diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_gather_materials.py b/io_scene_gltf2/blender/exp/gltf2_blender_gather_materials.py
index 5fff246a..1242c08d 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_materials.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_materials.py
@@ -20,6 +20,7 @@ from io_scene_gltf2.io.com.gltf2_io_extensions import Extension
 from io_scene_gltf2.blender.exp import gltf2_blender_gather_texture_info, gltf2_blender_export_keys
 from io_scene_gltf2.blender.exp import gltf2_blender_gather_material_normal_texture_info_class
 from io_scene_gltf2.blender.exp import gltf2_blender_gather_material_occlusion_texture_info_class
+from io_scene_gltf2.blender.exp import gltf2_blender_search_node_tree
 
 from io_scene_gltf2.blender.exp import gltf2_blender_gather_materials_pbr_metallic_roughness
 from io_scene_gltf2.blender.exp import gltf2_blender_generate_extras
@@ -38,6 +39,8 @@ def gather_material(blender_material, mesh_double_sided, export_settings):
     if not __filter_material(blender_material, export_settings):
         return None
 
+    orm_texture = __gather_orm_texture(blender_material, export_settings)
+
     material = gltf2_io.Material(
         alpha_cutoff=__gather_alpha_cutoff(blender_material, export_settings),
         alpha_mode=__gather_alpha_mode(blender_material, export_settings),
@@ -48,8 +51,8 @@ def gather_material(blender_material, mesh_double_sided, export_settings):
         extras=__gather_extras(blender_material, export_settings),
         name=__gather_name(blender_material, export_settings),
         normal_texture=__gather_normal_texture(blender_material, export_settings),
-        occlusion_texture=__gather_occlusion_texture(blender_material, export_settings),
-        pbr_metallic_roughness=__gather_pbr_metallic_roughness(blender_material, export_settings)
+        occlusion_texture=__gather_occlusion_texture(blender_material, orm_texture, export_settings),
+        pbr_metallic_roughness=__gather_pbr_metallic_roughness(blender_material, orm_texture, export_settings)
     )
 
     return material
@@ -150,7 +153,33 @@ def __gather_normal_texture(blender_material, export_settings):
         export_settings)
 
 
-def __gather_occlusion_texture(blender_material, export_settings):
+def __gather_orm_texture(blender_material, export_settings):
+    # Check for the presence of Occlusion, Roughness, Metallic sharing a single image.
+    # If not fully shared, return None, so the images will be cached and processed separately.
+
+    occlusion = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "Occlusion")
+    if occlusion is None or not __has_image_node_from_socket(occlusion):
+        occlusion = gltf2_blender_get.get_socket_or_texture_slot_old(blender_material, "Occlusion")
+        if occlusion is None or not __has_image_node_from_socket(occlusion):
+            return None
+
+    metallic_socket = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "Metallic")
+    roughness_socket = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "Roughness")
+    if metallic_socket is None or roughness_socket is None\
+            or not __has_image_node_from_socket(metallic_socket)\
+            or not __has_image_node_from_socket(roughness_socket):
+        metallic_roughness = gltf2_blender_get.get_socket_or_texture_slot_old(blender_material, "MetallicRoughness")
+        if metallic_roughness is None or not __has_image_node_from_socket(metallic_roughness):
+            return None
+        return (occlusion, metallic_roughness, metallic_roughness)
+
+    return (occlusion, roughness_socket, metallic_socket)
+
+def __gather_occlusion_texture(blender_material, orm_texture, export_settings):
+    if orm_texture is not None:
+        return gltf2_blender_gather_material_occlusion_texture_info_class.gather_material_occlusion_texture_info_class(
+            orm_texture,
+            export_settings)
     occlusion = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "Occlusion")
     if occlusion is None:
         occlusion = gltf2_blender_get.get_socket_or_texture_slot_old(blender_material, "Occlusion")
@@ -159,8 +188,17 @@ def __gather_occlusion_texture(blender_material, export_settings):
         export_settings)
 
 
-def __gather_pbr_metallic_roughness(blender_material, export_settings):
+def __gather_pbr_metallic_roughness(blender_material, orm_texture, export_settings):
     return gltf2_blender_gather_materials_pbr_metallic_roughness.gather_material_pbr_metallic_roughness(
         blender_material,
+        orm_texture,
         export_settings)
 
+def __has_image_node_from_socket(socket):
+    result = gltf2_blender_search_node_tree.from_socket(
+        socket,
+        gltf2_blender_search_node_tree.FilterByType(bpy.types.ShaderNodeTexImage))
+    if not result:
+        return False
+    return True
+
diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_gather_materials_pbr_metallic_roughness.py b/io_scene_gltf2/blender/exp/gltf2_blender_gather_materials_pbr_metallic_roughness.py
index 67e7fa12..ffbf9914 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_materials_pbr_metallic_roughness.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_materials_pbr_metallic_roughness.py
@@ -23,7 +23,7 @@ from io_scene_gltf2.io.com.gltf2_io_debug import print_console
 
 
 @cached
-def gather_material_pbr_metallic_roughness(blender_material, export_settings):
+def gather_material_pbr_metallic_roughness(blender_material, orm_texture, export_settings):
     if not __filter_pbr_material(blender_material, export_settings):
         return None
 
@@ -33,7 +33,7 @@ def gather_material_pbr_metallic_roughness(blender_material, export_settings):
         extensions=__gather_extensions(blender_material, export_settings),
         extras=__gather_extras(blender_material, export_settings),
         metallic_factor=__gather_metallic_factor(blender_material, export_settings),
-        metallic_roughness_texture=__gather_metallic_roughness_texture(blender_material, export_settings),
+        metallic_roughness_texture=__gather_metallic_roughness_texture(blender_material, orm_texture, export_settings),
         roughness_factor=__gather_roughness_factor(blender_material, export_settings)
     )
 
@@ -123,17 +123,20 @@ def __gather_metallic_factor(blender_material, export_settings):
     return None
 
 
-def __gather_metallic_roughness_texture(blender_material, export_settings):
-    metallic_socket = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "Metallic")
-    roughness_socket = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "Roughness")
-
-    if metallic_socket is None and roughness_socket is None:
-        metallic_roughness = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "MetallicRoughness")
-        if metallic_roughness is None:
-            metallic_roughness = gltf2_blender_get.get_socket_or_texture_slot_old(blender_material, "MetallicRoughness")
-        texture_input = (metallic_roughness,)
+def __gather_metallic_roughness_texture(blender_material, orm_texture, export_settings):
+    if orm_texture is not None:
+        texture_input = orm_texture
     else:
-        texture_input = (metallic_socket, roughness_socket)
+        metallic_socket = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "Metallic")
+        roughness_socket = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "Roughness")
+
+        if metallic_socket is None and roughness_socket is None:
+            metallic_roughness = gltf2_blender_get.get_socket_or_texture_slot(blender_material, "MetallicRoughness")
+            if metallic_roughness is None:
+                metallic_roughness = gltf2_blender_get.get_socket_or_texture_slot_old(blender_material, "MetallicRoughness")
+            texture_input = (metallic_roughness,)
+        else:
+            texture_input = (metallic_socket, roughness_socket)
 
     return gltf2_blender_gather_texture_info.gather_texture_info(texture_input, export_settings)
 
diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_get.py b/io_scene_gltf2/blender/exp/gltf2_blender_get.py
index 27d8a341..fef2e684 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_get.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_get.py
@@ -15,6 +15,7 @@
 import bpy
 from mathutils import Vector, Matrix
 
+from ..com.gltf2_blender_material_helpers import get_gltf_node_name
 from ...blender.com.gltf2_blender_conversion import texture_transform_blender_to_gltf
 from io_scene_gltf2.io.com import gltf2_io_debug
 
@@ -73,16 +74,17 @@ def get_socket_or_texture_slot(blender_material: bpy.types.Material, name: str):
 
 def get_socket_or_texture_slot_old(blender_material: bpy.types.Material, name: str):
     """
-    For a given material input name, retrieve the corresponding node tree socket in the special glTF Metallic Roughness nodes (which might be deprecated?).
+    For a given material input name, retrieve the corresponding node tree socket in the special glTF node group.
 
     :param blender_material: a blender material for which to get the socket/slot
     :param name: the name of the socket/slot
     :return: either a blender NodeSocket, if the material is a node tree or a blender Texture otherwise
     """
+    gltf_node_group_name = get_gltf_node_name().lower()
     if blender_material.node_tree and blender_material.use_nodes:
         nodes = [n for n in blender_material.node_tree.nodes if \
             isinstance(n, bpy.types.ShaderNodeGroup) and \
-            n.node_tree.name.startswith('glTF Metallic Roughness')]
+            (n.node_tree.name.startswith('glTF Metallic Roughness') or n.node_tree.name.lower() == gltf_node_group_name)]
         inputs = sum([[input for input in node.inputs if input.name == name] for node in nodes], [])
         if inputs:
             return inputs[0]