Enhancement: Multi-thread the convert_colorspace method

openpype/lib/transcoding.py

@@ -9,6 +9,8 @@
 
 import xml.etree.ElementTree
 
+from multiprocessing.pool import ThreadPool
+
 from .execute import run_subprocess
 from .vendor_bin_utils import (
     get_ffmpeg_tool_args,
@@ -1134,50 +1136,78 @@ def convert_colorspace(
     if logger is None:
         logger = logging.getLogger(__name__)
 
-    input_info = get_oiio_info_for_input(input_path, logger=logger)
+    def get_chunks(frame_start, frame_end, chunk_size):
+        """Chunk the frame numbers into groups."""
+        frames = list(range(frame_start, frame_end + 1))
+        chunk_list = [frames[i:i + chunk_size] for i in range(0, len(frames), chunk_size)]
+        return chunk_list
 
-    # Collect channels to export
-    input_arg, channels_arg = get_oiio_input_and_channel_args(input_info)
+    def process_conversion(in_path, out_path, start_frame=None, end_frame=None):
+        if start_frame is not None:
+            frame_group = f'.{start_frame}-{end_frame}#.'
+            frame_regex = r'\.\d+-\d+#\.'
+            in_path = re.sub(frame_regex, frame_group, in_path)
+            out_path = re.sub(frame_regex, frame_group, out_path)
 
-    # Prepare subprocess arguments
-    oiio_cmd = get_oiio_tool_args(
-        "oiiotool",
-        # Don't add any additional attributes
-        "--nosoftwareattrib",
-        "--colorconfig", config_path
-    )
+        input_info = get_oiio_info_for_input(in_path, logger=logger)
 
-    oiio_cmd.extend([
-        input_arg, input_path,
-        # Tell oiiotool which channels should be put to top stack
-        #   (and output)
-        "--ch", channels_arg,
-        # Use first subimage
-        "--subimage", "0"
-    ])
+        # Collect channels to export
+        input_arg, channels_arg = get_oiio_input_and_channel_args(input_info)
+
+        # Prepare subprocess arguments
+        oiio_cmd = get_oiio_tool_args(
+            "oiiotool",
+            # Don't add any additional attributes
+            "--nosoftwareattrib",
+            "--colorconfig", config_path
+        )
+
+        oiio_cmd.extend([
+            input_arg, in_path,
+            # Tell oiiotool which channels should be put to top stack
+            #   (and output)
+            "--ch", channels_arg,
+            # Use first subimage
+            "--subimage", "0"
+        ])
 
-    if all([target_colorspace, view, display]):
-        raise ValueError("Colorspace and both screen and display"
-                         " cannot be set together."
-                         "Choose colorspace or screen and display")
-    if not target_colorspace and not all([view, display]):
-        raise ValueError("Both screen and display must be set.")
+        if all([target_colorspace, view, display]):
+            raise ValueError("Colorspace and both screen and display"
+                             " cannot be set together."
+                             "Choose colorspace or screen and display")
+        if not target_colorspace and not all([view, display]):
+            raise ValueError("Both screen and display must be set.")
 
-    if additional_command_args:
-        oiio_cmd.extend(additional_command_args)
+        if additional_command_args:
+            oiio_cmd.extend(additional_command_args)
 
-    if target_colorspace:
-        oiio_cmd.extend(["--colorconvert",
-                         source_colorspace,
-                         target_colorspace])
-    if view and display:
-        oiio_cmd.extend(["--iscolorspace", source_colorspace])
-        oiio_cmd.extend(["--ociodisplay", display, view])
+        if target_colorspace:
+            oiio_cmd.extend(["--colorconvert",
+                             source_colorspace,
+                             target_colorspace])
+        if view and display:
+            oiio_cmd.extend(["--iscolorspace", source_colorspace])
+            oiio_cmd.extend(["--ociodisplay", display, view])
 
-    oiio_cmd.extend(["-o", output_path])
+        oiio_cmd.extend(["-o", out_path])
 
-    logger.debug("Conversion command: {}".format(" ".join(oiio_cmd)))
-    run_subprocess(oiio_cmd, logger=logger)
+        logger.debug("Conversion command: {}".format(" ".join(oiio_cmd)))
+        run_subprocess(oiio_cmd, logger=logger)
+
+    frame_data = input_path.split('.')[-2]
+    # If it's a range of frames, we multithread it
+    if '-' in frame_data:
+        start, end = frame_data[:-1].split('-')
+        # chunk_size could be editable
+        chunks = get_chunks(int(start), int(end), chunk_size=10)
+
+        pool = ThreadPool()
+        for chunk in chunks:
+            pool.apply_async(process_conversion, (input_path, output_path, chunk[0], chunk[-1]))
+        pool.close()
+        pool.join()
+    else:
+        process_conversion(input_path, output_path)
 
 
 def split_cmd_args(in_args):