Initial libaom encoder

This commit is contained in:
DataHoarder 2022-09-15 15:36:41 +02:00
parent 4898ae67cf
commit 52147f9138
Signed by: DataHoarder
SSH key fingerprint: SHA256:EnPQOqPpbCa7nzalCEJY2sd9iPluFIBuJu2rDFalACI
3 changed files with 302 additions and 0 deletions

25
encoder/libaom/libaom.c Normal file
View file

@ -0,0 +1,25 @@
#include "libaom.h"
aom_codec_err_t aom_codec_control_int(aom_codec_ctx_t *ctx, int ctrl_id, int v) {
return aom_codec_control(ctx, ctrl_id, v);
}
aom_codec_err_t aom_codec_control_uint(aom_codec_ctx_t *ctx, int ctrl_id, unsigned int v) {
return aom_codec_control(ctx, ctrl_id, v);
}
void* aom_get_pkt_buf(aom_codec_cx_pkt_t *pkt){
return pkt->data.frame.buf;
}
size_t aom_get_pkt_sz(aom_codec_cx_pkt_t *pkt){
return pkt->data.frame.sz;
}
aom_codec_pts_t aom_get_pkt_pts(aom_codec_cx_pkt_t *pkt){
return pkt->data.frame.pts;
}
aom_codec_frame_flags_t aom_get_pkt_flags(aom_codec_cx_pkt_t *pkt){
return pkt->data.frame.flags;
}

262
encoder/libaom/libaom.go Normal file
View file

@ -0,0 +1,262 @@
package libaom
/*
#cgo pkg-config: aom
#include "libaom.h"
*/
import "C"
import (
"encoding/binary"
"errors"
"fmt"
"git.gammaspectra.live/S.O.N.G/Ignite/frame"
"git.gammaspectra.live/S.O.N.G/Ignite/y4m"
"io"
"runtime"
"strconv"
"sync/atomic"
"unsafe"
)
type Encoder struct {
w io.Writer
cleaned atomic.Bool
cfg C.aom_codec_enc_cfg_t
codec C.aom_codec_ctx_t
raw *C.aom_image_t
}
var libaomVersion = "libaom-av1 " + C.GoString(C.aom_codec_version_str())
func Version() string {
return libaomVersion
}
func NewEncoder(w io.Writer, stream *y4m.Stream, settings map[string]any) (*Encoder, error) {
e := &Encoder{
w: w,
}
var aomErr C.aom_codec_err_t
encoder := C.aom_codec_av1_cx()
if encoder == nil {
return nil, errors.New("unsupported codec")
}
var usage C.uint = C.AOM_USAGE_GOOD_QUALITY
var speed C.int = 8
var bitrate C.uint = 800
var imageFormat C.aom_img_fmt_t
if aomErr = C.aom_codec_enc_config_default(encoder, &e.cfg, usage); aomErr != 0 {
return nil, errors.New("failed to get default codec config")
}
switch true {
case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 4 && stream.ColorSpace().Subsampling.B == 4:
imageFormat = C.AOM_IMG_FMT_I444
e.cfg.g_profile = 1
case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 2 && stream.ColorSpace().Subsampling.B == 2:
imageFormat = C.AOM_IMG_FMT_I422
e.cfg.g_profile = 2
case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 2 && stream.ColorSpace().Subsampling.B == 0:
imageFormat = C.AOM_IMG_FMT_I420
e.cfg.g_profile = 0
default:
return nil, errors.New("unsupported input chroma subsampling")
}
if stream.ColorSpace().BitDepth > 8 {
imageFormat |= C.AOM_IMG_FMT_HIGHBITDEPTH
}
e.cfg.g_input_bit_depth = C.uint(stream.ColorSpace().BitDepth)
e.cfg.g_bit_depth = C.aom_bit_depth_t(stream.ColorSpace().BitDepth)
if e.cfg.g_bit_depth >= 12 {
e.cfg.g_bit_depth = 12
e.cfg.g_profile = 2
}
width, height := stream.Resolution()
if e.raw = (*C.aom_image_t)(C.malloc(C.size_t(unsafe.Sizeof(C.aom_image_t{})))); e.raw == nil {
return nil, errors.New("error allocating memory")
}
if C.aom_img_alloc(e.raw, imageFormat, C.uint(width), C.uint(height), 1) == nil {
return nil, errors.New("failed to allocate image")
}
runtime.SetFinalizer(e, func(encoder *Encoder) {
encoder.Close()
})
e.cfg.g_w = C.uint(width)
e.cfg.g_h = C.uint(height)
e.cfg.g_timebase.num = C.int(stream.FrameRate().Denominator)
e.cfg.g_timebase.den = C.int(stream.FrameRate().Numerator)
e.cfg.rc_target_bitrate = bitrate
if aomErr = C.aom_codec_enc_init_ver(&e.codec, encoder, &e.cfg, 0, C.AOM_ENCODER_ABI_VERSION); aomErr != 0 {
return nil, fmt.Errorf("failed to initialize encoder: %s", C.GoString(e.codec.err_detail))
}
if aomErr = C.aom_codec_control_int(&e.codec, C.AOME_SET_CPUUSED, speed); aomErr != 0 {
return nil, errors.New("failed to set cpu-used")
}
if err := binary.Write(e.w, binary.LittleEndian, struct {
Magic [4]byte
Version uint16
HeaderSize uint16
FourCC uint32
Width uint16
Height uint16
Denominator uint32
Numerator uint32
Length uint32
Unused uint32
}{
Magic: [4]byte{'D', 'K', 'I', 'F'},
Version: 0,
HeaderSize: 32,
FourCC: 0x31305641,
Width: uint16(e.cfg.g_w),
Height: uint16(e.cfg.g_h),
Denominator: uint32(e.cfg.g_timebase.den),
Numerator: uint32(e.cfg.g_timebase.num),
Length: 0,
}); err != nil {
return nil, err
}
return e, nil
}
func (e *Encoder) Encode(pts int64, f frame.Frame) error {
//TODO: make this a Source channel
if int8F, ok := f.(frame.TypedFrame[uint8]); ok {
e.raw.planes[0] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeLuma()[0]))
e.raw.planes[1] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeCb()[0]))
e.raw.planes[2] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeCr()[0]))
} else if int16F, ok := f.(frame.TypedFrame[uint16]); ok {
e.raw.planes[0] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeLuma()[0]))
e.raw.planes[1] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeCb()[0]))
e.raw.planes[2] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeCr()[0]))
}
//cleanup pointers
defer func() {
e.raw.planes[0] = nil
e.raw.planes[1] = nil
e.raw.planes[2] = nil
}()
defer runtime.KeepAlive(f)
if _, err := e.encodeFrame(pts, e.raw); err != nil {
return err
}
return nil
}
func (e *Encoder) encodeFrame(pts int64, raw *C.aom_image_t) (pkts int, err error) {
//TODO: make this a Source channel
var aomErr C.aom_codec_err_t
if aomErr = C.aom_codec_encode(&e.codec, raw, C.long(pts), 1, 0); aomErr != C.AOM_CODEC_OK {
if aomErr == C.AOM_CODEC_INCAPABLE {
return 0, errors.New("error encoding frame: AOM_CODEC_INCAPABLE")
} else if aomErr == C.AOM_CODEC_INVALID_PARAM {
return 0, errors.New("error encoding frame: AOM_CODEC_INVALID_PARAM")
} else if aomErr == C.AOM_CODEC_ERROR {
return 0, errors.New("error encoding frame: AOM_CODEC_ERROR")
} else {
return 0, errors.New("error encoding frame")
}
}
var iter C.aom_codec_iter_t
for {
pkt := C.aom_codec_get_cx_data(&e.codec, &iter)
if pkt == nil {
break
}
pkts++
if pkt.kind == C.AOM_CODEC_CX_FRAME_PKT {
dataFrameSz := C.aom_get_pkt_sz(pkt)
dataFramePts := C.aom_get_pkt_pts(pkt)
buf := make([]byte, 4)
binary.LittleEndian.PutUint32(buf, uint32(dataFrameSz))
if _, err = e.w.Write(buf); err != nil {
return pkts, err
}
binary.LittleEndian.PutUint32(buf, uint32(dataFramePts&0xFFFFFFFF))
if _, err = e.w.Write(buf); err != nil {
return pkts, err
}
binary.LittleEndian.PutUint32(buf, uint32(dataFramePts>>32))
if _, err = e.w.Write(buf); err != nil {
return pkts, err
}
if _, err = e.w.Write(unsafe.Slice((*byte)(C.aom_get_pkt_buf(pkt)), int(dataFrameSz))); err != nil {
return pkts, err
}
}
}
return pkts, nil
}
func (e *Encoder) Flush() error {
var pkts int
var err error
for {
if pkts, err = e.encodeFrame(-1, nil); err != nil {
return err
}
if pkts == 0 {
break
}
}
return nil
}
func (e *Encoder) Close() {
if e.cleaned.Swap(true) == false {
if e.raw != nil {
C.aom_img_free(e.raw)
C.free(unsafe.Pointer(e.raw))
e.raw = nil
}
C.aom_codec_destroy(&e.codec)
}
}
func (e *Encoder) Version() string {
return Version()
}
func getSettingString(m map[string]any, name string, fallback string) string {
if v, ok := m[name]; ok {
if val, ok := v.(string); ok {
return val
}
if val, ok := v.(int); ok {
return strconv.Itoa(val)
}
if val, ok := v.(int64); ok {
return strconv.Itoa(int(val))
}
}
return fallback
}

15
encoder/libaom/libaom.h Normal file
View file

@ -0,0 +1,15 @@
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <aom/aom_encoder.h>
#include <aom/aomcx.h>
aom_codec_err_t aom_codec_control_int(aom_codec_ctx_t *ctx, int ctrl_id, int v);
aom_codec_err_t aom_codec_control_uint(aom_codec_ctx_t *ctx, int ctrl_id, unsigned int v);
void* aom_get_pkt_buf(aom_codec_cx_pkt_t *pkt);
size_t aom_get_pkt_sz(aom_codec_cx_pkt_t *pkt);
aom_codec_pts_t aom_get_pkt_pts(aom_codec_cx_pkt_t *pkt);
aom_codec_frame_flags_t aom_get_pkt_flags(aom_codec_cx_pkt_t *pkt);