From e636ccbdfa41de33786e3c2c61dcff0ba953962b Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Thu, 9 Nov 2023 09:52:48 -0500 Subject: [PATCH] Add color_range to CodecContext/Frame (#17) Co-authored-by: Johan Jeppsson --- av/video/codeccontext.pyx | 13 ++++++++ av/video/frame.pyx | 26 ++++++++++++++++ av/video/reformatter.pxd | 3 +- av/video/reformatter.pyx | 56 +++++++++++++++++++++++----------- include/libavcodec/avcodec.pxd | 3 ++ include/libavutil/avutil.pxd | 20 ++++++++++++ 6 files changed, 103 insertions(+), 18 deletions(-) diff --git a/av/video/codeccontext.pyx b/av/video/codeccontext.pyx index 891f5ccaf..f0f5d2edd 100644 --- a/av/video/codeccontext.pyx +++ b/av/video/codeccontext.pyx @@ -169,3 +169,16 @@ cdef class VideoCodecContext(CodecContext): property coded_height: def __get__(self): return self.ptr.coded_height + + @property + def color_range(self): + """ + Color range of context. + + Wraps :ffmpeg:`AVFrame.color_range`. + """ + def __get__(self): + return self.ptr.color_range + + def __set__(self, value): + self.ptr.color_range = value diff --git a/av/video/frame.pyx b/av/video/frame.pyx index 3e058961c..768046cf9 100644 --- a/av/video/frame.pyx +++ b/av/video/frame.pyx @@ -186,6 +186,32 @@ cdef class VideoFrame(Frame): def pict_type(self, value): self.ptr.pict_type = PictureType[value].value + @property + def colorspace(self): + """Colorspace of frame. + + Wraps :ffmpeg:`AVFrame.colorspace`. + + """ + return self.ptr.colorspace + + @colorspace.setter + def colorspace(self, value): + self.ptr.colorspace = value + + @property + def color_range(self): + """Color range of frame. + + Wraps :ffmpeg:`AVFrame.color_range`. + + """ + return self.ptr.color_range + + @color_range.setter + def color_range(self, value): + self.ptr.color_range = value + def reformat(self, *args, **kwargs): """reformat(width=None, height=None, format=None, src_colorspace=None, dst_colorspace=None, interpolation=None) diff --git a/av/video/reformatter.pxd b/av/video/reformatter.pxd index 579402046..c256c3eee 100644 --- a/av/video/reformatter.pxd +++ b/av/video/reformatter.pxd @@ -8,4 +8,5 @@ cdef class VideoReformatter: cdef _reformat(self, VideoFrame frame, int width, int height, lib.AVPixelFormat format, int src_colorspace, - int dst_colorspace, int interpolation) + int dst_colorspace, int interpolation, + int src_color_range, int dst_color_range) diff --git a/av/video/reformatter.pyx b/av/video/reformatter.pyx index 1d3f08065..69fc50292 100644 --- a/av/video/reformatter.pyx +++ b/av/video/reformatter.pyx @@ -42,8 +42,15 @@ Colorspace = define_enum('Colorspace', __name__, ( )) -cdef class VideoReformatter: +ColorRange = define_enum('ColorRange', __name__, ( + ('UNSPECIFIED', lib.AVCOL_RANGE_UNSPECIFIED, "Unspecified"), + ('MPEG', lib.AVCOL_RANGE_MPEG, "MPEG (limited) YUV range, 219*2^(n-8)"), + ('JPEG', lib.AVCOL_RANGE_JPEG, "JPEG (full) YUV range, 2^n-1"), + ('NB', lib.AVCOL_RANGE_NB, "Not part of ABI"), +)) + +cdef class VideoReformatter: """An object for reformatting size and pixel format of :class:`.VideoFrame`. It is most efficient to have a reformatter object for each set of parameters @@ -57,7 +64,8 @@ cdef class VideoReformatter: def reformat(self, VideoFrame frame not None, width=None, height=None, format=None, src_colorspace=None, dst_colorspace=None, - interpolation=None): + interpolation=None, src_color_range=None, + dst_color_range=None): """Create a new :class:`VideoFrame` with the given width/height/format/colorspace. Returns the same frame untouched if nothing needs to be done to it. @@ -66,19 +74,25 @@ cdef class VideoReformatter: :param int height: New height, or ``None`` for the same height. :param format: New format, or ``None`` for the same format. :type format: :class:`.VideoFormat` or ``str`` - :param src_colorspace: Current colorspace, or ``None`` for ``DEFAULT``. + :param src_colorspace: Current colorspace, or ``None`` for the frame colorspace. :type src_colorspace: :class:`Colorspace` or ``str`` - :param dst_colorspace: Desired colorspace, or ``None`` for ``DEFAULT``. + :param dst_colorspace: Desired colorspace, or ``None`` for the frame colorspace. :type dst_colorspace: :class:`Colorspace` or ``str`` :param interpolation: The interpolation method to use, or ``None`` for ``BILINEAR``. :type interpolation: :class:`Interpolation` or ``str`` + :param src_color_range: Current color range, or ``None`` for the frame color range. + :type src_color_range: :class:`color range` or ``str`` + :param dst_color_range: Desired color range, or ``None`` for the frame color range. + :type dst_color_range: :class:`color range` or ``str`` """ cdef VideoFormat video_format = VideoFormat(format if format is not None else frame.format) - cdef int c_src_colorspace = (Colorspace[src_colorspace] if src_colorspace is not None else Colorspace.DEFAULT).value - cdef int c_dst_colorspace = (Colorspace[dst_colorspace] if dst_colorspace is not None else Colorspace.DEFAULT).value + cdef int c_src_colorspace = (Colorspace[src_colorspace].value if src_colorspace is not None else frame.colorspace) + cdef int c_dst_colorspace = (Colorspace[dst_colorspace].value if dst_colorspace is not None else frame.colorspace) cdef int c_interpolation = (Interpolation[interpolation] if interpolation is not None else Interpolation.BILINEAR).value + cdef int c_src_color_range = (ColorRange[src_color_range].value if src_color_range is not None else frame.color_range) + cdef int c_dst_color_range = (ColorRange[dst_color_range].value if dst_color_range is not None else frame.color_range) return self._reformat( frame, @@ -88,11 +102,14 @@ cdef class VideoReformatter: c_src_colorspace, c_dst_colorspace, c_interpolation, + c_src_color_range, + c_dst_color_range, ) cdef _reformat(self, VideoFrame frame, int width, int height, lib.AVPixelFormat dst_format, int src_colorspace, - int dst_colorspace, int interpolation): + int dst_colorspace, int interpolation, + int src_color_range, int dst_color_range): if frame.ptr.format < 0: raise ValueError("Frame does not have format set.") @@ -104,7 +121,8 @@ cdef class VideoReformatter: dst_format == src_format and width == frame.ptr.width and height == frame.ptr.height and - dst_colorspace == src_colorspace + dst_colorspace == src_colorspace and + src_color_range == dst_color_range ): return frame @@ -126,24 +144,28 @@ cdef class VideoReformatter: NULL ) - # We want to change the colorspace transforms. We do that by grabbing - # all of the current settings, changing a couple, and setting them all. - # We need a lot of state here. + # We want to change the colorspace/color_range transforms. + # We do that by grabbing all of the current settings, changing a + # couple, and setting them all. We need a lot of state here. cdef const int *inv_tbl cdef const int *tbl - cdef int src_range, dst_range, brightness, contrast, saturation + cdef int src_colorspace_range, dst_colorspace_range + cdef int brightness, contrast, saturation cdef int ret - if src_colorspace != dst_colorspace: + if ( + src_colorspace != dst_colorspace or + src_color_range != dst_color_range + ): with nogil: # Casts for const-ness, because Cython isn't expressive enough. ret = lib.sws_getColorspaceDetails( self.ptr, &inv_tbl, - &src_range, + &src_colorspace_range, &tbl, - &dst_range, + &dst_colorspace_range, &brightness, &contrast, &saturation @@ -164,9 +186,9 @@ cdef class VideoReformatter: ret = lib.sws_setColorspaceDetails( self.ptr, inv_tbl, - src_range, + src_color_range, tbl, - dst_range, + dst_color_range, brightness, contrast, saturation diff --git a/include/libavcodec/avcodec.pxd b/include/libavcodec/avcodec.pxd index c63a7cc4b..71e9186c4 100644 --- a/include/libavcodec/avcodec.pxd +++ b/include/libavcodec/avcodec.pxd @@ -202,6 +202,7 @@ cdef extern from "libavcodec/avcodec.h" nogil: int gop_size # The number of pictures in a group of pictures, or 0 for intra_only. int max_b_frames int has_b_frames + AVColorRange color_range # Audio. AVSampleFormat sample_fmt @@ -355,6 +356,8 @@ cdef extern from "libavcodec/avcodec.h" nogil: AVDictionary *metadata int flags int decode_error_flags + AVColorRange color_range + AVColorSpace colorspace cdef AVFrame* avcodec_alloc_frame() diff --git a/include/libavutil/avutil.pxd b/include/libavutil/avutil.pxd index 27ade47b1..fbf6235eb 100644 --- a/include/libavutil/avutil.pxd +++ b/include/libavutil/avutil.pxd @@ -35,6 +35,26 @@ cdef extern from "libavutil/avutil.h" nogil: # This is nice, but only in FFMpeg: # AV_ROUND_PASS_MINMAX + cdef enum AVColorSpace: + AVCOL_SPC_RGB + AVCOL_SPC_BT709 + AVCOL_SPC_UNSPECIFIED + AVCOL_SPC_RESERVED + AVCOL_SPC_FCC + AVCOL_SPC_BT470BG + AVCOL_SPC_SMPTE170M + AVCOL_SPC_SMPTE240M + AVCOL_SPC_YCOCG + AVCOL_SPC_BT2020_NCL + AVCOL_SPC_BT2020_CL + AVCOL_SPC_NB + + cdef enum AVColorRange: + AVCOL_RANGE_UNSPECIFIED + AVCOL_RANGE_MPEG + AVCOL_RANGE_JPEG + AVCOL_RANGE_NB + cdef double M_PI cdef void* av_malloc(size_t size)