Add color_range to CodecContext/Frame

Change-Id: I42451b3b84ea7ca4ae645566d1ef4709e0326f92

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

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)
 

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)

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,
                     <int**>&inv_tbl,
-                    &src_range,
+                    &src_colorspace_range,
                     <int**>&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

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()

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