From 9bf538073b6b77f5b8fc10be38b8fdc1499625d2 Mon Sep 17 00:00:00 2001 From: emotion3459 <176516814+emotion3459@users.noreply.github.com> Date: Mon, 14 Oct 2024 15:27:03 -0400 Subject: [PATCH] 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 --- tests/functions/test_funcs.py | 34 +++++++++++++++------------ vstools/exceptions/color.py | 44 +++++++++++++++++++++++++++++++++++ vstools/functions/funcs.py | 41 ++++++++++++++++++++------------ 3 files changed, 89 insertions(+), 30 deletions(-) diff --git a/tests/functions/test_funcs.py b/tests/functions/test_funcs.py index 6715f790..952ffa8e 100644 --- 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 index b186ccc0..bcbe3f5c 100644 --- 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 index 6d2c0d47..87659d43 100644 --- 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)]