Ignite/encoder/libx264/libx264.go

280 lines
7.3 KiB
Go

//go:build cgo && !disable_codec_libx264
package libx264
/*
#cgo pkg-config: x264
#include "libx264.h"
*/
import "C"
import (
"errors"
"fmt"
"git.gammaspectra.live/S.O.N.G/Ignite/frame"
"io"
"runtime"
"strconv"
"sync/atomic"
"unsafe"
)
type Encoder struct {
w io.Writer
cleaned atomic.Bool
params C.x264_param_t
pictureIn *C.x264_picture_t
pictureOut C.x264_picture_t
h *C.x264_t
}
var x264Version = fmt.Sprintf("x264 core:%d%s", C.X264_BUILD, C.GoString(C.Version()))
func Version() string {
return x264Version
}
func NewEncoder(w io.Writer, properties frame.StreamProperties, settings map[string]any) (*Encoder, error) {
e := &Encoder{
w: w,
}
preset := C.CString(getSettingString(settings, "preset", "medium"))
defer C.free(unsafe.Pointer(preset))
var tune *C.char
if strVal := getSettingString(settings, "tune", ""); strVal != "" {
tune = C.CString(strVal)
defer C.free(unsafe.Pointer(tune))
}
if C.x264_param_default_preset(&e.params, preset, tune) < 0 {
return nil, errors.New("error setting preset or tune")
}
defaultProfile := "high"
switch true {
case properties.ColorSpace.ChromaSampling.J == 4 && properties.ColorSpace.ChromaSampling.A == 4 && properties.ColorSpace.ChromaSampling.B == 4:
e.params.i_csp = C.X264_CSP_I444
defaultProfile = "high444"
case properties.ColorSpace.ChromaSampling.J == 4 && properties.ColorSpace.ChromaSampling.A == 2 && properties.ColorSpace.ChromaSampling.B == 2:
e.params.i_csp = C.X264_CSP_I422
defaultProfile = "high422"
case properties.ColorSpace.ChromaSampling.J == 4 && properties.ColorSpace.ChromaSampling.A == 2 && properties.ColorSpace.ChromaSampling.B == 0:
e.params.i_csp = C.X264_CSP_I420
defaultProfile = "high"
case properties.ColorSpace.ChromaSampling.J == 4 && properties.ColorSpace.ChromaSampling.A == 0 && properties.ColorSpace.ChromaSampling.B == 0:
e.params.i_csp = C.X264_CSP_I400
defaultProfile = "high"
default:
return nil, errors.New("unsupported input chroma subsampling")
}
if properties.ColorSpace.BitDepth > 8 {
e.params.i_csp |= C.X264_CSP_HIGH_DEPTH
if defaultProfile == "high" {
defaultProfile = "high10"
}
}
profile := C.CString(getSettingString(settings, "profile", defaultProfile))
defer C.free(unsafe.Pointer(profile))
e.params.i_bitdepth = C.int(properties.ColorSpace.BitDepth)
e.params.i_width = C.int(properties.Width)
e.params.i_height = C.int(properties.Height)
e.params.b_vfr_input = 0
e.params.b_repeat_headers = 1
e.params.b_annexb = 1
e.params.i_fps_num = C.uint32_t(properties.FrameRate.Numerator)
e.params.i_fps_den = C.uint32_t(properties.FrameRate.Denominator)
if properties.FullColorRange {
e.params.vui.b_fullrange = 1
} else {
e.params.vui.b_fullrange = 0
}
for k, v := range settings {
if k == "profile" || k == "preset" || k == "tune" {
continue
}
if err := func() error {
var strVal *C.char
if val, ok := v.(string); ok {
strVal = C.CString(val)
} else if val, ok := v.(int); ok {
strVal = C.CString(strconv.FormatInt(int64(val), 10))
} else if val, ok := v.(int64); ok {
strVal = C.CString(strconv.FormatInt(val, 10))
} else if val, ok := v.(uint); ok {
strVal = C.CString(strconv.FormatUint(uint64(val), 10))
} else if val, ok := v.(uint64); ok {
strVal = C.CString(strconv.FormatUint(val, 10))
}
if strVal != nil {
defer C.free(unsafe.Pointer(strVal))
} else {
return fmt.Errorf("could not get parameter %s", k)
}
strKey := C.CString(k)
defer C.free(unsafe.Pointer(strKey))
if ret := C.x264_param_parse(&e.params, strKey, strVal); ret != 0 {
if ret == C.X264_PARAM_BAD_NAME {
return fmt.Errorf("bad parameter name %s", k)
} else if ret == C.X264_PARAM_BAD_VALUE {
return fmt.Errorf("bad parameter value %s for %s", C.GoString(strVal), k)
} else {
return fmt.Errorf("error setting parameter %s", k)
}
}
return nil
}(); err != nil {
return nil, err
}
}
if e.params.rc.b_stat_read == 0 && e.params.rc.b_stat_write == 1 && C.GoString(preset) != "placebo" {
//doing first pass
C.x264_param_apply_fastfirstpass(&e.params)
}
if C.x264_param_apply_profile(&e.params, profile) < 0 {
return nil, errors.New("error setting profile")
}
runtime.SetFinalizer(e, func(encoder *Encoder) {
encoder.Close()
})
if e.pictureIn = (*C.x264_picture_t)(C.malloc(C.size_t(unsafe.Sizeof(C.x264_picture_t{})))); e.pictureIn == nil {
return nil, errors.New("error allocating memory")
}
if C.x264_picture_alloc(e.pictureIn, e.params.i_csp, e.params.i_width, e.params.i_height) < 0 {
return nil, errors.New("error allocating input picture")
}
if e.h = C.x264_encoder_open(&e.params); e.h == nil {
return nil, errors.New("error opening encoder")
}
return e, nil
}
func (e *Encoder) EncodeStream(stream *frame.Stream) error {
for f := range stream.Channel() {
if err := e.Encode(f); err != nil {
return err
}
}
return e.Flush()
}
func (e *Encoder) Encode(f frame.Frame) error {
var nal *C.x264_nal_t
var iNal C.int
var frameSize C.int
if f8, ok := f.(frame.TypedFrame[uint8]); ok {
e.pictureIn.img.plane[0] = (*C.uint8_t)(unsafe.Pointer(&f8.GetNativeLuma()[0]))
e.pictureIn.img.plane[1] = (*C.uint8_t)(unsafe.Pointer(&f8.GetNativeCb()[0]))
e.pictureIn.img.plane[2] = (*C.uint8_t)(unsafe.Pointer(&f8.GetNativeCr()[0]))
} else if f16, ok := f.(frame.TypedFrame[uint16]); ok {
e.pictureIn.img.plane[0] = (*C.uint8_t)(unsafe.Pointer(&f16.GetNativeLuma()[0]))
e.pictureIn.img.plane[1] = (*C.uint8_t)(unsafe.Pointer(&f16.GetNativeCb()[0]))
e.pictureIn.img.plane[2] = (*C.uint8_t)(unsafe.Pointer(&f16.GetNativeCr()[0]))
}
defer runtime.KeepAlive(f)
//cleanup pointers
defer func() {
e.pictureIn.img.plane[0] = nil
e.pictureIn.img.plane[1] = nil
e.pictureIn.img.plane[2] = nil
}()
e.pictureIn.i_pts = C.int64_t(f.PTS())
if frameSize = C.x264_encoder_encode(e.h, &nal, &iNal, e.pictureIn, &e.pictureOut); frameSize < 0 {
return errors.New("error encoding frame")
}
if frameSize == 0 {
return nil
}
if iNal != 1 {
//return errors.New("more than one NAL present")
}
buf := unsafe.Slice((*byte)(nal.p_payload), int(frameSize))
if _, err := e.w.Write(buf); err != nil {
return err
}
return nil
}
func (e *Encoder) Flush() error {
var nal *C.x264_nal_t
var iNal C.int
var frameSize C.int
for C.x264_encoder_delayed_frames(e.h) > 0 {
if frameSize = C.x264_encoder_encode(e.h, &nal, &iNal, nil, &e.pictureOut); frameSize < 0 {
return errors.New("error encoding frame")
}
if frameSize == 0 {
return nil
}
if iNal != 1 {
//return errors.New("more than one NAL present")
}
buf := unsafe.Slice((*byte)(nal.p_payload), int(frameSize))
if _, err := e.w.Write(buf); err != nil {
return err
}
}
return nil
}
func (e *Encoder) Version() string {
return Version()
}
func (e *Encoder) Close() {
if e.cleaned.Swap(true) == false {
if e.h != nil {
C.x264_encoder_close(e.h)
e.h = nil
}
if e.pictureIn != nil {
C.x264_picture_clean(e.pictureIn)
C.free(unsafe.Pointer(e.pictureIn))
e.pictureIn = nil
}
}
}
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
}