Added Helper function GrayScale()

HomeOfVapourSynthEvolution · Apr 8, 2016 · 8c838a5 · 8c838a5
1 parent 8fbba92
commit 8c838a5
Showing 1 changed file with 77 additions and 32 deletions.
diff --git a/mvsfunc.py b/mvsfunc.py
@@ -42,8 +42,9 @@
 ##     CheckVersion
 ##     GetMatrix
 ##     zDepth
-##     GetPlane
 ##     PlaneAverage
+##     GetPlane
+##     GrayScale
 ################################################################################################################################
 
 
@@ -2452,20 +2453,23 @@ def zDepth(clip, sample=None, depth=None, range=None, range_in=None, dither_type
 
 
 ################################################################################################################################
-## Helper function: GetPlane()
+## Helper function: PlaneAverage()
 ################################################################################################################################
-## Extract specific plane and store it in a Gray clip.
+## Evaluate normalized average value of the specified plane and store it as a frame property.
+## Mainly as a wrap function to support both std.PlaneAverage and std.PlaneStats with the same interface.
 ################################################################################################################################
 ## Parameters
 ##     clip {clip}: the source clip
 ##         can be of any constant format
-##     plane {int}: the plane to extract
+##     plane {int}: the plane to evaluate
 ##         default: 0
+##     prop {str}: the frame property name to be written
+##         default: 'PlaneAverage'
 ################################################################################################################################
-def GetPlane(clip, plane=None):
+def PlaneAverage(clip, plane=None, prop=None):
     # Set VS core and function name
     core = vs.get_core()
-    funcName = 'GetPlane'
+    funcName = 'PlaneAverage'
 
     if not isinstance(clip, vs.VideoNode):
         raise TypeError(funcName + ': \"clip\" must be a clip!')
@@ -2482,29 +2486,46 @@ def GetPlane(clip, plane=None):
     elif plane < 0 or plane > sNumPlanes:
         raise ValueError(funcName + ': valid range of \"plane\" is [0, {})!'.format(sNumPlanes))
 
+    if prop is None:
+        prop = 'PlaneAverage'
+    elif not isinstance(prop, str):
+        raise TypeError(funcName + ': \"prop\" must be a str!')
+
     # Process
-    return core.std.ShufflePlanes(clip, plane, vs.GRAY)
+    if core.std.get_functions().__contains__('PlaneAverage'):
+        clip = core.std.PlaneAverage(clip, plane=plane, prop=prop)
+    elif core.std.get_functions().__contains__('PlaneStats'):
+        clip = core.std.PlaneStats(clip, plane=plane, prop='PlaneStats')
+        def _PlaneAverageTransfer(n, f):
+            fout = f.copy()
+            fout.props.__setattr__(prop, f.props.PlaneStatsAverage)
+            del fout.props.PlaneStatsAverage
+            del fout.props.PlaneStatsMinMax
+            return fout
+        clip = core.std.ModifyFrame(clip, clip, selector=_PlaneAverageTransfer)
+    else:
+        raise AttributeError(funcName + ': Available plane average function not found!')
+
+    # output
+    return clip
 ################################################################################################################################
 
 
 ################################################################################################################################
-## Helper function: PlaneAverage()
+## Helper function: GetPlane()
 ################################################################################################################################
-## Evaluate normalized average value of the specified plane and store it as a frame property.
-## Mainly as a wrap function to support both std.PlaneAverage and std.PlaneStats with the same interface.
+## Extract specific plane and store it in a Gray clip.
 ################################################################################################################################
 ## Parameters
 ##     clip {clip}: the source clip
 ##         can be of any constant format
-##     plane {int}: the plane to evaluate
+##     plane {int}: the plane to extract
 ##         default: 0
-##     prop {str}: the frame property name to be written
-##         default: 'PlaneAverage'
 ################################################################################################################################
-def PlaneAverage(clip, plane=None, prop=None):
+def GetPlane(clip, plane=None):
     # Set VS core and function name
     core = vs.get_core()
-    funcName = 'PlaneAverage'
+    funcName = 'GetPlane'
 
     if not isinstance(clip, vs.VideoNode):
         raise TypeError(funcName + ': \"clip\" must be a clip!')
@@ -2521,27 +2542,51 @@ def PlaneAverage(clip, plane=None, prop=None):
     elif plane < 0 or plane > sNumPlanes:
         raise ValueError(funcName + ': valid range of \"plane\" is [0, {})!'.format(sNumPlanes))
 
-    if prop is None:
-        prop = 'PlaneAverage'
-    elif not isinstance(prop, str):
-        raise TypeError(funcName + ': \"prop\" must be a str!')
+    # Process
+    return core.std.ShufflePlanes(clip, plane, vs.GRAY)
+################################################################################################################################
+
+
+################################################################################################################################
+## Helper function: GrayScale()
+################################################################################################################################
+## Convert the give clip to gray-scale.
+################################################################################################################################
+## Parameters
+##     clip {clip}: the source clip
+##         can be of any constant format
+##     matrix {int|str}: for RGB input only, same as the one in ToYUV()
+################################################################################################################################
+def GrayScale(clip, matrix=None):
+    # Set VS core and function name
+    core = vs.get_core()
+    funcName = 'GrayScale'
+
+    if not isinstance(clip, vs.VideoNode):
+        raise TypeError(funcName + ': \"clip\" must be a clip!')
+
+    # Get properties of input clip
+    sFormat = clip.format
+
+    sColorFamily = sFormat.color_family
+    sIsRGB = sColorFamily == vs.RGB
+    sIsYUV = sColorFamily == vs.YUV
+    sIsGRAY = sColorFamily == vs.GRAY
+    sIsYCOCG = sColorFamily == vs.YCOCG
+    if sColorFamily == vs.COMPAT:
+        raise ValueError(funcName + ': color family *COMPAT* is not supported!')
 
     # Process
-    if core.std.get_functions().__contains__('PlaneAverage'):
-        clip = core.std.PlaneAverage(clip, plane=plane, prop=prop)
-    elif core.std.get_functions().__contains__('PlaneStats'):
-        clip = core.std.PlaneStats(clip, plane=plane, prop='PlaneStats')
-        def _PlaneAverageTransfer(n, f):
-            fout = f.copy()
-            fout.props.__setattr__(prop, f.props.PlaneStatsAverage)
-            del fout.props.PlaneStatsAverage
-            del fout.props.PlaneStatsMinMax
-            return fout
-        clip = core.std.ModifyFrame(clip, clip, selector=_PlaneAverageTransfer)
+    if sIsGRAY:
+        pass
+    elif sIsRGB:
+        clip = ToYUV(clip, matrix=matrix, full=True)
+        clip = core.std.ShufflePlanes(clip, [0,0,0], sColorFamily)
     else:
-        raise AttributeError(funcName + ': Available plane average function not found!')
+        blank = clip.std.BlankClip()
+        clip = core.std.ShufflePlanes([clip,blank,blank], [0,1,2], sColorFamily)
 
-    # output
+    # Output
     return clip
 ################################################################################################################################