Skip to content

Commit

Permalink
use PyBUF
Browse files Browse the repository at this point in the history
  • Loading branch information
lim committed Jun 5, 2023
1 parent 6d99555 commit 385b209
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 64 deletions.
Binary file modified extension/recorder.cp310-win_amd64.pyd
View file
Binary file not shown.
Binary file modified extension/recorder.cpython-310-x86_64-linux-gnu.so
View file
Binary file not shown.
20 changes: 1 addition & 19 deletions extension/recorder.pxd
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from libc.stdint cimport int64_t, uint64_t, uint32_t, uint8_t, UINT_LEAST32_MAX
from libc.string cimport memcpy, memmove
from libc.string cimport memcpy, memmove, strcmp
from libc.stdlib cimport malloc, free
from libc.time cimport time
from cpython.buffer cimport PyBUF_SIMPLE, PyBuffer_Release, PyObject_CheckBuffer, PyObject_GetBuffer
Expand Down Expand Up @@ -166,28 +166,10 @@ cdef class Recorder(object):
cdef AVCodecContext *audio_codec_ctx
cdef video_packet_merger merger

cdef void init_packet(self, AVPacket * packet, uint64_t pts, int length, const uint8_t *data)

cdef void packet_merger_init(self)

cdef void packet_merger_merge(self, AVPacket *packet)

cdef void packet_merger_destroy(self)

cdef AVCodecID get_avcodec_id(self, char *codec_name)

cpdef bint add_video_stream(self, char *codec_name, int width, int height) except False

cpdef bint add_audio_stream(self, char *codec_name) except False

cpdef bint write_video_header(self, uint64_t pts, int length, const uint8_t *data) except False

cpdef bint write_audio_header(self, uint64_t pts, int length, const uint8_t *data) except False

cpdef bint write_header(self) except False

cpdef bint write_video_packet(self, uint64_t pts, int length, const uint8_t *data) except False

cpdef bint write_audio_packet(self, uint64_t pts, int length, const uint8_t *data) except False

cpdef int close_container(self) except 0
85 changes: 40 additions & 45 deletions extension/recorder.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,23 +5,15 @@ cdef AVRational RECORD_TIME_BASE = {"num":1, "den":1000000}


cdef class ByteSource(object):
def __cinit__(self, owner):
self.owner = owner
try:
self.ptr = owner
except TypeError:
pass
else:
self.length = len(owner)
def __cinit__(self, data):
self.ptr = data
# Can very likely use PyBUF_ND instead of PyBUF_SIMPLE
res = PyObject_GetBuffer(data, &self.view, PyBUF_SIMPLE)
if not res:
self.has_view = True
self.ptr = <unsigned char *>self.view.buf
self.length = self.view.len
return
if PyObject_CheckBuffer(owner):
# Can very likely use PyBUF_ND instead of PyBUF_SIMPLE
res = PyObject_GetBuffer(owner, &self.view, PyBUF_SIMPLE)
if not res:
self.has_view = True
self.ptr = <unsigned char *>self.view.buf
self.length = self.view.len
return
raise TypeError('expected bytes, bytearray or memoryview')

def __dealloc__(self):
Expand Down Expand Up @@ -63,17 +55,6 @@ cdef class Recorder(object):
def finish_time(self):
return self.finish_time

cdef void init_packet(self, AVPacket * packet, uint64_t pts, int length, const uint8_t *data):
av_new_packet(packet, length)
memcpy(packet.data, data, packet.size)
if (SC_PACKET_FLAG_CONFIG & pts):
packet.pts = AV_NOPTS_VALUE
else:
packet.pts = pts & SC_PACKET_PTS_MASK
if pts & SC_PACKET_FLAG_KEY_FRAME:
packet.flags |= AV_PKT_FLAG_KEY
packet.dts = packet.pts

cdef void packet_merger_init(self):
self.merger.config = NULL

Expand All @@ -99,22 +80,22 @@ cdef class Recorder(object):
free(self.merger.config)

cdef AVCodecID get_avcodec_id(self, char *codec_name,):
if codec_name == "h264":
if strcmp(codec_name, "h264")==0:
return AVCodecID.AV_CODEC_ID_H264
elif codec_name == "h265":
elif strcmp(codec_name, "h265")==0:
return AVCodecID.AV_CODEC_ID_HEVC
elif codec_name == "av1":
elif strcmp(codec_name, "av1")==0:
return AVCodecID.AV_CODEC_ID_AV1
elif codec_name == "opus":
elif strcmp(codec_name, "opus")==0:
return AVCodecID.AV_CODEC_ID_OPUS
elif codec_name == "aac":
elif strcmp(codec_name, "aac")==0:
return AVCodecID.AV_CODEC_ID_AAC
elif codec_name == "raw":
elif strcmp(codec_name, "raw")==0:
return AVCodecID.AV_CODEC_ID_PCM_S16LE
else:
return AVCodecID.AV_CODEC_ID_NONE

cpdef bint add_video_stream(self, char *codec_name, int width, int height) except False:
def add_video_stream(self, char *codec_name, int width, int height):
# 1.get codec id
cdef AVCodecID video_codec_id = self.get_avcodec_id(codec_name)
# 2.get codec and set codec_ctx
Expand All @@ -131,7 +112,7 @@ cdef class Recorder(object):
avcodec_parameters_from_context(video_stream.codecpar, self.video_codec_ctx)
return True

cpdef bint add_audio_stream(self, char *codec_name) except False:
def add_audio_stream(self, char *codec_name):
# 1.get codec id
cdef AVCodecID audio_codec_id = self.get_avcodec_id(codec_name)
# 2.get codec and set codec_ctx
Expand All @@ -148,10 +129,24 @@ cdef class Recorder(object):
avcodec_parameters_from_context(audio_stream.codecpar, self.audio_codec_ctx)
return True

cpdef bint write_video_header(self, uint64_t pts, int length, const uint8_t *data) except False:
def init_packet(self, bint is_video, uint64_t pts, int length, object data):
packet = self.video_packet if is_video else self.audio_packet
av_new_packet(packet, length)
byte_source = ByteSource(data)
memcpy(packet.data, byte_source.ptr, packet.size)
del byte_source
if (SC_PACKET_FLAG_CONFIG & pts):
packet.pts = AV_NOPTS_VALUE
else:
packet.pts = pts & SC_PACKET_PTS_MASK
if pts & SC_PACKET_FLAG_KEY_FRAME:
packet.flags |= AV_PKT_FLAG_KEY
packet.dts = packet.pts

def write_video_header(self, uint64_t pts, int length, object data):
self.video_packet = av_packet_alloc()
self.packet_merger_init()
self.init_packet(self.video_packet, pts, length, data)
self.init_packet(True, pts, length, data)
self.packet_merger_merge(self.video_packet)
cdef uint8_t *extradata = <uint8_t *> av_malloc(self.video_packet.size * sizeof(uint8_t))
memcpy(extradata, self.video_packet.data, self.video_packet.size)
Expand All @@ -160,22 +155,22 @@ cdef class Recorder(object):
av_packet_free(&self.video_packet)
return True

cpdef bint write_audio_header(self, uint64_t pts, int length, const uint8_t *data) except False:
def write_audio_header(self, uint64_t pts, int length, object data):
self.audio_packet = av_packet_alloc()
self.init_packet(self.audio_packet, pts, length, data)
self.init_packet(False, pts, length, data)
cdef uint8_t *extradata = <uint8_t *> av_malloc(self.audio_packet.size * sizeof(uint8_t))
memcpy(extradata, self.audio_packet.data, self.audio_packet.size)
self.container.streams[1].codecpar.extradata = extradata
self.container.streams[1].codecpar.extradata_size = self.audio_packet.size
av_packet_free(&self.audio_packet)
return True

cpdef bint write_header(self) except False:
def write_header(self):
return (avformat_write_header(self.container, NULL) >= 0)

cpdef bint write_video_packet(self, uint64_t pts, int length, const uint8_t *data) except False:
def write_video_packet(self, uint64_t pts, int length, object data):
self.video_packet = av_packet_alloc()
self.init_packet(self.video_packet, pts, length, data)
self.init_packet(True, pts, length, data)
if self.pts_origin == AV_NOPTS_VALUE:
self.pts_origin = self.video_packet.pts
# config packet
Expand All @@ -198,9 +193,9 @@ cdef class Recorder(object):
self.video_packet = NULL
return True

cpdef bint write_audio_packet(self, uint64_t pts, int length, const uint8_t *data) except False:
def write_audio_packet(self, uint64_t pts, int length, object data):
self.audio_packet = av_packet_alloc()
self.init_packet(self.audio_packet, pts, length, data)
self.init_packet(False, pts, length, data)
if self.pts_origin == AV_NOPTS_VALUE:
self.pts_origin = self.audio_packet.pts
self.audio_packet.stream_index = 1
Expand All @@ -212,7 +207,7 @@ cdef class Recorder(object):
av_packet_free(&self.video_packet)
return True

cpdef int close_container(self) except 0:
def close_container(self):
cdef int duration
if self.previous_video_packet:
self.previous_video_packet.duration = 100000
Expand Down

0 comments on commit 385b209

Please sign in to comment.