Advertisement
fast62

Untitled

Nov 8th, 2021
986
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. """
  2. VSUtil. A collection of general-purpose Vapoursynth functions to be reused in modules and scripts.
  3. """
  4. from functools import reduce
  5. from typing import Callable, TypeVar, Union, List, Tuple, Optional
  6. import vapoursynth as vs
  7. import mimetypes
  8. import os
  9.  
  10. core = vs.core
  11. T = TypeVar("T")
  12.  
  13.  
  14. def get_subsampling(clip: vs.VideoNode) -> str:
  15.     """
  16.    Returns the subsampling of a clip in human-readable format.
  17.    """
  18.     if clip.format.subsampling_w == 1 and clip.format.subsampling_h == 1:
  19.         css = '420'
  20.     elif clip.format.subsampling_w == 1 and clip.format.subsampling_h == 0:
  21.         css = '422'
  22.     elif clip.format.subsampling_w == 0 and clip.format.subsampling_h == 0:
  23.         css = '444'
  24.     elif clip.format.subsampling_w == 2 and clip.format.subsampling_h == 2:
  25.         css = '410'
  26.     elif clip.format.subsampling_w == 2 and clip.format.subsampling_h == 0:
  27.         css = '411'
  28.     elif clip.format.subsampling_w == 0 and clip.format.subsampling_h == 1:
  29.         css = '440'
  30.     else:
  31.         raise ValueError('Unknown subsampling')
  32.     return css
  33.  
  34.  
  35. def get_depth(clip: vs.VideoNode) -> int:
  36.     """
  37.    Returns the bitdepth of a clip as an integer.
  38.    """
  39.     return clip.format.bits_per_sample
  40.  
  41.  
  42. def get_plane_size(frame: Union[vs.VideoFrame, vs.VideoNode], planeno: int) -> Tuple[int, int]:
  43.     """
  44.    Calculates the size of the plane
  45.    
  46.    :param frame:    The frame
  47.    :param planeno:  The plane
  48.    :return: (width, height)
  49.    """
  50.     # Add additional checks on Video Nodes as their size and format can be variable.
  51.     if isinstance(frame, vs.VideoNode):
  52.         if frame.width == 0:
  53.             raise ValueError("Cannot calculate plane size of variable size clip. Pass a frame instead.")
  54.         if frame.format is None:
  55.             raise ValueError("Cannot calculate plane size of variable format clip. Pass a frame instead.")
  56.  
  57.     width, height = frame.width, frame.height
  58.     if planeno != 0:
  59.         width >>= frame.format.subsampling_w
  60.         height >>= frame.format.subsampling_h
  61.     return width, height
  62.  
  63.  
  64. def iterate(base: T, function: Callable[[T], T], count: int) -> T:
  65.     """
  66.    Utility function that executes a given function for a given number of times.
  67.    """
  68.     return reduce(lambda v, _: function(v), range(count), base)
  69.  
  70.  
  71. def insert_clip(clip: vs.VideoNode, insert: vs.VideoNode, start_frame: int) -> vs.VideoNode:
  72.     """
  73.    Convenience method to insert a shorter clip into a longer one.
  74.    The inserted clip cannot go beyond the last frame of the source clip or an exception is raised.
  75.    """
  76.     if start_frame == 0:
  77.         return insert + clip[insert.num_frames:]
  78.     pre = clip[:start_frame]
  79.     frame_after_insert = start_frame + insert.num_frames
  80.     if frame_after_insert > clip.num_frames:
  81.         raise ValueError('Inserted clip is too long')
  82.     if frame_after_insert == clip.num_frames:
  83.         return pre + insert
  84.     post = clip[start_frame + insert.num_frames:]
  85.     return pre + insert + post
  86.  
  87.  
  88. def fallback(value: Optional[T], fallback_value: T) -> T:
  89.     """
  90.    Utility function that returns a value or a fallback if the value is None.
  91.    """
  92.     return fallback_value if value is None else value
  93.  
  94.  
  95. def plane(clip: vs.VideoNode, planeno: int) -> vs.VideoNode:
  96.     """
  97.    Extract the plane with the given index from the clip.
  98.    
  99.    :param clip:     The clip to extract the plane from.
  100.    :param planeno:  The planeno that specifies which plane to extract.
  101.    :return: A grayscale clip that only contains the given plane.
  102.    """
  103.     return core.std.ShufflePlanes(clip, planeno, vs.GRAY)
  104.  
  105.  
  106. def get_y(clip: vs.VideoNode) -> vs.VideoNode:
  107.     """
  108.    Helper to get the luma of a VideoNode.
  109.    """
  110.     return plane(clip, 0)
  111.  
  112.  
  113. def split(clip: vs.VideoNode) -> List[vs.VideoNode]:
  114.     """
  115.    Returns a list of planes for the given input clip.
  116.    """
  117.     return [plane(clip, x) for x in range(clip.format.num_planes)]
  118.  
  119.  
  120. def join(planes: List[vs.VideoNode], family=vs.YUV) -> vs.VideoNode:
  121.     """
  122.    Joins the supplied list of planes into a YUV video node.
  123.    """
  124.     return core.std.ShufflePlanes(clips=planes, planes=[0], colorfamily=family)
  125.  
  126.  
  127. def frame2clip(frame: vs.VideoFrame, *, enforce_cache=True) -> vs.VideoNode:
  128.     """
  129.    Converts a vapoursynth frame to a clip.
  130.    
  131.    :param frame:         The frame to wrap.
  132.    :param enforce_cache: Always add a cache. (Even if the vapoursynth module has this feature disabled)
  133.    :returns:     A one-frame VideoNode that yields the frame passed to the function.
  134.    """
  135.     bc = core.std.BlankClip(
  136.         width=frame.width,
  137.         height=frame.height,
  138.         length=1,
  139.         fpsnum=1,
  140.         fpsden=1,
  141.         format=frame.format.id
  142.     )
  143.     frame = frame.copy()
  144.     result = bc.std.ModifyFrame([bc], lambda n, f: frame.copy())
  145.  
  146.     # Forcefully add a cache to Modify-Frame if caching is disabled on the core.
  147.     # This will ensure that the filter will not include GIL characteristics.
  148.     if not core.add_cache and enforce_cache:
  149.         result = result.std.Cache(size=1, fixed=True)
  150.     return result
  151.  
  152.  
  153. def get_w(height: int, aspect_ratio: float = 16 / 9, only_even: bool = True) -> int:
  154.     """
  155.    Calculates the width for a clip with the given height and aspect ratio.
  156.    only_even is True by default because it imitates the math behind most standard resolutions (e.g. 854x480).
  157.    """
  158.     width = height * aspect_ratio
  159.     if only_even:
  160.         width = round(width / 2) * 2
  161.     return int(round(width))
  162.  
  163.  
  164. def is_image(filename: str) -> bool:
  165.     """
  166.    Returns true if a filename refers to an image.
  167.    """
  168.     return mimetypes.types_map.get(os.path.splitext(filename)[-1], "").startswith("image/")
Advertisement
RAW Paste Data Copied
Advertisement