get_matrix: Slight updates, new helper function

docs/submodules/helpers.rst

@@ -4,18 +4,20 @@ lvsfunc.helpers
 Basic helpers used throughout lvsfunc.
 These can be used by users as well.
 
-------------------
+---------------
 
 .. autosummary::
 
     lvsfunc.helpers._check_has_nvidia
     lvsfunc.helpers._check_index_exists
     lvsfunc.helpers._generate_dgi
     lvsfunc.helpers._get_dgidx
+    lvsfunc.helpers._get_matrix_from_res
     lvsfunc.helpers._load_dgi
     lvsfunc.helpers._tail
 
 .. automodule:: lvsfunc.helpers
     :members:
     :undoc-members:
     :show-inheritance:
+    :member-order: bysource

docs/submodules/misc.rst

@@ -1,5 +1,5 @@
 lvsfunc.misc
--------------------
+------------
 
 Miscellaneous functions and wrappers that don't really have a place elsewhere.
 

lvsfunc/helpers.py

@@ -9,12 +9,18 @@
 import vapoursynth as vs
 from vsutil import is_image
 
-from .types import IndexExists, VSIdxFunction
+from .types import IndexExists, Matrix, VSIdxFunction
 
 core = vs.core
 
 __all__: List[str] = [
-    '_check_has_nvidia', '_check_index_exists', '_generate_dgi', '_get_dgidx', '_load_dgi', '_tail'
+    '_check_has_nvidia',
+    '_check_index_exists',
+    '_generate_dgi',
+    '_get_dgidx',
+    '_get_matrix_from_res',
+    '_load_dgi',
+    '_tail'
 ]
 
 
@@ -93,3 +99,19 @@ def _load_dgi(path: str, film_thr: float, src_filter: VSIdxFunction,
         props |= dict(dgi_fieldop=1, _FieldBased=0)
 
     return src_filter(path, **index_args).std.SetFrameProps(**props)
+
+
+def _get_matrix_from_res(frame: vs.VideoFrame | vs.VideoNode) -> Matrix:
+    """Return matrix based on the frame dimensions"""
+    if isinstance(frame, vs.VideoNode):
+        frame = frame.get_frame(0)
+
+    w, h = frame.width, frame.height
+
+    if frame.format.color_family == vs.RGB:
+        return Matrix(0)
+    elif w <= 1024 and h <= 576:
+        return Matrix(6)
+    elif w <= 2048 and h <= 1536:
+        return Matrix(1)
+    return Matrix(9)

lvsfunc/util.py

@@ -14,6 +14,7 @@
 
 from .exceptions import (InvalidFormatError, InvalidMatrixError, MatrixError, VariableFormatError,
                          VariableResolutionError)
+from .helpers import _get_matrix_from_res
 from .types import CURVES, Coefs, F, Matrix, Range, T, _VideoNode
 
 core = vs.core
@@ -377,39 +378,17 @@ def get_matrix(frame: vs.VideoNode | vs.VideoFrame, strict: bool = False) -> Mat
     matrix = get_prop(frame, "_Matrix", int)
 
     match matrix:
-        case 0: return Matrix.RGB
-        case 1: return Matrix.BT709
         case 2 if strict: raise MatrixError("get_matrix", matrix, "{func}: 'Matrix is undefined.'")
+        case 2: return _get_matrix_from_res(frame)
         case 3: raise MatrixError("get_matrix", matrix, "{func}: 'Matrix is reserved.'")
-        case 4: return Matrix.FCC
-        case 5: return Matrix.BT470BG
-        case 6: return Matrix.SMPTE170M
-        case 7: return Matrix.SMPTE240M
-        case 8:
-            if core.version_number() < 55:
-                return Matrix.YCGCO
-            else:
-                raise MatrixError("get_matrix", matrix, "{func}: 'VapourSynth no longer supports {matrix}.'")
-        case 9: return Matrix.BT2020NC
-        case 10: return Matrix.BT2020C
-        case 11: return Matrix.SMPTE2085
-        case 12: return Matrix.CHROMA_DERIVED_NC
-        case 13: return Matrix.CHROMA_DERIVED_C
-        case 14: return Matrix.ICTCP
-        case _ if matrix > 14: raise MatrixError("get_matrix", matrix, "{func}: 'This matrix is current unsupported. "
-                                                 "If you believe this to be in error, please leave an issue "
-                                                 "in the lvsfunc GitHub repository.'")
-        case _ if strict: raise MatrixError("get_matrix", matrix, "{func}: 'Some kind of error occured!'")
-
-    w, h = frame.width, frame.height
-
-    if frame.format.color_family == vs.RGB:
-        return Matrix.RGB
-    elif w <= 1024 and h <= 576:
-        return Matrix.SMPTE170M
-    elif w <= 2048 and h <= 1536:
-        return Matrix.BT709
-    return Matrix.BT2020NC
+        case 8 if core.version_number() >= 55:
+            raise MatrixError("get_matrix", matrix, "{func}: 'VapourSynth no longer supports {matrix}.'")
+        case _ if matrix > 14:
+            raise MatrixError("get_matrix", matrix, "{func}: 'This matrix is current unsupported. "
+                              "If you believe this to be in error, please leave an issue "
+                              "in the lvsfunc GitHub repository.'")
+
+    return Matrix(matrix)
 
 
 def get_matrix_curve(matrix: Matrix) -> CURVES: