245 lines
6.4 KiB
Go
245 lines
6.4 KiB
Go
package x264
|
|
|
|
/*
|
|
#cgo pkg-config: x264
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <x264.h>
|
|
*/
|
|
import "C"
|
|
import (
|
|
"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
|
|
params C.x264_param_t
|
|
pictureIn *C.x264_picture_t
|
|
pictureOut *C.x264_picture_t
|
|
h *C.x264_t
|
|
}
|
|
|
|
func NewEncoder(w io.Writer, stream *y4m.Stream, settings map[string]any) (*Encoder, error) {
|
|
e := &Encoder{
|
|
w: w,
|
|
}
|
|
|
|
preset := C.CString(getSettingString(settings, "preset", "medium"))
|
|
defer C.free(unsafe.Pointer(preset))
|
|
|
|
if C.x264_param_default_preset(&e.params, preset, nil) < 0 {
|
|
return nil, errors.New("error setting preset")
|
|
}
|
|
|
|
defaultProfile := "high"
|
|
|
|
switch true {
|
|
case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 4 && stream.ColorSpace().Subsampling.B == 4:
|
|
e.params.i_csp = C.X264_CSP_I444
|
|
defaultProfile = "high444"
|
|
case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 2 && stream.ColorSpace().Subsampling.B == 2:
|
|
e.params.i_csp = C.X264_CSP_I422
|
|
defaultProfile = "high422"
|
|
case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 2 && stream.ColorSpace().Subsampling.B == 0:
|
|
e.params.i_csp = C.X264_CSP_I420
|
|
defaultProfile = "high"
|
|
default:
|
|
return nil, errors.New("unsupported input chroma subsampling")
|
|
|
|
}
|
|
|
|
profile := C.CString(getSettingString(settings, "profile", defaultProfile))
|
|
defer C.free(unsafe.Pointer(profile))
|
|
|
|
if stream.ColorSpace().BitDepth > 8 {
|
|
e.params.i_csp |= C.X264_CSP_HIGH_DEPTH
|
|
if defaultProfile == "high" {
|
|
defaultProfile = "high10"
|
|
}
|
|
}
|
|
e.params.i_bitdepth = C.int(stream.ColorSpace().BitDepth)
|
|
|
|
width, height := stream.Resolution()
|
|
e.params.i_width = C.int(width)
|
|
e.params.i_height = C.int(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(stream.FrameRate().Numerator)
|
|
e.params.i_fps_den = C.uint32_t(stream.FrameRate().Denominator)
|
|
|
|
if stream.ColorRange() {
|
|
e.params.vui.b_fullrange = 1
|
|
} else {
|
|
e.params.vui.b_fullrange = 0
|
|
}
|
|
|
|
if C.x264_param_apply_profile(&e.params, profile) < 0 {
|
|
return nil, errors.New("error setting profile")
|
|
}
|
|
|
|
for k, v := range settings {
|
|
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.Itoa(val))
|
|
} else if val, ok := v.(int64); ok {
|
|
strVal = C.CString(strconv.Itoa(int(val)))
|
|
}
|
|
|
|
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
|
|
}
|
|
}
|
|
|
|
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 e.pictureOut = (*C.x264_picture_t)(C.malloc(C.size_t(unsafe.Sizeof(C.x264_picture_t{})))); e.pictureOut == 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) Encode(pts int64, f frame.Frame) error {
|
|
//TODO: make this a Source channel
|
|
var nal *C.x264_nal_t
|
|
var i_nal C.int
|
|
var frame_size C.int
|
|
|
|
if int8F, ok := f.(frame.TypedFrame[uint8]); ok {
|
|
e.pictureIn.img.plane[0] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeLuma()[0]))
|
|
e.pictureIn.img.plane[1] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeCb()[0]))
|
|
e.pictureIn.img.plane[2] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeCr()[0]))
|
|
} else if int16F, ok := f.(frame.TypedFrame[uint16]); ok {
|
|
e.pictureIn.img.plane[0] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeLuma()[0]))
|
|
e.pictureIn.img.plane[1] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeCb()[0]))
|
|
e.pictureIn.img.plane[2] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeCr()[0]))
|
|
}
|
|
|
|
//cleanup pointers
|
|
defer func() {
|
|
e.pictureIn.img.plane[0] = nil
|
|
e.pictureIn.img.plane[1] = nil
|
|
e.pictureIn.img.plane[2] = nil
|
|
|
|
}()
|
|
defer runtime.KeepAlive(f)
|
|
|
|
e.pictureIn.i_pts = C.int64_t(pts)
|
|
|
|
if frame_size = C.x264_encoder_encode(e.h, &nal, &i_nal, e.pictureIn, e.pictureOut); frame_size < 0 {
|
|
return errors.New("error encoding frame")
|
|
}
|
|
|
|
if frame_size == 0 {
|
|
return nil
|
|
}
|
|
|
|
buf := unsafe.Slice((*byte)(nal.p_payload), int(frame_size))
|
|
if _, err := e.w.Write(buf); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (e *Encoder) Flush() error {
|
|
var nal *C.x264_nal_t
|
|
var i_nal C.int
|
|
var frame_size C.int
|
|
|
|
for C.x264_encoder_delayed_frames(e.h) > 0 {
|
|
if frame_size = C.x264_encoder_encode(e.h, &nal, &i_nal, nil, e.pictureOut); frame_size <= 0 {
|
|
return errors.New("error encoding frame")
|
|
}
|
|
|
|
buf := unsafe.Slice((*byte)(nal.p_payload), int(frame_size))
|
|
if _, err := e.w.Write(buf); err != nil {
|
|
return err
|
|
}
|
|
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
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
|
|
}
|
|
if e.pictureOut != nil {
|
|
C.free(unsafe.Pointer(e.pictureOut))
|
|
e.pictureOut = 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
|
|
}
|
|
|
|
func getSettingInt(m map[string]any, name string, fallback int) int {
|
|
return fallback
|
|
}
|