mirror of
https://github.com/sunnypilot/sunnypilot.git
synced 2026-07-23 23:02:04 +08:00
Adeeb Shihadeh
262 lines
9.3 KiB
Python
262 lines
9.3 KiB
Python
import ctypes
|
|
import errno
|
|
import os
|
|
|
|
import ffmpeg
|
|
import numpy as np
|
|
|
|
|
|
AV_INPUT_BUFFER_PADDING_SIZE = 64
|
|
AV_LOG_QUIET = -8
|
|
SWS_FAST_BILINEAR = 1
|
|
|
|
|
|
class FFmpegError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class AVPacket(ctypes.Structure):
|
|
# Public prefix of AVPacket. Only data and size are modified here; the packet
|
|
# remains non-refcounted and points at Decoder._packet_buffer.
|
|
_fields_ = [
|
|
("buf", ctypes.c_void_p),
|
|
("pts", ctypes.c_int64),
|
|
("dts", ctypes.c_int64),
|
|
("data", ctypes.POINTER(ctypes.c_uint8)),
|
|
("size", ctypes.c_int),
|
|
]
|
|
|
|
|
|
class AVFrame(ctypes.Structure):
|
|
# Public prefix of AVFrame through format. Stable within a libavutil major
|
|
_fields_ = [
|
|
("data", ctypes.POINTER(ctypes.c_uint8) * 8),
|
|
("linesize", ctypes.c_int * 8),
|
|
("extended_data", ctypes.POINTER(ctypes.POINTER(ctypes.c_uint8))),
|
|
("width", ctypes.c_int),
|
|
("height", ctypes.c_int),
|
|
("nb_samples", ctypes.c_int),
|
|
("format", ctypes.c_int),
|
|
]
|
|
|
|
|
|
def _bind(fn, restype, *argtypes):
|
|
fn.restype = restype
|
|
fn.argtypes = list(argtypes)
|
|
return fn
|
|
|
|
|
|
def _load_libraries():
|
|
avutil = ctypes.CDLL(os.path.join(ffmpeg.LIB_DIR, "libavutil.so.59"), mode=ctypes.RTLD_GLOBAL)
|
|
avcodec = ctypes.CDLL(os.path.join(ffmpeg.LIB_DIR, "libavcodec.so.61"), mode=ctypes.RTLD_GLOBAL)
|
|
swscale = ctypes.CDLL(os.path.join(ffmpeg.LIB_DIR, "libswscale.so.8"), mode=ctypes.RTLD_GLOBAL)
|
|
|
|
c_int, c_char_p, c_void_p, c_size_t = ctypes.c_int, ctypes.c_char_p, ctypes.c_void_p, ctypes.c_size_t
|
|
c_uint8_p = ctypes.POINTER(ctypes.c_uint8)
|
|
c_void_p_p = ctypes.POINTER(c_void_p)
|
|
|
|
_bind(avutil.av_log_set_level, None, c_int)
|
|
_bind(avutil.av_opt_set, c_int, c_void_p, c_char_p, c_char_p, c_int)
|
|
_bind(avutil.av_strerror, c_int, c_int, c_char_p, c_size_t)
|
|
_bind(avutil.av_get_pix_fmt, c_int, c_char_p)
|
|
|
|
_bind(avcodec.avcodec_find_decoder_by_name, c_void_p, c_char_p)
|
|
_bind(avcodec.avcodec_alloc_context3, c_void_p, c_void_p)
|
|
_bind(avcodec.avcodec_open2, c_int, c_void_p, c_void_p, c_void_p)
|
|
_bind(avcodec.avcodec_free_context, None, c_void_p_p)
|
|
_bind(avcodec.avcodec_flush_buffers, None, c_void_p)
|
|
_bind(avcodec.avcodec_send_packet, c_int, c_void_p, ctypes.POINTER(AVPacket))
|
|
_bind(avcodec.avcodec_receive_frame, c_int, c_void_p, ctypes.POINTER(AVFrame))
|
|
_bind(avcodec.av_packet_alloc, ctypes.POINTER(AVPacket))
|
|
_bind(avcodec.av_packet_free, None, ctypes.POINTER(ctypes.POINTER(AVPacket)))
|
|
_bind(avcodec.av_frame_alloc, ctypes.POINTER(AVFrame))
|
|
_bind(avcodec.av_frame_free, None, ctypes.POINTER(ctypes.POINTER(AVFrame)))
|
|
_bind(avcodec.av_frame_unref, None, ctypes.POINTER(AVFrame))
|
|
|
|
_bind(swscale.sws_getCachedContext, c_void_p,
|
|
c_void_p, c_int, c_int, c_int, c_int, c_int, c_int, c_int, c_void_p, c_void_p, c_void_p)
|
|
_bind(swscale.sws_scale, c_int,
|
|
c_void_p, ctypes.POINTER(c_uint8_p), ctypes.POINTER(c_int),
|
|
c_int, c_int, ctypes.POINTER(c_uint8_p), ctypes.POINTER(c_int))
|
|
_bind(swscale.sws_freeContext, None, c_void_p)
|
|
|
|
avutil.av_log_set_level(AV_LOG_QUIET)
|
|
return avutil, avcodec, swscale
|
|
_avutil, _avcodec, _swscale = _load_libraries()
|
|
|
|
_DataArray = ctypes.POINTER(ctypes.c_uint8) * 4
|
|
_LinesizeArray = ctypes.c_int * 4
|
|
|
|
AV_PIX_FMT_NV12 = _avutil.av_get_pix_fmt(b"nv12")
|
|
assert AV_PIX_FMT_NV12 >= 0
|
|
|
|
|
|
def _error_string(code: int) -> str:
|
|
buf = ctypes.create_string_buffer(256)
|
|
if _avutil.av_strerror(code, buf, len(buf)) == 0:
|
|
return buf.value.decode(errors="replace")
|
|
return f"FFmpeg error {code}"
|
|
|
|
|
|
def _check(code: int, operation: str) -> None:
|
|
if code < 0:
|
|
raise FFmpegError(f"{operation}: {_error_string(code)}")
|
|
|
|
|
|
class Decoder:
|
|
def __init__(self, codec_name: str = "hevc"):
|
|
self.closed = True
|
|
self._sws_context = ctypes.c_void_p()
|
|
self._packet_buffer = bytearray()
|
|
self._packet_address = 0
|
|
self._packet_data = None
|
|
self._output = np.empty(0, dtype=np.uint8)
|
|
self._dst_data = _DataArray()
|
|
self._dst_linesize = _LinesizeArray()
|
|
self.width = 0
|
|
self.height = 0
|
|
self._src_format = -1
|
|
|
|
codec = _avcodec.avcodec_find_decoder_by_name(codec_name.encode())
|
|
if not codec:
|
|
raise FFmpegError(f"decoder not found: {codec_name}")
|
|
|
|
self._context = ctypes.c_void_p(_avcodec.avcodec_alloc_context3(codec))
|
|
if not self._context:
|
|
raise MemoryError("avcodec_alloc_context3 failed")
|
|
|
|
self._packet = _avcodec.av_packet_alloc()
|
|
if not self._packet:
|
|
_avcodec.avcodec_free_context(ctypes.byref(self._context))
|
|
raise MemoryError("av_packet_alloc failed")
|
|
|
|
self._frame = _avcodec.av_frame_alloc()
|
|
if not self._frame:
|
|
_avcodec.av_packet_free(ctypes.byref(self._packet))
|
|
_avcodec.avcodec_free_context(ctypes.byref(self._context))
|
|
raise MemoryError("av_frame_alloc failed")
|
|
|
|
try:
|
|
# Frame threading holds decoded frames to populate worker pipelines.
|
|
# Slice threads can reduce decode time without adding that frame queue;
|
|
# four was the latency minimum on the replay camera workload.
|
|
_check(_avutil.av_opt_set(self._context, b"threads", b"4", 0), "set decoder threads")
|
|
_check(_avutil.av_opt_set(self._context, b"thread_type", b"slice", 0), "set decoder thread type")
|
|
_check(_avutil.av_opt_set(self._context, b"flags", b"+low_delay", 0), "set low-delay mode")
|
|
_check(_avcodec.avcodec_open2(self._context, codec, None), "open decoder")
|
|
except Exception:
|
|
_avcodec.av_frame_free(ctypes.byref(self._frame))
|
|
_avcodec.av_packet_free(ctypes.byref(self._packet))
|
|
_avcodec.avcodec_free_context(ctypes.byref(self._context))
|
|
raise
|
|
|
|
self.closed = False
|
|
|
|
def __enter__(self):
|
|
self._ensure_open()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
self.close()
|
|
|
|
def _ensure_open(self) -> None:
|
|
if self.closed:
|
|
raise RuntimeError("decoder is closed")
|
|
|
|
def _prepare_packet(self, data) -> None:
|
|
size = len(data)
|
|
required = size + AV_INPUT_BUFFER_PADDING_SIZE
|
|
if len(self._packet_buffer) < required:
|
|
# Grow-only; the address stays valid until the next reallocation.
|
|
self._packet_buffer = bytearray(required)
|
|
self._packet_address = ctypes.addressof(ctypes.c_uint8.from_buffer(self._packet_buffer))
|
|
self._packet_data = ctypes.cast(self._packet_address, ctypes.POINTER(ctypes.c_uint8))
|
|
|
|
self._packet_buffer[:size] = data
|
|
ctypes.memset(self._packet_address + size, 0, AV_INPUT_BUFFER_PADDING_SIZE)
|
|
self._packet.contents.data = self._packet_data
|
|
self._packet.contents.size = size
|
|
|
|
def _prepare_output(self, frame: AVFrame) -> None:
|
|
width, height, src_format = frame.width, frame.height, frame.format
|
|
if width <= 0 or height <= 0 or width % 2 or height % 2:
|
|
raise FFmpegError(f"unsupported frame dimensions: {width}x{height}")
|
|
if (width, height, src_format) == (self.width, self.height, self._src_format):
|
|
return
|
|
|
|
sws_context = _swscale.sws_getCachedContext(
|
|
self._sws_context, width, height, src_format,
|
|
width, height, AV_PIX_FMT_NV12, SWS_FAST_BILINEAR,
|
|
None, None, None,
|
|
)
|
|
if not sws_context:
|
|
raise FFmpegError("sws_getCachedContext failed")
|
|
self._sws_context = ctypes.c_void_p(sws_context)
|
|
|
|
self.width, self.height = width, height
|
|
self._src_format = src_format
|
|
self._output = np.empty(width * height * 3 // 2, dtype=np.uint8)
|
|
output_address = self._output.ctypes.data
|
|
self._dst_data = _DataArray(
|
|
ctypes.cast(output_address, ctypes.POINTER(ctypes.c_uint8)),
|
|
ctypes.cast(output_address + width * height, ctypes.POINTER(ctypes.c_uint8)),
|
|
None,
|
|
None,
|
|
)
|
|
self._dst_linesize = _LinesizeArray(width, width, 0, 0)
|
|
|
|
def _receive(self) -> np.ndarray | None:
|
|
"""Return one NV12 frame, or None if the decoder needs more input.
|
|
|
|
The returned buffer is reused on the next successful decode; callers must
|
|
use or copy it before calling decode again.
|
|
"""
|
|
result = _avcodec.avcodec_receive_frame(self._context, self._frame)
|
|
if result == -errno.EAGAIN:
|
|
return None
|
|
_check(result, "receive decoded frame")
|
|
|
|
try:
|
|
frame = self._frame.contents
|
|
self._prepare_output(frame)
|
|
rows = _swscale.sws_scale(
|
|
self._sws_context, frame.data, frame.linesize, 0, frame.height,
|
|
self._dst_data, self._dst_linesize,
|
|
)
|
|
if rows != frame.height:
|
|
raise FFmpegError(f"convert decoded frame: produced {rows} of {frame.height} rows")
|
|
return self._output
|
|
finally:
|
|
_avcodec.av_frame_unref(self._frame)
|
|
|
|
def decode(self, data) -> np.ndarray | None:
|
|
self._ensure_open()
|
|
if len(data) == 0:
|
|
return None
|
|
|
|
self._prepare_packet(data)
|
|
result = _avcodec.avcodec_send_packet(self._context, self._packet)
|
|
# The packet buffer is ours, not FFmpeg's. Clear the borrowed pointer so
|
|
# packet teardown can never attempt to release it.
|
|
self._packet.contents.data = None
|
|
self._packet.contents.size = 0
|
|
_check(result, "send packet to decoder")
|
|
return self._receive()
|
|
|
|
def reset(self) -> None:
|
|
"""Discard decoder state after a stream discontinuity."""
|
|
self._ensure_open()
|
|
_avcodec.avcodec_flush_buffers(self._context)
|
|
|
|
def close(self) -> None:
|
|
if self.closed:
|
|
return
|
|
self.closed = True
|
|
_swscale.sws_freeContext(self._sws_context)
|
|
_avcodec.av_frame_free(ctypes.byref(self._frame))
|
|
_avcodec.av_packet_free(ctypes.byref(self._packet))
|
|
_avcodec.avcodec_free_context(ctypes.byref(self._context))
|
|
|
|
def __del__(self):
|
|
self.close()
|