Skip to content

Commit

Permalink
L0rem1psum feat hardware frame ctx (#92)
Browse files Browse the repository at this point in the history
* feat: add hardware frame context

* fix: pr comments

* feat(example): add hardware encoding example

* Clean up

* Clean up

---------

Co-authored-by: Wu Wenxuan <[email protected]>
  • Loading branch information
asticode and l0rem1psum authored Oct 23, 2024
1 parent cd2a16d commit bc14852
Show file tree
Hide file tree
Showing 5 changed files with 241 additions and 1 deletion.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
16 changes: 15 additions & 1 deletion codec_context.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -367,5 +382,4 @@ func goAstiavCodecContextGetFormat(cc *C.AVCodecContext, pfsCPtr *C.enum_AVPixel

// Callback
return C.enum_AVPixelFormat(c(pfs))

}
170 changes: 170 additions & 0 deletions examples/hardware_encoding/main.go
Original file line number Diff line number Diff line change
@@ -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")
}
}
4 changes: 4 additions & 0 deletions frame.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)))
}
Expand Down
51 changes: 51 additions & 0 deletions hardware_frame_context.go
Original file line number Diff line number Diff line change
@@ -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))
}

0 comments on commit bc14852

Please sign in to comment.