Fix FunctionUtil colorfamily conversion edgecases (#156)

* Update funcs.py * Update funcs.py * Update funcs.py * Update funcs.py * Update funcs.py * Update funcs.py * Update funcs.py * Update funcs.py * Update funcs.py * Update funcs.py * Update funcs.py * Remove unnecessary test * InvalidColorspacePathError: New exception * FunctionUtil: Check invalid colorspace path * Update FunctionUtil test * FunctionUtil: Fix YUV/RGB => RGB/YUV with planes=0 * Remove FunctionUtil GRAY to YUV test * Update funcs.py * Update funcs.py * Minor changes --------- Co-authored-by: LightArrowsEXE <Lightarrowsreboot@gmail.com>
Jaded-Encoding-Thaumaturgy · Oct 14, 2024 · 9bf5380 · 9bf5380
1 parent 07f7a07
commit 9bf5380
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 30 deletions.
diff --git a/tests/functions/test_funcs.py b/tests/functions/test_funcs.py
@@ -1,7 +1,9 @@
 from typing import Callable, cast
 from unittest import TestCase
 
-from vstools import FunctionUtil, UndefinedMatrixError, fallback, iterate, kwargs_fallback, vs
+from vstools import (
+    FunctionUtil, InvalidColorspacePathError, UndefinedMatrixError, fallback, iterate, kwargs_fallback, vs
+)
 
 
 class TestFuncs(TestCase):
@@ -101,19 +103,6 @@ def test_functionutil_color_family_conversion_gray_to_gray(self) -> None:
         self.assertEqual(result.work_clip.format.color_family, vs.GRAY)
         self.assertFalse(result.cfamily_converted)
 
-    def test_functionutil_color_family_conversion_gray_to_yuv(self) -> None:
-        clip = vs.core.std.BlankClip(format=vs.GRAY8)
-        result = FunctionUtil(clip, 'FunctionUtilTest', color_family=vs.YUV, matrix=1)
-        self.assertEqual(result.work_clip.format.color_family, vs.YUV)
-        self.assertTrue(result.cfamily_converted)
-        self.assertEqual(result.work_clip.format.subsampling_w, 0)
-        self.assertEqual(result.work_clip.format.subsampling_h, 0)
-
-    def test_functionutil_color_family_conversion_gray_to_yuv_without_matrix(self) -> None:
-        clip = vs.core.std.BlankClip(format=vs.GRAY8)
-        with self.assertRaises(UndefinedMatrixError):
-            FunctionUtil(clip, 'FunctionUtilTest', color_family=vs.YUV)
-
     def test_functionutil_color_family_conversion_gray_to_rgb(self) -> None:
         clip = vs.core.std.BlankClip(format=vs.GRAY8)
         result = FunctionUtil(clip, 'FunctionUtilTest', color_family=vs.RGB, matrix=1)
@@ -158,7 +147,7 @@ def test_functionutil_color_family_conversion_rgb_to_rgb(self) -> None:
 
     def test_functionutil_color_conversions_yuv_to_rgb_without_matrix(self) -> None:
         yuv_clip = vs.core.std.BlankClip(format=vs.YUV420P8)
-        with self.assertRaises(UndefinedMatrixError):
+        with self.assertRaises(InvalidColorspacePathError):
             FunctionUtil(yuv_clip, 'FunctionUtilTest', color_family=vs.RGB)
 
     def test_functionutil_color_conversions_yuv_to_rgb_with_matrix(self) -> None:
@@ -246,3 +235,18 @@ def test_functionutil_num_planes_rgb(self) -> None:
         clip_rgb = vs.core.std.BlankClip(format=vs.RGB24)
         result_rgb = FunctionUtil(clip_rgb, 'FunctionUtilTest')
         self.assertEqual(result_rgb.num_planes, 3)
+
+    def test_functionutil_planes_0_yuv_to_rgb(self) -> None:
+        clip = vs.core.std.BlankClip(format=vs.YUV420P8)
+        func_util = FunctionUtil(clip, 'FunctionUtilTest', planes=0, color_family=vs.RGB, matrix=1)
+        self.assertTrue(func_util.cfamily_converted)
+        self.assertEqual(func_util.work_clip.format.color_family, vs.GRAY)
+        self.assertEqual(func_util.norm_planes, [0])
+
+    def test_functionutil_planes_0_rgb_to_yuv(self) -> None:
+        clip = vs.core.std.BlankClip(format=vs.RGB24)
+        func_util = FunctionUtil(clip, 'FunctionUtilTest', planes=0, color_family=vs.YUV, matrix=1)
+        self.assertTrue(func_util.cfamily_converted)
+        self.assertEqual(func_util.work_clip.format.color_family, vs.GRAY)
+        self.assertEqual(func_util.norm_planes, [0])
+
diff --git a/vstools/exceptions/color.py b/vstools/exceptions/color.py
@@ -2,9 +2,12 @@
 
 from typing import Any
 
+import vapoursynth as vs
 from stgpytools import CustomPermissionError, CustomValueError, FuncExceptT, SupportsString
 
 __all__ = [
+    'InvalidColorspacePathError',
+
     'UndefinedMatrixError',
     'UndefinedTransferError',
     'UndefinedPrimariesError',
@@ -23,6 +26,47 @@
     'UnsupportedColorRangeError'
 ]
 
+########################################################
+# Colorspace
+
+class InvalidColorspacePathError(CustomValueError):
+    """Raised when there is no path between two colorspaces."""
+
+    def __init__(
+        self, func: FuncExceptT, message: SupportsString | None = None,
+        **kwargs: Any
+    ) -> None:
+        def_msg = 'Unable to convert between colorspaces! '
+        def_msg += 'Please provide more colorspace information (e.g., matrix, transfer, primaries).'
+
+        if isinstance(message, vs.Error):
+            error_msg = str(message)
+            if 'Resize error:' in error_msg:
+                kwargs['reason'] = error_msg[error_msg.find('(') + 1:error_msg.rfind(')')]
+                message = def_msg
+
+        super().__init__(message or def_msg, func, **kwargs)
+
+    @staticmethod
+    def check(func: FuncExceptT, to_check: vs.VideoNode) -> None:
+        """
+        Check if there's a valid colorspace path for the given clip.
+
+        :param func:        Function returned for custom error handling.
+                            This should only be set by VS package developers.
+        :param to_check:    Value to check. Must be a VideoNode.
+
+        :raises InvalidColorspacePathError: If there's no valid colorspace path.
+        """
+
+        try:
+            to_check.get_frame(0)
+        except vs.Error as e:
+            if 'no path between colorspaces' in str(e):
+                raise InvalidColorspacePathError(func, e)
+            raise
+
+
 ########################################################
 # Matrix
 

diff --git a/vstools/functions/funcs.py b/vstools/functions/funcs.py
@@ -5,6 +5,8 @@
 import vapoursynth as vs
 from stgpytools import FuncExceptT, T, cachedproperty, fallback, iterate, kwargs_fallback, normalize_seq, to_arr
 
+from vstools.exceptions.color import InvalidColorspacePathError
+
 from ..enums import (
     ColorRange, ColorRangeT, Matrix, MatrixT, Transfer, TransferT, Primaries, PrimariesT,
     ChromaLocation, ChromaLocationT, FieldBased, FieldBasedT
@@ -91,6 +93,8 @@ def __init__(
 
         if color_family is not None:
             color_family = [get_color_family(c) for c in to_arr(color_family)]
+            if not set(color_family) & {vs.YUV, vs.RGB}:
+                planes = 0
 
         if isinstance(bitdepth, tuple):
             bitdepth = range(bitdepth[0], bitdepth[1] + 1)
@@ -109,7 +113,7 @@ def __init__(
         self._chromaloc = chromaloc
         self._order = order
 
-        self.norm_planes = normalize_planes(self.norm_clip, planes)
+        self.norm_planes = normalize_planes(self.norm_clip, self.planes)
 
         super().__init__(self.norm_planes)
 
@@ -119,9 +123,9 @@ def __init__(
     def norm_clip(self) -> ConstantFormatVideoNode:
         """Get a "normalized" clip. This means color space and bitdepth are converted if necessary."""
 
-        from .. import get_depth
-
         if isinstance(self.bitdepth, (range, set)) and self.clip.format.bits_per_sample not in self.bitdepth:
+            from .. import get_depth
+
             src_depth = get_depth(self.clip)
             target_depth = next((bits for bits in self.bitdepth if bits >= src_depth), max(self.bitdepth))
 
@@ -138,22 +142,29 @@ def norm_clip(self) -> ConstantFormatVideoNode:
         if not self.allowed_cfamilies or cfamily in self.allowed_cfamilies:
             return clip
 
-        if cfamily is vs.YUV and vs.GRAY in self.allowed_cfamilies:
-            return plane(clip, 0)
+        if cfamily is vs.RGB:
+            if not self._matrix:
+                raise UndefinedMatrixError(
+                    'You must specify a matrix for RGB to '
+                    f'{'/'.join(cf.name for cf in sorted(self.allowed_cfamilies, key=lambda x: x.name))} conversions!',
+                    self.func
+                )
 
-        self.cfamily_converted = True
+            self.cfamily_converted = True
 
-        if cfamily is vs.YUV:
-            return clip.resize.Bicubic(format=clip.format.replace(color_family=vs.RGB, subsampling_h=0, subsampling_w=0))
+            clip = clip.resize.Bicubic(format=clip.format.replace(color_family=vs.YUV), matrix=self._matrix)
 
-        if not self._matrix:
-            raise UndefinedMatrixError(
-                'You must specify a matrix for RGB to '
-                f'{'/'.join(cf.name for cf in sorted(self.allowed_cfamilies, key=lambda x: x.name))} conversions!',
-                self.func
+        elif cfamily in (vs.YUV, vs.GRAY) and not set(self.allowed_cfamilies) & {vs.YUV, vs.GRAY}:
+            self.cfamily_converted = True
+
+            clip = clip.resize.Bicubic(
+                format=clip.format.replace(color_family=vs.RGB, subsampling_h=0, subsampling_w=0),
+                matrix_in=self._matrix
             )
 
-        return clip.resize.Bicubic(format=clip.format.replace(color_family=vs.YUV), matrix=self._matrix)
+            InvalidColorspacePathError.check(self.func, clip)
+
+        return clip
 
     @cachedproperty
     def work_clip(self) -> ConstantFormatVideoNode:
@@ -165,7 +176,7 @@ def work_clip(self) -> ConstantFormatVideoNode:
     def chroma_planes(self) -> list[vs.VideoNode]:
         """Get a list of all chroma planes in the normalised clip."""
 
-        if self == [0] or self.norm_clip.format.num_planes == 1:
+        if self != [0] or self.norm_clip.format.num_planes == 1:
             return []
 
         return [plane(self.norm_clip, i) for i in (1, 2)]