Skip to content

Commit

Permalink
Add color_range to CodecContext/Frame
Browse files Browse the repository at this point in the history
  • Loading branch information
johanjeppsson authored Nov 28, 2023
1 parent 1de6a50 commit 8aa5fe7
Show file tree
Hide file tree
Showing 7 changed files with 121 additions and 19 deletions.
13 changes: 13 additions & 0 deletions av/video/codeccontext.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,19 @@ cdef class VideoCodecContext(CodecContext):
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

property max_b_frames:
"""
The maximum run of consecutive B frames when encoding a video.
Expand Down
26 changes: 26 additions & 0 deletions av/video/frame.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,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)
Expand Down
3 changes: 2 additions & 1 deletion av/video/reformatter.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -9,4 +9,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)
55 changes: 37 additions & 18 deletions av/video/reformatter.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ Interpolation = define_enum('Interpolation', __name__, (
))

Colorspace = define_enum('Colorspace', __name__, (

('ITU709', lib.SWS_CS_ITU709),
('FCC', lib.SWS_CS_FCC),
('ITU601', lib.SWS_CS_ITU601),
Expand All @@ -41,9 +40,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:

"""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
Expand All @@ -57,7 +61,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.
Expand All @@ -66,19 +71,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,
Expand All @@ -88,11 +99,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.")
Expand All @@ -104,7 +118,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

Expand All @@ -126,24 +141,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,
<int**>&inv_tbl,
&src_range,
&src_colorspace_range,
<int**>&tbl,
&dst_range,
&dst_colorspace_range,
&brightness,
&contrast,
&saturation
Expand All @@ -164,9 +183,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
Expand Down
3 changes: 3 additions & 0 deletions include/libavcodec/avcodec.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,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
Expand Down Expand Up @@ -317,6 +318,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()

Expand Down
20 changes: 20 additions & 0 deletions include/libavutil/avutil.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,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)
Expand Down
20 changes: 20 additions & 0 deletions tests/test_colorspace.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 8aa5fe7

Please sign in to comment.