diff --git a/av/frame.pxd b/av/frame.pxd index e0d5b42..57fbc47 100644 --- a/av/frame.pxd +++ b/av/frame.pxd @@ -19,3 +19,5 @@ cdef class Frame(object): cdef _copy_internal_attributes(self, Frame source, bint data_layout=?) cdef _init_user_attributes(self) + + cdef _add_cc_data(self, buffer, lib.AVFrameSideDataType type) diff --git a/av/frame.pyx b/av/frame.pyx index 3717ddc..94b3d8b 100644 --- a/av/frame.pyx +++ b/av/frame.pyx @@ -2,7 +2,7 @@ from av.utils cimport avrational_to_fraction, to_avrational from fractions import Fraction -from av.sidedata.sidedata import SideDataContainer +from av.sidedata.sidedata import SideDataContainer, SideData cdef class Frame(object): @@ -66,6 +66,13 @@ cdef class Frame(object): self._time_base = dst + def add(self, buffer, type): + self._add_cc_data(buffer, type) + + cdef _add_cc_data(self, buffer, lib.AVFrameSideDataType type): + data = SideData(self, buffer, type) + # Need to update self._side_data + property dts: """ The decoding timestamp in :attr:`time_base` units for this frame. diff --git a/av/sidedata/sidedata.pxd b/av/sidedata/sidedata.pxd index f30d8fe..d736850 100644 --- a/av/sidedata/sidedata.pxd +++ b/av/sidedata/sidedata.pxd @@ -11,6 +11,7 @@ cdef class SideData(Buffer): cdef Frame frame cdef lib.AVFrameSideData *ptr cdef _Dictionary metadata + cdef int writable cdef SideData wrap_side_data(Frame frame, int index) diff --git a/av/sidedata/sidedata.pyx b/av/sidedata/sidedata.pyx index ec7de59..befd548 100644 --- a/av/sidedata/sidedata.pyx +++ b/av/sidedata/sidedata.pyx @@ -4,6 +4,9 @@ from collections.abc import Mapping from av.sidedata.motionvectors import MotionVectors +from libc.string cimport memcpy + +from av.bytesource cimport ByteSource, bytesource cdef object _cinit_bypass_sentinel = object() @@ -52,6 +55,15 @@ cdef class SideData(Buffer): self.frame = frame self.ptr = frame.ptr.side_data[index] self.metadata = wrap_dictionary(self.ptr.metadata) + self.writable = False + + def __init__(self, Frame frame, bytes, lib.AVFrameSideDataType type): + cdef ByteSource source = bytesource(bytes) + self.frame = frame + self.ptr = lib.av_frame_new_side_data(frame.ptr, type, source.length) + self.metadata = wrap_dictionary(self.ptr.metadata) + self.writable = True + memcpy(self.ptr.data, source.ptr, source.length) cdef size_t _buffer_size(self): return self.ptr.size @@ -60,7 +72,7 @@ cdef class SideData(Buffer): return self.ptr.data cdef bint _buffer_writable(self): - return False + return self.writable def __repr__(self): return f'self.ptr.data:0x}>' @@ -99,6 +111,8 @@ cdef class _SideDataContainer(object): type_ = Type.get(key) return self._by_type[type_] + def add(self, buffer, type): + self.frame.add(buffer, type) class SideDataContainer(_SideDataContainer, Mapping): pass diff --git a/examples/numpy/generate_video.py b/examples/numpy/generate_video.py index 80145d5..939d414 100644 --- a/examples/numpy/generate_video.py +++ b/examples/numpy/generate_video.py @@ -1,15 +1,41 @@ import numpy as np import av +from av.sidedata.sidedata import Type +class CC608Data(): + cc_type_byte = 0xfc # Field 1 + odd_parity = bytearray([0] * 128) + idx = 0 + for idx in range(128): + if idx.bit_count() & 1: + odd_parity[idx] = idx + else: + odd_parity[idx] = idx | 0x80 -duration = 4 + def __init__(self): + self.ba = bytearray() + + def appendPair(self, cc_data_1, cc_data_2): + self.ba.append(self.cc_type_byte) + self.ba.append(self.odd_parity[cc_data_1]) + self.ba.append(self.odd_parity[cc_data_2]) + + def appendString(self, string): + bytes = bytearray(string.encode('ascii')) + if len(bytes) & 1: + bytes.append(0) + for i in range(0, len(bytes), 2): + self.appendPair(bytes[i], bytes[i+1]) + + +duration = 20 fps = 24 total_frames = duration * fps container = av.open("test.mp4", mode="w") -stream = container.add_stream("mpeg4", rate=fps) +stream = container.add_stream("libx264", rate=fps) stream.width = 480 stream.height = 320 stream.pix_fmt = "yuv420p" @@ -24,6 +50,18 @@ for frame_i in range(total_frames): img = np.round(255 * img).astype(np.uint8) frame = av.VideoFrame.from_ndarray(img, format="rgb24") + + if (frame_i % fps) == 0: + # Add pop-on closed caption + data = CC608Data() + data.appendPair(0x14, 0x20) # RCL + data.appendPair(0x14, 0x2e) # ENM + data.appendPair(0x14, 0x60) # Row 15, White + data.appendString("CC1: 00:%02d:%02d" % ((frame_i/(60*fps))%60, (frame_i/fps)%60)) + data.appendPair(0x14, 0x2c) # EDM + data.appendPair(0x14, 0x2f) # EOC + frame.side_data.add(data.ba, Type.A53_CC) + for packet in stream.encode(frame): container.mux(packet) diff --git a/include/libavutil/frame.pxd b/include/libavutil/frame.pxd index aa8dc3a..7dd1d8e 100644 --- a/include/libavutil/frame.pxd +++ b/include/libavutil/frame.pxd @@ -12,3 +12,4 @@ cdef extern from "libavutil/frame.h" nogil: cdef int av_frame_copy(AVFrame *dst, const AVFrame *src) cdef int av_frame_copy_props(AVFrame *dst, const AVFrame *src) cdef AVFrameSideData* av_frame_get_side_data(AVFrame *frame, AVFrameSideDataType type) + cdef AVFrameSideData* av_frame_new_side_data(AVFrame *frame, AVFrameSideDataType type, int size)