Merge pull request #10 from ynput/enhancement/AY-6681_extract_maketx

Generate .tx files for exported texture sets

client/ayon_substancepainter/plugins/publish/extract_maketx.py

@@ -0,0 +1,187 @@
+import copy
+import os
+
+from ayon_core.pipeline import KnownPublishError, publish
+from ayon_core.lib import (
+    ToolNotFoundError,
+    get_oiio_tool_args,
+    run_subprocess,
+)
+from ayon_core.pipeline.colorspace import (
+    get_ocio_config_colorspaces
+)
+
+
+def convert_to_tx(
+    source,
+    ocio_config_path=None,
+    colorspace=None,
+    target_colorspace=None,
+    staging_dir=None,
+    log=None
+):
+    """Process the texture.
+
+    This function requires the `maketx` executable to be available in an
+    OpenImageIO toolset detectable by AYON.
+
+    Args:
+        source (str): Path to source file.
+        ocio_config_path (str): Path to the OCIO config file.
+        colorspace (str): Colorspace of the source file.
+        target_colorspace (str): Target colorspace
+        staging_dir (str): Output directory to write to.
+        log (logging.Logger): Python logger.
+
+    Returns:
+        str: The resulting texture path.
+
+    """
+
+    try:
+        maketx_args = get_oiio_tool_args("maketx")
+    except ToolNotFoundError:
+        raise KnownPublishError(
+            "OpenImageIO is not available on the machine")
+
+    # Define .tx filepath in staging if source file is not .tx
+    fname, ext = os.path.splitext(os.path.basename(source))
+    if ext == ".tx":
+        return source
+
+    # Hardcoded default arguments for maketx conversion based on Arnold's
+    # txManager in Maya
+    args = [
+        # unpremultiply before conversion (recommended when alpha present)
+        "--unpremult",
+        # use oiio-optimized settings for tile-size, planarconfig, metadata
+        "--oiio",
+        "--filter", "lanczos3",
+    ]
+
+    if ocio_config_path:
+        args.extend(["--colorconvert", colorspace, target_colorspace])
+        args.extend(["--colorconfig", ocio_config_path])
+
+    subprocess_args = maketx_args + [
+        "-v",  # verbose
+        "-u",  # update mode
+        # --checknan doesn't influence the output file but aborts the
+        # conversion if it finds any. So we can avoid it for the file hash
+        "--checknan",
+        source
+    ]
+
+    subprocess_args.extend(args)
+    # if self.extra_args:
+    #     subprocess_args.extend(self.extra_args)
+
+    destination = os.path.join(staging_dir, fname + ".tx")
+    subprocess_args.extend(["-o", destination])
+
+    # We want to make sure we are explicit about what OCIO config gets
+    # used. So when we supply no --colorconfig flag that no fallback to
+    # an OCIO env var occurs.
+    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
+
+    return destination
+
+
+class ExtractMakeTX(publish.Extractor,
+                    publish.ColormanagedPyblishPluginMixin,
+                    publish.OptionalPyblishPluginMixin):
+    """Extract MakeTX
+
+    This requires color management to be enabled so that the MakeTX file
+    generation is converted to the correct render colorspace.
+
+    Adds an extra `tx` representation to the instance.
+
+    """
+
+    label = "Extract TX"
+    hosts = ["substancepainter"]
+    families = ["image"]
+    settings_category = "substancepainter"
+
+    # Run directly after textures export
+    order = publish.Extractor.order - 0.099
+
+    def process(self, instance):
+        if not self.is_active(instance.data):
+            return
+
+        representations: "list[dict]" = instance.data["representations"]
+
+        # If a tx representation is present we skip extraction
+        if any(repre["name"] == "tx" for repre in representations):
+            return
+
+        for representation in list(representations):
+            tx_representation = copy.deepcopy(representation)
+            tx_representation["name"] = "tx"
+            tx_representation["ext"] = "tx"
+
+            colorspace_data: dict = tx_representation.get("colorspaceData", {})
+            if not colorspace_data:
+                self.log.debug(
+                    "Skipping .tx conversion for representation "
+                    f"{representation['name']} because it has no colorspace "
+                    "data.")
+                continue
+
+            colorspace: str = colorspace_data["colorspace"]
+            ocio_config_path: str = colorspace_data["config"]["path"]
+            target_colorspace = self.get_target_colorspace(ocio_config_path)
+
+            source_files = representation["files"]
+            is_sequence = isinstance(source_files, (list, tuple))
+            if not is_sequence:
+                source_files = [source_files]
+
+            # Generate the TX files
+            tx_files = []
+            staging_dir = instance.data["stagingDir"]
+            for source_filename in source_files:
+                source_filepath = os.path.join(staging_dir, source_filename)
+                self.log.debug(f"Converting to .tx: {source_filepath}")
+                tx_filepath = convert_to_tx(
+                    source_filepath,
+                    ocio_config_path=ocio_config_path,
+                    colorspace=colorspace,
+                    target_colorspace=target_colorspace,
+                    staging_dir=staging_dir,
+                    log=self.log
+                )
+                tx_filename = os.path.basename(tx_filepath)
+                tx_files.append(tx_filename)
+
+            # Make sure to store again as single file it was also in the
+            # original representation
+            if not is_sequence:
+                tx_files = tx_files[0]
+
+            tx_representation["files"] = tx_files
+
+            representations.append(tx_representation)
+
+            # Only ever one `tx` representation is needed
+            break
+
+        else:
+            self.log.warning(
+                "No .tx file conversions occurred. This may happen because"
+                " no representations were found with colorspace data."
+            )
+
+    def get_target_colorspace(self, ocio_path: str) -> str:
+        ocio_colorspaces = get_ocio_config_colorspaces(ocio_path)
+        return ocio_colorspaces["roles"]["rendering"]["colorspace"]

server/settings/main.py

@@ -2,6 +2,7 @@
 from .imageio import ImageIOSettings, DEFAULT_IMAGEIO_SETTINGS
 from .creator_plugins import CreatorsModel, DEFAULT_CREATOR_SETTINGS
 from .load_plugins import LoadersModel, DEFAULT_LOADER_SETTINGS
+from .publish_plugins import PublishersModel, DEFAULT_PUBLISH_SETTINGS
 
 
 class ShelvesSettingsModel(BaseSettingsModel):
@@ -23,12 +24,14 @@ class SubstancePainterSettings(BaseSettingsModel):
         default_factory=DEFAULT_CREATOR_SETTINGS, title="Creators")
     load: LoadersModel = SettingsField(
         default_factory=DEFAULT_LOADER_SETTINGS, title="Loaders")
+    publish: PublishersModel = SettingsField(
+        default_factory=PublishersModel, title="Publishers")
 
 
 DEFAULT_SPAINTER_SETTINGS = {
     "imageio": DEFAULT_IMAGEIO_SETTINGS,
     "shelves": [],
     "create": DEFAULT_CREATOR_SETTINGS,
     "load": DEFAULT_LOADER_SETTINGS,
-
+    "publish": DEFAULT_PUBLISH_SETTINGS,
 }

server/settings/publish_plugins.py

@@ -0,0 +1,23 @@
+from ayon_server.settings import BaseSettingsModel, SettingsField
+
+
+class BasicEnabledModel(BaseSettingsModel):
+    enabled: bool = SettingsField(title="Enabled")
+    optional: bool = SettingsField(title="Optional")
+    active: bool = SettingsField(title="Active")
+
+
+class PublishersModel(BaseSettingsModel):
+    ExtractMakeTX: BasicEnabledModel = SettingsField(
+        default_factory=BasicEnabledModel,
+        title="Extract Make TX",
+    )
+
+
+DEFAULT_PUBLISH_SETTINGS = {
+    "ExtractMakeTX": {
+        "enabled": True,
+        "optional": True,
+        "active": True,
+    },
+}