Expose ratecontrol options, issue #26

gen2brain · Apr 28, 2023 · 5e0d7bf · 5e0d7bf
1 parent fe508eb
commit 5e0d7bf
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 6 deletions.
diff --git a/encode.go b/encode.go
@@ -37,6 +37,12 @@ type Options struct {
 	Preset string
 	// Profiles: baseline, main, high, high10, high422, high444.
 	Profile string
+	// RateControl: cqp, crf, abr.
+	RateControl string
+	// RateConstant.
+	RateConstant float32
+	// RateMax.
+	RateMax float32
 	// Log level.
 	LogLevel int32
 }
@@ -75,7 +81,7 @@ func NewEncoder(w io.Writer, opts *Options) (e *Encoder, err error) {
 
 	param := x264c.Param{}
 
-	if e.opts.Preset != "" && e.opts.Profile != "" {
+	if e.opts.Preset != "" && e.opts.Tune != "" {
 		ret := x264c.ParamDefaultPreset(&param, e.opts.Preset, e.opts.Tune)
 		if ret < 0 {
 			err = fmt.Errorf("x264: invalid preset/tune name")
@@ -108,6 +114,32 @@ func NewEncoder(w io.Writer, opts *Options) (e *Encoder, err error) {
 		}
 	}
 
+	if e.opts.RateControl != "" {
+		switch e.opts.RateControl {
+		case "cqp":
+			param.Rc.IRcMethod = x264c.RcCqp
+			if e.opts.RateConstant != 0 {
+				param.Rc.IQpConstant = int32(e.opts.RateConstant)
+			}
+			if e.opts.RateMax != 0 {
+				param.Rc.IQpMax = int32(e.opts.RateMax)
+			}
+		case "crf":
+			param.Rc.IRcMethod = x264c.RcCrf
+			if e.opts.RateConstant != 0 {
+				param.Rc.FRfConstant = e.opts.RateConstant
+			}
+			if e.opts.RateMax != 0 {
+				param.Rc.FRfConstantMax = e.opts.RateMax
+			}
+		case "abr":
+			param.Rc.IRcMethod = x264c.RcAbr
+			if e.opts.RateMax != 0 {
+				param.Rc.IVbvMaxBitrate = int32(e.opts.RateMax)
+			}
+		}
+	}
+
 	var picIn x264c.Picture
 	x264c.PictureInit(&picIn)
 	e.picIn = picIn

diff --git a/encode_test.go b/encode_test.go
@@ -5,7 +5,6 @@ import (
 	"image"
 	"image/color"
 	"image/draw"
-	"io/ioutil"
 	"os"
 	"path/filepath"
 	"testing"
@@ -33,7 +32,7 @@ func TestEncode(t *testing.T) {
 	draw.Draw(img, img.Bounds(), image.Black, image.ZP, draw.Src)
 
 	for i := 0; i < opts.Width/2; i++ {
-		img.Set(i, opts.Height/2, color.RGBA{255, 0, 0, 255})
+		img.Set(i, opts.Height/2, color.RGBA{R: 255, A: 255})
 
 		err = enc.Encode(img)
 		if err != nil {
@@ -51,7 +50,7 @@ func TestEncode(t *testing.T) {
 		t.Error(err)
 	}
 
-	err = ioutil.WriteFile(filepath.Join(os.TempDir(), "test.264"), buf.Bytes(), 0644)
+	err = os.WriteFile(filepath.Join(os.TempDir(), "test.264"), buf.Bytes(), 0644)
 	if err != nil {
 		t.Error(err)
 	}
@@ -79,7 +78,7 @@ func TestEncodeFlush(t *testing.T) {
 	draw.Draw(img, img.Bounds(), image.Black, image.ZP, draw.Src)
 
 	for i := 0; i < opts.Width/2; i++ {
-		img.Set(i, opts.Height/2, color.RGBA{255, 0, 0, 255})
+		img.Set(i, opts.Height/2, color.RGBA{R: 255, A: 255})
 
 		err = enc.Encode(img)
 		if err != nil {
@@ -97,7 +96,55 @@ func TestEncodeFlush(t *testing.T) {
 		t.Error(err)
 	}
 
-	err = ioutil.WriteFile(filepath.Join(os.TempDir(), "test.high.264"), buf.Bytes(), 0644)
+	err = os.WriteFile(filepath.Join(os.TempDir(), "test.high.264"), buf.Bytes(), 0644)
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func TestEncodeCrf(t *testing.T) {
+	buf := bytes.NewBuffer(make([]byte, 0))
+
+	opts := &Options{
+		Width:        640,
+		Height:       480,
+		FrameRate:    25,
+		Tune:         "zerolatency",
+		Preset:       "veryfast",
+		Profile:      "baseline",
+		RateControl:  "crf",
+		RateConstant: 18,
+		LogLevel:     LogDebug,
+	}
+
+	enc, err := NewEncoder(buf, opts)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	img := NewYCbCr(image.Rect(0, 0, opts.Width, opts.Height))
+	draw.Draw(img, img.Bounds(), image.Black, image.ZP, draw.Src)
+
+	for i := 0; i < opts.Width/2; i++ {
+		img.Set(i, opts.Height/2, color.RGBA{R: 255, A: 255})
+
+		err = enc.Encode(img)
+		if err != nil {
+			t.Error(err)
+		}
+	}
+
+	err = enc.Flush()
+	if err != nil {
+		t.Error(err)
+	}
+
+	err = enc.Close()
+	if err != nil {
+		t.Error(err)
+	}
+
+	err = os.WriteFile(filepath.Join(os.TempDir(), "test.crf.264"), buf.Bytes(), 0644)
 	if err != nil {
 		t.Error(err)
 	}