L0rem1psum feat hardware frame ctx (#92)

* feat: add hardware frame context

* fix: pr comments

* feat(example): add hardware encoding example

* Clean up

* Clean up

---------

Co-authored-by: Wu Wenxuan <sunshine.xuan5960@gmail.com>

README.md

@@ -29,6 +29,7 @@ Examples are located in the [examples](examples) directory and mirror as much as
 |Filtering|[see](examples/filtering/main.go)|[see](https://github.com/FFmpeg/FFmpeg/blob/n7.0/doc/examples/filtering_video.c)
 |Frame data manipulation|[see](examples/frame_data_manipulation/main.go)|X
 |Hardware Decoding|[see](examples/hardware_decoding/main.go)|[see](https://github.com/FFmpeg/FFmpeg/blob/n7.0/doc/examples/hw_decode.c)
+|Hardware Encoding|[see](examples/hardware_encoding/main.go)|[see](https://github.com/FFmpeg/FFmpeg/blob/n7.0/doc/examples/vaapi_encode.c)
 |Remuxing|[see](examples/remuxing/main.go)|[see](https://github.com/FFmpeg/FFmpeg/blob/n7.0/doc/examples/remuxing.c)
 |Scaling|[see](examples/scaling/main.go)|[see](https://github.com/FFmpeg/FFmpeg/blob/n7.0/doc/examples/scaling_video.c)
 |Transcoding|[see](examples/transcoding/main.go)|[see](https://github.com/FFmpeg/FFmpeg/blob/n7.0/doc/examples/transcoding.c)

codec_context.go

@@ -12,6 +12,7 @@ type CodecContext struct {
 	c *C.AVCodecContext
 	// We need to store this to unref it properly
 	hdc *HardwareDeviceContext
+	hfc *HardwareFrameContext
 }
 
 func newCodecContextFromC(c *C.AVCodecContext) *CodecContext {
@@ -38,6 +39,10 @@ func (cc *CodecContext) Free() {
 		C.av_buffer_unref(&cc.hdc.c)
 		cc.hdc = nil
 	}
+	if cc.hfc != nil {
+		C.av_buffer_unref(&cc.hfc.c)
+		cc.hfc = nil
+	}
 	if cc.c != nil {
 		// Make sure to clone the classer before freeing the object since
 		// the C free method may reset the pointer
@@ -317,6 +322,16 @@ func (cc *CodecContext) SetHardwareDeviceContext(hdc *HardwareDeviceContext) {
 	}
 }
 
+func (cc *CodecContext) SetHardwareFrameContext(hfc *HardwareFrameContext) {
+	if cc.hfc != nil {
+		C.av_buffer_unref(&cc.hfc.c)
+	}
+	cc.hfc = hfc
+	if cc.hfc != nil {
+		cc.c.hw_frames_ctx = C.av_buffer_ref(cc.hfc.c)
+	}
+}
+
 func (cc *CodecContext) ExtraHardwareFrames() int {
 	return int(cc.c.extra_hw_frames)
 }
@@ -367,5 +382,4 @@ func goAstiavCodecContextGetFormat(cc *C.AVCodecContext, pfsCPtr *C.enum_AVPixel
 
 	// Callback
 	return C.enum_AVPixelFormat(c(pfs))
-
 }

examples/hardware_encoding/main.go

@@ -0,0 +1,170 @@
+package main
+
+import (
+	"errors"
+	"flag"
+	"fmt"
+	"log"
+	"strings"
+
+	"github.com/asticode/go-astiav"
+)
+
+var (
+	encoderCodecName        = flag.String("c", "", "the encoder codec name (e.g. h264_nvenc)")
+	hardwareDeviceName      = flag.String("n", "", "the hardware device name (e.g. 0)")
+	hardwareDeviceTypeName  = flag.String("t", "", "the hardware device type (e.g. cuda)")
+	hardwarePixelFormatName = flag.String("hpf", "", "the hardware pixel format name (e.g. cuda)")
+	height                  = flag.Int("h", 1080, "the height")
+	softwarePixelFormatName = flag.String("spf", "", "the software pixel format name (e.g. nv12)")
+	width                   = flag.Int("w", 1920, "the width")
+)
+
+func main() {
+	// Handle ffmpeg logs
+	astiav.SetLogLevel(astiav.LogLevelDebug)
+	astiav.SetLogCallback(func(c astiav.Classer, l astiav.LogLevel, fmt, msg string) {
+		var cs string
+		if c != nil {
+			if cl := c.Class(); cl != nil {
+				cs = " - class: " + cl.String()
+			}
+		}
+		log.Printf("ffmpeg log: %s%s - level: %d\n", strings.TrimSpace(msg), cs, l)
+	})
+
+	// Parse flags
+	flag.Parse()
+
+	// Usage
+	if *hardwareDeviceTypeName == "" || *encoderCodecName == "" || *hardwarePixelFormatName == "" {
+		log.Println("Usage: <binary path> -t <hardware device type> -c <encoder codec> -hpf <hardware pixel format> [-n <hardware device name> -w <width> -h <height>]")
+		return
+	}
+
+	// Get hardware device type
+	hardwareDeviceType := astiav.FindHardwareDeviceTypeByName(*hardwareDeviceTypeName)
+	if hardwareDeviceType == astiav.HardwareDeviceTypeNone {
+		log.Fatal(errors.New("main: hardware device not found"))
+	}
+
+	// Create hardware device context
+	hardwareDeviceContext, err := astiav.CreateHardwareDeviceContext(hardwareDeviceType, *hardwareDeviceName, nil)
+	if err != nil {
+		log.Fatal(fmt.Errorf("main: creating hardware device context failed: %w", err))
+	}
+
+	// Find encoder codec
+	encCodec := astiav.FindEncoderByName(*encoderCodecName)
+	if encCodec == nil {
+		log.Fatal("main: encoder codec is nil")
+	}
+
+	// Alloc codec context
+	encCodecContext := astiav.AllocCodecContext(encCodec)
+	if encCodecContext == nil {
+		log.Fatal("main: codec context is nil")
+	}
+	defer encCodecContext.Free()
+
+	// Get hardware pixel format
+	hardwarePixelFormat := astiav.FindPixelFormatByName(*hardwarePixelFormatName)
+	if hardwarePixelFormat == astiav.PixelFormatNone {
+		log.Fatal("main: hardware pixel format not found")
+	}
+
+	// Set codec context
+	encCodecContext.SetWidth(*width)
+	encCodecContext.SetHeight(*height)
+	encCodecContext.SetTimeBase(astiav.NewRational(1, 25))
+	encCodecContext.SetFramerate(encCodecContext.TimeBase().Invert())
+	encCodecContext.SetPixelFormat(hardwarePixelFormat)
+
+	// Alloc hardware frame context
+	hardwareFrameContext := astiav.AllocHardwareFrameContext(hardwareDeviceContext)
+	if hardwareFrameContext == nil {
+		log.Fatal("main: hardware frame context is nil")
+	}
+
+	// Get software pixel format
+	softwarePixelFormat := astiav.FindPixelFormatByName(*softwarePixelFormatName)
+	if softwarePixelFormat == astiav.PixelFormatNone {
+		log.Fatal("main: software pixel format not found")
+	}
+
+	// Set hardware frame content
+	hardwareFrameContext.SetPixelFormat(hardwarePixelFormat)
+	hardwareFrameContext.SetSoftwarePixelFormat(softwarePixelFormat)
+	hardwareFrameContext.SetWidth(*width)
+	hardwareFrameContext.SetHeight(*height)
+	hardwareFrameContext.SetInitialPoolSize(20)
+
+	// Initialize hardware frame context
+	if err := hardwareFrameContext.Initialize(); err != nil {
+		log.Fatal(fmt.Errorf("main: initializing hardware frame context failed: %w", err))
+	}
+
+	// Update encoder codec context hardware frame context
+	encCodecContext.SetHardwareFrameContext(hardwareFrameContext)
+
+	// Open codec context
+	if err := encCodecContext.Open(encCodec, nil); err != nil {
+		log.Fatal(fmt.Errorf("main: opening codec context failed: %w", err))
+	}
+
+	// Alloc software frame
+	softwareFrame := astiav.AllocFrame()
+	defer softwareFrame.Free()
+
+	// Set software frame
+	softwareFrame.SetWidth(*width)
+	softwareFrame.SetHeight(*height)
+	softwareFrame.SetPixelFormat(softwarePixelFormat)
+
+	// Alloc software frame buffer
+	if err := softwareFrame.AllocBuffer(0); err != nil {
+		log.Fatal(fmt.Errorf("main: allocating buffer failed: %w", err))
+	}
+
+	// Fill software frame with black
+	if err = softwareFrame.ImageFillBlack(); err != nil {
+		log.Fatal(fmt.Errorf("main: filling software frame with black failed: %w", err))
+	}
+
+	// Alloc hardware frame
+	hardwareFrame := astiav.AllocFrame()
+	defer hardwareFrame.Free()
+
+	// Alloc hardware frame buffer
+	if err := hardwareFrame.AllocHardwareBuffer(hardwareFrameContext); err != nil {
+		log.Fatal(fmt.Errorf("main: allocating hardware buffer failed: %w", err))
+	}
+
+	// Transfer software frame data to hardware frame
+	if err := softwareFrame.TransferHardwareData(hardwareFrame); err != nil {
+		log.Fatal(fmt.Errorf("main: transferring hardware data failed: %w", err))
+	}
+
+	// Encode frame
+	if err := encCodecContext.SendFrame(hardwareFrame); err != nil {
+		log.Fatal(fmt.Errorf("main: sending frame failed: %w", err))
+	}
+
+	// Alloc packet
+	pkt := astiav.AllocPacket()
+	defer pkt.Free()
+
+	// Loop
+	for {
+		// Receive packet
+		if err = encCodecContext.ReceivePacket(pkt); err != nil {
+			if errors.Is(err, astiav.ErrEof) || errors.Is(err, astiav.ErrEagain) {
+				break
+			}
+			log.Fatal(fmt.Errorf("main: receiving packet failed: %w", err))
+		}
+
+		// Log
+		log.Println("new packet")
+	}
+}

frame.go

@@ -33,6 +33,10 @@ func (f *Frame) AllocBuffer(align int) error {
 	return newError(C.av_frame_get_buffer(f.c, C.int(align)))
 }
 
+func (f *Frame) AllocHardwareBuffer(hfc *HardwareFrameContext) error {
+	return newError(C.av_hwframe_get_buffer(hfc.c, f.c, 0))
+}
+
 func (f *Frame) AllocImage(align int) error {
 	return newError(C.av_image_alloc(&f.c.data[0], &f.c.linesize[0], f.c.width, f.c.height, (C.enum_AVPixelFormat)(f.c.format), C.int(align)))
 }

hardware_frame_context.go

@@ -0,0 +1,51 @@
+package astiav
+
+//#include <libavcodec/avcodec.h>
+import "C"
+import (
+	"unsafe"
+)
+
+// https://github.com/FFmpeg/FFmpeg/blob/n7.0/libavutil/hwcontext.h#L115
+type HardwareFrameContext struct {
+	c *C.struct_AVBufferRef
+}
+
+func newHardwareFrameContextFromC(c *C.struct_AVBufferRef) *HardwareFrameContext {
+	if c == nil {
+		return nil
+	}
+	return &HardwareFrameContext{c: c}
+}
+
+func AllocHardwareFrameContext(hdc *HardwareDeviceContext) *HardwareFrameContext {
+	return newHardwareFrameContextFromC(C.av_hwframe_ctx_alloc(hdc.c))
+}
+
+func (hfc *HardwareFrameContext) data() *C.AVHWFramesContext {
+	return (*C.AVHWFramesContext)(unsafe.Pointer((hfc.c.data)))
+}
+
+func (hfc *HardwareFrameContext) SetWidth(width int) {
+	hfc.data().width = C.int(width)
+}
+
+func (hfc *HardwareFrameContext) SetHeight(height int) {
+	hfc.data().height = C.int(height)
+}
+
+func (hfc *HardwareFrameContext) SetPixelFormat(format PixelFormat) {
+	hfc.data().format = C.enum_AVPixelFormat(format)
+}
+
+func (hfc *HardwareFrameContext) SetSoftwarePixelFormat(swFormat PixelFormat) {
+	hfc.data().sw_format = C.enum_AVPixelFormat(swFormat)
+}
+
+func (hfc *HardwareFrameContext) SetInitialPoolSize(initialPoolSize int) {
+	hfc.data().initial_pool_size = C.int(initialPoolSize)
+}
+
+func (hfc *HardwareFrameContext) Initialize() error {
+	return newError(C.av_hwframe_ctx_init(hfc.c))
+}