implement the new extractor and options for exporting texture as sing…

…le output

client/ayon_substancepainter/plugins/create/create_textures.py

@@ -53,7 +53,9 @@ def create(self, product_name, instance_data, pre_create_data):
             "exportPadding",
             "exportDilationDistance",
             "useCustomExportPreset",
-            "exportChannel"
+            "exportChannel",
+            "exportTextureSets",
+            "exportTextureSetsAsOneOutput"
         ]:
             if key in pre_create_data:
                 creator_attributes[key] = pre_create_data[key]
@@ -152,6 +154,11 @@ def get_instance_attr_defs(self):
                     label="Review",
                     tooltip="Mark as reviewable",
                     default=True),
+            BoolDef("exportTextureSetsAsOneOutput",
+                    label="Export Texture Sets As One Texture Output",
+                    tooltip="Export multiple texture set(s) "
+                            "as one Texture Output",
+                    default=False),
             EnumDef("exportTextureSets",
                     items=export_texture_set_enum,
                     multiselection=True,

client/ayon_substancepainter/plugins/publish/collect_textureset_images.py

@@ -4,6 +4,7 @@
 import pyblish.api
 import ayon_api
 
+from collections import defaultdict
 import substance_painter.textureset
 from ayon_core.pipeline import publish
 from ayon_substancepainter.api.lib import (
@@ -44,14 +45,29 @@ def process(self, instance):
 
         # Let's break the instance into multiple instances to integrate
         # a product per generated texture or texture UDIM sequence
-        for (texture_set_name, stack_name), template_maps in maps.items():
-            self.log.info(f"Processing {texture_set_name}/{stack_name}")
-            for template, outputs in template_maps.items():
-                self.log.info(f"Processing {template}")
-                self.create_image_instance(instance, template, outputs,
-                                           task_entity=task_entity,
-                                           texture_set_name=texture_set_name,
-                                           stack_name=stack_name)
+        creator_attr = instance.data["creator_attributes"]
+        if creator_attr.get("exportTextureSetsAsOneOutput", False):
+            texture_sets_by_map_identifier = defaultdict(list)
+            for (texture_set_name, stack_name), template_maps in maps.items():
+                for template, outputs in template_maps.items():
+                    self.log.info(f"Processing {template}")
+                    map_identifier = strip_template(template)
+                    map_identifier = f"{map_identifier}"
+                    texture_sets_by_map_identifier[map_identifier].extend(outputs)
+            for map_identifier, outputs in texture_sets_by_map_identifier.items():
+                self.log.info(f"Processing {map_identifier}")
+                self.create_image_instance_by_map_filtering(
+                    instance, outputs, task_entity, map_identifier)
+
+        else:
+            for (texture_set_name, stack_name), template_maps in maps.items():
+                self.log.info(f"Processing {texture_set_name}/{stack_name}")
+                for template, outputs in template_maps.items():
+                    self.log.info(f"Processing {template}")
+                    self.create_image_instance(instance, template, outputs,
+                                                task_entity=task_entity,
+                                                texture_set_name=texture_set_name,
+                                                stack_name=stack_name)
 
     def create_image_instance(self, instance, template, outputs,
                               task_entity, texture_set_name, stack_name):
@@ -156,6 +172,110 @@ def create_image_instance(self, instance, template, outputs,
         image_instance.data["textureSetName"] = texture_set_name
         image_instance.data["textureStackName"] = stack_name
 
+        # Store color space with the instance
+        # Note: The extractor will assign it to the representation
+        colorspace = outputs[0].get("colorSpace")
+        if colorspace:
+            self.log.debug(f"{image_product_name} colorspace: {colorspace}")
+            image_instance.data["colorspace"] = colorspace
+
+        # Store the instance in the original instance as a member
+        instance.append(image_instance)
+
+    def create_image_instance_by_map_filtering(self, instance, outputs,
+                                               task_entity, map_identifier):
+        """Create a new instance per image based on map filtering.
+
+        The new instances will be of product type `image`.
+        **Only used for exporting multiple texture sets as one texture output
+
+        """
+
+        context = instance.context
+        first_filepath = outputs[0]["filepath"]
+        fnames = [os.path.basename(output["filepath"]) for output in outputs]
+        ext = os.path.splitext(first_filepath)[1]
+        assert ext.lstrip("."), f"No extension: {ext}"
+        # Function to remove textureSet from filepath
+        def remove_texture_set_token(filepath, texture_set):
+            return filepath.replace(texture_set, '')
+
+        fnames_without_textureSet = [
+            remove_texture_set_token(output["output"], output["textureSet"])
+            for output in outputs
+        ]
+
+        task_name = task_type = None
+        if task_entity:
+            task_name = task_entity["name"]
+            task_type = task_entity["taskType"]
+
+        # TODO: The product type actually isn't 'texture' currently but
+        #   for now this is only done so the product name starts with
+        #   'texture'
+        image_product_name = get_product_name(
+            context.data["projectName"],
+            task_name,
+            task_type,
+            context.data["hostName"],
+            product_type="texture",
+            variant=instance.data["variant"] + f".{map_identifier}",
+            project_settings=context.data["project_settings"]
+        )
+        image_product_group_name = get_product_name(
+            context.data["projectName"],
+            task_name,
+            task_type,
+            context.data["hostName"],
+            product_type="texture",
+            variant=instance.data["variant"],
+            project_settings=context.data["project_settings"]
+        )
+
+        # Prepare representation
+        representation = {
+            "name": ext.lstrip("."),
+            "ext": ext.lstrip("."),
+            #TODO: strip the texture_sets.
+            "files": (
+                fnames_without_textureSet
+                if len(fnames_without_textureSet) > 1
+                else fnames_without_textureSet[0]
+            ),
+        }
+
+        # Mark as UDIM explicitly if it has UDIM tiles.
+        if bool(outputs[0].get("udim")):
+            # The representation for a UDIM sequence should have a `udim` key
+            # that is a list of all udim tiles (str) like: ["1001", "1002"]
+            # strings. See CollectTextures plug-in and Integrators.
+            representation["udim"] = [output["udim"] for output in outputs]
+
+        # Set up the representation for thumbnail generation
+        # TODO: Simplify this once thumbnail extraction is refactored
+        staging_dir = os.path.dirname(first_filepath)
+        representation["tags"] = ["review"]
+        representation["stagingDir"] = staging_dir
+        # Clone the instance
+        product_type = "image"
+        image_instance = context.create_instance(image_product_name)
+        image_instance[:] = instance[:]
+        image_instance.data.update(copy.deepcopy(dict(instance.data)))
+        image_instance.data["name"] = image_product_name
+        image_instance.data["label"] = image_product_name
+        image_instance.data["productName"] = image_product_name
+        image_instance.data["productType"] = product_type
+        image_instance.data["family"] = product_type
+        image_instance.data["families"] = [product_type, "textures"]
+        if instance.data["creator_attributes"].get("review"):
+            image_instance.data["families"].append("review")
+        image_instance.data["image_outputs"] = fnames
+        image_instance.data["representations"] = [representation]
+
+        # Group the textures together in the loader
+        image_instance.data["productGroup"] = image_product_group_name
+
+
         # Store color space with the instance
         # Note: The extractor will assign it to the representation
         colorspace = outputs[0].get("colorSpace")

client/ayon_substancepainter/plugins/publish/extract_map_as_single_output.py

@@ -0,0 +1,96 @@
+import clique
+import os
+
+from ayon_core.pipeline import publish
+from ayon_core.lib import (
+    get_oiio_tool_args,
+    run_subprocess,
+)
+
+
+def convert_texture_maps_for_udim_export(staging_dir, image_outputs, has_udim=False):
+    if has_udim:
+        collections, remainder = clique.assemble(image_outputs, minimum_items=1)
+        return [
+            os.path.join(
+                staging_dir,
+                collection.format(pattern="{head}{padding}{tail}")
+            )
+            for collection in collections
+        ]
+    else:
+        return [
+            os.path.join(staging_dir, output) for output in image_outputs
+        ]
+
+
+def convert_texture_maps_as_single_output(staging_dir, source_image_outputs,
+                                          dest_image_outputs, has_udim=False,
+                                          log=None):
+    oiio_tool_args = get_oiio_tool_args("oiiotool")
+
+    source_maps = convert_texture_maps_for_udim_export(
+        staging_dir, source_image_outputs, has_udim=has_udim)
+    dest_map = next(convert_texture_maps_for_udim_export(
+        staging_dir, dest_image_outputs, has_udim=has_udim
+        ), None)
+
+    log.info(f"{source_maps} composited as {dest_map}")
+    oiio_cmd = oiio_tool_args + source_maps + [
+        "--over", "-o",
+       dest_map
+    ]
+
+    subprocess_args = " ".join(oiio_cmd)
+
+    env = os.environ.copy()
+    env.pop("OCIO", None)
+    log.info(" ".join(subprocess_args))
+    try:
+        run_subprocess(subprocess_args, env=env)
+    except Exception:
+        log.error("Texture maketx conversion failed", exc_info=True)
+        raise
+
+
+class ExtractTexturesAsSingleOutput(publish.Extractor):
+    """Extract Texture As Single Output
+
+    Combine the multliple texture sets into one single texture output.
+
+    """
+
+    label = "Extract Texture Sets as Single Texture Output"
+    hosts = ["substancepainter"]
+    families = ["image"]
+    settings_category = "substancepainter"
+
+    # Run directly after textures export
+    order = publish.Extractor.order - 0.099
+
+    def process(self, instance):
+        if "exportTextureSetsAsOneOutput" not in instance.data["creator_attributes"]:
+            self.log.debug(
+                "Skipping to export texture sets as single texture output.."
+            )
+            return
+
+        representations: "list[dict]" = instance.data["representations"]
+
+        staging_dir = instance.data["stagingDir"]
+        source_image_outputs = instance.data["image_outputs"]
+        has_udim = False
+        dest_image_outputs = []
+        for representation in list(representations):
+            dest_files = representation["files"]
+            is_sequence = isinstance(dest_files, (list, tuple))
+            if not is_sequence:
+                dest_image_outputs = [dest_image_outputs]
+            if "udim" in representation:
+                has_udim = True
+
+        convert_texture_maps_as_single_output(
+            staging_dir, source_image_outputs,
+            dest_image_outputs, has_udim=has_udim,
+            log=self.log
+        )