From 3c23f6cc3f833ff085ef1c146511595ef5e16dc7 Mon Sep 17 00:00:00 2001
From: Roy Nieterau <roy_nieterau@hotmail.com>
Date: Mon, 25 Mar 2024 14:00:57 +0100
Subject: [PATCH] Houdini: Improve load image

- Support setting OCIO colorspace
- Support single frame file versus sequence
- Support 'switch' folder functionality
- Allow loading more product types
---
 .../hosts/houdini/plugins/load/load_image.py  | 96 ++++++++++++++-----
 1 file changed, 73 insertions(+), 23 deletions(-)

diff --git a/client/ayon_core/hosts/houdini/plugins/load/load_image.py b/client/ayon_core/hosts/houdini/plugins/load/load_image.py
index b77e4f662a..72bb873eff 100644
--- a/client/ayon_core/hosts/houdini/plugins/load/load_image.py
+++ b/client/ayon_core/hosts/houdini/plugins/load/load_image.py
@@ -1,4 +1,5 @@
 import os
+import re
 
 from ayon_core.pipeline import (
     load,
@@ -44,7 +45,14 @@ def get_image_avalon_container():
 class ImageLoader(load.LoaderPlugin):
     """Load images into COP2"""
 
-    product_types = {"imagesequence"}
+    product_types = {
+        "imagesequence",
+        "review",
+        "render",
+        "plate",
+        "image",
+        "online",
+    }
     label = "Load Image (COP2)"
     representations = ["*"]
     order = -10
@@ -55,10 +63,8 @@ class ImageLoader(load.LoaderPlugin):
     def load(self, context, name=None, namespace=None, data=None):
 
         # Format file name, Houdini only wants forward slashes
-        file_path = self.filepath_from_context(context)
-        file_path = os.path.normpath(file_path)
-        file_path = file_path.replace("\\", "/")
-        file_path = self._get_file_sequence(file_path)
+        path = self.filepath_from_context(context)
+        path = self.format_path(path, representation=context["representation"])
 
         # Get the root node
         parent = get_image_avalon_container()
@@ -70,7 +76,10 @@ class ImageLoader(load.LoaderPlugin):
         node = parent.createNode("file", node_name=node_name)
         node.moveToGoodPosition()
 
-        node.setParms({"filename1": file_path})
+        parms = {"filename1": path}
+        parms.update(self.get_colorspace_parms(context["representation"]))
+
+        node.setParms(parms)
 
         # Imprint it manually
         data = {
@@ -93,16 +102,17 @@ class ImageLoader(load.LoaderPlugin):
 
         # Update the file path
         file_path = get_representation_path(repre_entity)
-        file_path = file_path.replace("\\", "/")
-        file_path = self._get_file_sequence(file_path)
+        file_path = self.format_path(file_path, repre_entity)
+
+        parms = {
+            "filename1": file_path,
+            "representation": repre_entity["id"],
+        }
+
+        parms.update(self.get_colorspace_parms(repre_entity))
 
         # Update attributes
-        node.setParms(
-            {
-                "filename1": file_path,
-                "representation": repre_entity["id"],
-            }
-        )
+        node.setParms(parms)
 
     def remove(self, container):
 
@@ -119,14 +129,54 @@ class ImageLoader(load.LoaderPlugin):
         if not parent.children():
             parent.destroy()
 
-    def _get_file_sequence(self, file_path):
-        root = os.path.dirname(file_path)
-        files = sorted(os.listdir(root))
+    @staticmethod
+    def format_path(path, representation):
+        """Format file path correctly for single image or sequence."""
+        if not os.path.exists(path):
+            raise RuntimeError("Path does not exist: %s" % path)
 
-        first_fname = files[0]
-        prefix, padding, suffix = first_fname.rsplit(".", 2)
-        fname = ".".join([prefix, "$F{}".format(len(padding)), suffix])
-        return os.path.join(root, fname).replace("\\", "/")
+        ext = os.path.splitext(path)[-1]
 
-    def switch(self, container, context):
-        self.update(container, context)
+        is_sequence = bool(representation["context"].get("frame"))
+        # The path is either a single file or sequence in a folder.
+        if not is_sequence:
+            filename = path
+        else:
+            filename = re.sub(r"(.*)\.(\d+){}$".format(re.escape(ext)),
+                              "\\1.$F4{}".format(ext),
+                              path)
+
+            filename = os.path.join(path, filename)
+
+        filename = os.path.normpath(filename)
+        filename = filename.replace("\\", "/")
+
+        return filename
+
+    def get_colorspace_parms(self, representation: dict) -> dict:
+        """Return the color space parameters.
+
+        Returns the values for the colorspace parameters on the node if there
+        is colorspace data on the representation.
+
+        Arguments:
+            representation (dict): The representation entity.
+
+        Returns:
+            dict: Parm to value mapping if colorspace data is defined.
+
+        """
+
+        data = representation.get("data", {}).get("colorspaceData", {})
+        if not data:
+            return {}
+
+        colorspace = data["colorspace"]
+        if colorspace:
+            return {
+                "colorspace": 3,  # Use OpenColorIO
+                "ocio_space": colorspace
+            }
+
+    def switch(self, container, representation):
+        self.update(container, representation)