From cfac5fd0699516279457b15cfaa60bbe9ca2c326 Mon Sep 17 00:00:00 2001 From: Johan Jeppsson Date: Mon, 6 Jul 2020 20:48:13 +0200 Subject: [PATCH 1/2] Add color_range to CodecContext/Frame Change-Id: I42451b3b84ea7ca4ae645566d1ef4709e0326f92 --- av/video/codeccontext.pyx | 13 ++++++++ av/video/frame.pyx | 26 ++++++++++++++++ av/video/reformatter.pxd | 3 +- av/video/reformatter.pyx | 55 ++++++++++++++++++++++++---------- include/libavcodec/avcodec.pxd | 3 ++ include/libavutil/avutil.pxd | 19 ++++++++++++ 6 files changed, 102 insertions(+), 17 deletions(-) diff --git a/av/video/codeccontext.pyx b/av/video/codeccontext.pyx index f5eb9d07b..47857de25 100644 --- a/av/video/codeccontext.pyx +++ b/av/video/codeccontext.pyx @@ -158,3 +158,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 62601ce0e..c2b299e77 100644 --- a/av/video/frame.pyx +++ b/av/video/frame.pyx @@ -184,6 +184,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 25135c27a..e6d5367d3 100644 --- a/av/video/reformatter.pxd +++ b/av/video/reformatter.pxd @@ -9,4 +9,5 @@ cdef class VideoReformatter(object): 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 3ad995fb9..a5f10916a 100644 --- a/av/video/reformatter.pyx +++ b/av/video/reformatter.pyx @@ -42,6 +42,14 @@ Colorspace = define_enum('Colorspace', __name__, ( )) +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(object): """An object for reformatting size and pixel format of :class:`.VideoFrame`. @@ -57,7 +65,8 @@ cdef class VideoReformatter(object): 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 +75,25 @@ cdef class VideoReformatter(object): :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 +103,14 @@ cdef class VideoReformatter(object): 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 +122,8 @@ cdef class VideoReformatter(object): 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 +145,28 @@ cdef class VideoReformatter(object): 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 +187,9 @@ cdef class VideoReformatter(object): 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 92d860d65..116e9c45a 100644 --- a/include/libavcodec/avcodec.pxd +++ b/include/libavcodec/avcodec.pxd @@ -206,6 +206,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 @@ -319,6 +320,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 fa1bdf203..5c789f722 100644 --- a/include/libavutil/avutil.pxd +++ b/include/libavutil/avutil.pxd @@ -39,6 +39,25 @@ 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 From fecb65b7dc8e090be5d83c16399b90cc378afc2b Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Tue, 28 Nov 2023 02:19:41 -0500 Subject: [PATCH 2/2] Add unit test --- tests/test_colorspace.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 tests/test_colorspace.py diff --git a/tests/test_colorspace.py b/tests/test_colorspace.py new file mode 100644 index 000000000..13784fec3 --- /dev/null +++ b/tests/test_colorspace.py @@ -0,0 +1,20 @@ +from av.video.reformatter import ColorRange, Colorspace +import av + +from .common import TestCase, fate_suite + + +class TestColorSpace(TestCase): + def test_color_range(self): + container = av.open( + fate_suite("amv/MTV_high_res_320x240_sample_Penguin_Joke_MTV_from_WMV.amv") + ) + stream = container.streams.video[0] + + self.assertEqual(stream.color_range, None) + + for packet in container.demux(stream): + for frame in packet.decode(): + self.assertEqual(frame.color_range, ColorRange.JPEG) # a.k.a "pc" + self.assertEqual(frame.colorspace, Colorspace.ITU601) + return