Implement bitrate controller in vpx and h264 codecs (#467)

Add bitrate control to vpx and h264 encoders

Co-authored-by: Jingyang Kang <3drxkjy@gmail.com>

pkg/codec/vpx/vpx.go

@@ -74,6 +74,7 @@ type encoder struct {
 	frame           []byte
 	deadline        int
 	requireKeyFrame bool
+	targetBitrate   int
 	isKeyFrame      bool
 
 	mu     sync.Mutex
@@ -200,14 +201,15 @@ func newEncoder(r video.Reader, p prop.Media, params Params, codecIface *C.vpx_c
 	}
 	t0 := time.Now()
 	return &encoder{
-		r:          video.ToI420(r),
-		codec:      codec,
-		raw:        rawNoBuffer,
-		cfg:        cfg,
-		tStart:     t0,
-		tLastFrame: t0,
-		deadline:   int(params.Deadline / time.Microsecond),
-		frame:      make([]byte, 1024),
+		r:             video.ToI420(r),
+		codec:         codec,
+		raw:           rawNoBuffer,
+		cfg:           cfg,
+		tStart:        t0,
+		tLastFrame:    t0,
+		deadline:      int(params.Deadline / time.Microsecond),
+		frame:         make([]byte, 1024),
+		targetBitrate: params.BitRate,
 	}, nil
 }
 
@@ -261,6 +263,16 @@ func (e *encoder) Read() ([]byte, func(), error) {
 	if duration == 0 {
 		duration = 1
 	}
+
+	targetVpxBitrate := C.uint(float32(e.targetBitrate / 1000)) // convert to kilobits / second
+	if e.cfg.rc_target_bitrate != targetVpxBitrate && targetVpxBitrate >= 1 {
+		e.cfg.rc_target_bitrate = targetVpxBitrate
+		rc := C.vpx_codec_enc_config_set(e.codec, e.cfg)
+		if rc != C.VPX_CODEC_OK {
+			return nil, func() {}, fmt.Errorf("vpx_codec_enc_config_set failed (%d)", rc)
+		}
+	}
+
 	var flags int
 	if e.requireKeyFrame {
 		flags = flags | C.VPX_EFLAG_FORCE_KF
@@ -303,6 +315,13 @@ func (e *encoder) ForceKeyFrame() error {
 	return nil
 }
 
+func (e *encoder) SetBitRate(bitrate int) error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	e.targetBitrate = bitrate
+	return nil
+}
+
 func (e *encoder) Controller() codec.EncoderController {
 	return e
 }

pkg/codec/vpx/vpx_test.go

@@ -232,9 +232,76 @@ func TestRequestKeyFrame(t *testing.T) {
 	}
 }
 
-func TestShouldImplementBitRateControl(t *testing.T) {
-	t.SkipNow() // TODO: Implement bit rate control
+func TestSetBitrate(t *testing.T) {
+	for name, factory := range map[string]func() (codec.VideoEncoderBuilder, error){
+		"VP8": func() (codec.VideoEncoderBuilder, error) {
+			p, err := NewVP8Params()
+			return &p, err
+		},
+		"VP9": func() (codec.VideoEncoderBuilder, error) {
+			p, err := NewVP9Params()
+			// Disable latency to ease test and begin to receive packets for each input frame
+			p.LagInFrames = 0
+			return &p, err
+		},
+	} {
+		factory := factory
+		t.Run(name, func(t *testing.T) {
+			param, err := factory()
+			if err != nil {
+				t.Fatal(err)
+			}
 
+			var initialWidth, initialHeight, width, height int = 320, 240, 320, 240
+
+			var cnt uint32
+			r, err := param.BuildVideoEncoder(
+				video.ReaderFunc(func() (image.Image, func(), error) {
+					i := atomic.AddUint32(&cnt, 1)
+					if i == 3 {
+						return nil, nil, io.EOF
+					}
+					return image.NewYCbCr(
+						image.Rect(0, 0, width, height),
+						image.YCbCrSubsampleRatio420,
+					), func() {}, nil
+				}),
+				prop.Media{
+					Video: prop.Video{
+						Width:       initialWidth,
+						Height:      initialHeight,
+						FrameRate:   1,
+						FrameFormat: frame.FormatI420,
+					},
+				},
+			)
+			if err != nil {
+				t.Fatal(err)
+			}
+			_, rel, err := r.Read()
+			if err != nil {
+				t.Fatal(err)
+			}
+			rel()
+			err = r.Controller().(codec.BitRateController).SetBitRate(1000) // 1000 bit/second is ridiculously low, but this is a testcase.
+			if err != nil {
+				t.Fatal(err)
+			}
+			_, rel, err = r.Read()
+			if err != nil {
+				t.Fatal(err)
+			}
+			rel()
+			_, _, err = r.Read()
+			if err != io.EOF {
+				t.Fatal(err)
+			}
+		})
+
+	}
+}
+
+func TestShouldImplementBitRateControl(t *testing.T) {
 	e := &encoder{}
 	if _, ok := e.Controller().(codec.BitRateController); !ok {
 		t.Error()

pkg/codec/x264/bridge.h

@@ -8,6 +8,7 @@
 #define ERR_ALLOC_PICTURE -3
 #define ERR_OPEN_ENGINE -4
 #define ERR_ENCODE -5
+#define ERR_BITRATE_RECONFIG -6
 
 typedef struct Slice {
   unsigned char *data;
@@ -78,6 +79,22 @@ fail:
   return NULL;
 }
 
+#define RC_MARGIN 10000 /* 1kilobits / second*/
+static int apply_target_bitrate(Encoder *e, int target_bitrate) {
+  int target_encoder_bitrate = (int)target_bitrate / 1000;
+  if (e->param.rc.i_bitrate == target_encoder_bitrate || target_encoder_bitrate <= 1) {
+    return 0; // if no change to bitrate or target bitrate is too small, we return no error (0)
+  }
+
+  e->param.rc.i_bitrate = target_encoder_bitrate;
+  e->param.rc.f_rate_tolerance = 0.1;
+  e->param.rc.i_vbv_max_bitrate = target_encoder_bitrate + RC_MARGIN / 2;
+  e->param.rc.i_vbv_buffer_size = e->param.rc.i_vbv_max_bitrate;
+  e->param.rc.f_vbv_buffer_init = 0.6;
+  int success = x264_encoder_reconfig(e->h, &e->param);
+  return success; // 0 on success or negative on error
+}
+
 Slice enc_encode(Encoder *e, uint8_t *y, uint8_t *cb, uint8_t *cr, int *rc) {
   x264_nal_t *nal;
   int i_nal;

pkg/codec/x264/x264.go

@@ -39,6 +39,8 @@ func (e cerror) Error() string {
 		return errOpenEngine.Error()
 	case C.ERR_ENCODE:
 		return errEncode.Error()
+	case C.ERR_BITRATE_RECONFIG:
+		return errSetBitrate.Error()
 	default:
 		return "unknown error"
 	}
@@ -58,6 +60,7 @@ var (
 	errAllocPicture  = fmt.Errorf("failed to alloc picture")
 	errOpenEngine    = fmt.Errorf("failed to open x264")
 	errEncode        = fmt.Errorf("failed to encode")
+	errSetBitrate    = fmt.Errorf("failed to change x264 encoder bitrate")
 )
 
 func newEncoder(r video.Reader, p prop.Media, params Params) (codec.ReadCloser, error) {
@@ -133,6 +136,16 @@ func (e *encoder) ForceKeyFrame() error {
 	return nil
 }
 
+func (e *encoder) SetBitRate(bitrate int) error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	errNum := C.apply_target_bitrate(e.engine, C.int(bitrate))
+	if err := errFromC(errNum); err != nil {
+		return err
+	}
+	return nil
+}
+
 func (e *encoder) Controller() codec.EncoderController {
 	return e
 }

pkg/codec/x264/x264_test.go

@@ -7,9 +7,34 @@ import (
 	"github.com/pion/mediadevices/pkg/codec"
 	"github.com/pion/mediadevices/pkg/codec/internal/codectest"
 	"github.com/pion/mediadevices/pkg/frame"
+	"github.com/pion/mediadevices/pkg/io/video"
 	"github.com/pion/mediadevices/pkg/prop"
 )
 
+func getTestVideoEncoder() (codec.ReadCloser, error) {
+	p, err := NewParams()
+	if err != nil {
+		return nil, err
+	}
+	p.BitRate = 200000
+	enc, err := p.BuildVideoEncoder(video.ReaderFunc(func() (image.Image, func(), error) {
+		return image.NewYCbCr(
+			image.Rect(0, 0, 256, 144),
+			image.YCbCrSubsampleRatio420,
+		), nil, nil
+	}), prop.Media{
+		Video: prop.Video{
+			Width:       256,
+			Height:      144,
+			FrameFormat: frame.FormatI420,
+		},
+	})
+	if err != nil {
+		return nil, err
+	}
+	return enc, nil
+}
+
 func TestEncoder(t *testing.T) {
 	t.Run("SimpleRead", func(t *testing.T) {
 		p, err := NewParams()
@@ -69,19 +94,53 @@ func TestEncoder(t *testing.T) {
 }
 
 func TestShouldImplementKeyFrameControl(t *testing.T) {
-	t.SkipNow() // TODO: Implement key frame control
-
 	e := &encoder{}
 	if _, ok := e.Controller().(codec.KeyFrameController); !ok {
 		t.Error()
 	}
 }
 
-func TestShouldImplementBitRateControl(t *testing.T) {
-	t.SkipNow() // TODO: Implement bit rate control
+func TestNoErrorOnForceKeyFrame(t *testing.T) {
+	enc, err := getTestVideoEncoder()
+	if err != nil {
+		t.Error(err)
+	}
+	kfc, ok := enc.Controller().(codec.KeyFrameController)
+	if !ok {
+		t.Error()
+	}
+	if err := kfc.ForceKeyFrame(); err != nil {
+		t.Error(err)
+	}
+	_, rel, err := enc.Read() // try to read the encoded frame
+	rel()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
 
+func TestShouldImplementBitRateControl(t *testing.T) {
 	e := &encoder{}
 	if _, ok := e.Controller().(codec.BitRateController); !ok {
 		t.Error()
 	}
 }
+
+func TestNoErrorOnSetBitRate(t *testing.T) {
+	enc, err := getTestVideoEncoder()
+	if err != nil {
+		t.Error(err)
+	}
+	brc, ok := enc.Controller().(codec.BitRateController)
+	if !ok {
+		t.Error()
+	}
+	if err := brc.SetBitRate(1000); err != nil { // 1000 bit/second is ridiculously low, but this is a testcase.
+		t.Error(err)
+	}
+	_, rel, err := enc.Read() // try to read the encoded frame
+	rel()
+	if err != nil {
+		t.Fatal(err)
+	}
+}