diff --git a/codec_context.go b/codec_context.go index dc8768c..fa34ccc 100644 --- a/codec_context.go +++ b/codec_context.go @@ -3,7 +3,34 @@ package astiav //#cgo pkg-config: libavcodec libavutil //#include //#include +/* +extern enum AVPixelFormat goAstiavCodecContextGetFormat(AVCodecContext *ctx, enum AVPixelFormat *pix_fmts, int pix_fmts_size); + +static inline enum AVPixelFormat astiavCodecContextGetFormat(AVCodecContext *ctx, const enum AVPixelFormat *pix_fmts) +{ + int pix_fmts_size = 0; + while (*pix_fmts != AV_PIX_FMT_NONE) { + pix_fmts_size++; + pix_fmts++; + } + pix_fmts -= pix_fmts_size; + return goAstiavCodecContextGetFormat(ctx, (enum AVPixelFormat*)(pix_fmts), pix_fmts_size); +} +static inline void astiavSetCodecContextGetFormat(AVCodecContext *ctx) +{ + ctx->get_format = astiavCodecContextGetFormat; +} +static inline void astiavResetCodecContextGetFormat(AVCodecContext *ctx) +{ + ctx->get_format = NULL; +} + +*/ import "C" +import ( + "sync" + "unsafe" +) // https://github.com/FFmpeg/FFmpeg/blob/n5.0/libavcodec/avcodec.h#L383 type CodecContext struct { @@ -263,5 +290,50 @@ func (cc *CodecContext) SendFrame(f *Frame) error { } func (cc *CodecContext) SetHardwareDeviceContext(hdc *HardwareDeviceContext) { - cc.c.hw_device_ctx = hdc.c + cc.c.hw_device_ctx = C.av_buffer_ref(hdc.c) +} + +type CodecContextPixelFormatCallback func(pfs []PixelFormat) PixelFormat + +var ( + codecContextPixelFormatCallbacks = make(map[*C.struct_AVCodecContext]CodecContextPixelFormatCallback) + codecContextPixelFormatCallbacksMutex = &sync.Mutex{} +) + +func (cc *CodecContext) SetPixelFormatCallback(c CodecContextPixelFormatCallback) { + // Lock + codecContextPixelFormatCallbacksMutex.Lock() + defer codecContextPixelFormatCallbacksMutex.Unlock() + + // Update callback + if c == nil { + C.astiavResetCodecContextGetFormat(cc.c) + delete(codecContextPixelFormatCallbacks, cc.c) + } else { + C.astiavSetCodecContextGetFormat(cc.c) + codecContextPixelFormatCallbacks[cc.c] = c + } +} + +//export goAstiavCodecContextGetFormat +func goAstiavCodecContextGetFormat(cc *C.struct_AVCodecContext, pfsCPtr *C.enum_AVPixelFormat, pfsCSize C.int) C.enum_AVPixelFormat { + // Lock + codecContextPixelFormatCallbacksMutex.Lock() + defer codecContextPixelFormatCallbacksMutex.Unlock() + + // Get callback + c, ok := codecContextPixelFormatCallbacks[cc] + if !ok { + return C.enum_AVPixelFormat(PixelFormatNone) + } + + // Get pixel formats + var pfs []PixelFormat + for _, v := range unsafe.Slice(pfsCPtr, pfsCSize) { + pfs = append(pfs, PixelFormat(v)) + } + + // Callback + return C.enum_AVPixelFormat(c(pfs)) + } diff --git a/examples/hardware_decoding/main.go b/examples/hardware_decoding/main.go index b032117..893dd6f 100644 --- a/examples/hardware_decoding/main.go +++ b/examples/hardware_decoding/main.go @@ -11,7 +11,9 @@ import ( ) var ( - hardwareDeviceTypeName = flag.String("d", "", "the hardware device type like: cuda") + decoderCodecName = flag.String("c", "", "the decoder 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)") input = flag.String("i", "", "the input path") ) @@ -35,7 +37,7 @@ func main() { // Usage if *input == "" || *hardwareDeviceTypeName == "" { - log.Println("Usage: -d -i ") + log.Println("Usage: -t -i [-n -c ]") return } @@ -87,7 +89,14 @@ func main() { s := &stream{inputStream: is} // Find decoder - if s.decCodec = astiav.FindDecoder(is.CodecParameters().CodecID()); s.decCodec == nil { + if *decoderCodecName != "" { + s.decCodec = astiav.FindDecoderByName(*decoderCodecName) + } else { + s.decCodec = astiav.FindDecoder(is.CodecParameters().CodecID()) + } + + // No codec + if s.decCodec == nil { log.Fatal(errors.New("main: codec is nil")) } @@ -97,11 +106,8 @@ func main() { } defer s.decCodecContext.Free() - // Get codec hardware configs - hardwareConfigs := s.decCodec.HardwareConfigs(hardwareDeviceType) - // Loop through codec hardware configs - for _, p := range hardwareConfigs { + for _, p := range s.decCodec.HardwareConfigs(hardwareDeviceType) { // Valid hardware config if p.MethodFlags().Has(astiav.CodecHardwareConfigMethodFlagHwDeviceCtx) && p.HardwareDeviceType() == hardwareDeviceType { s.hardwarePixelFormat = p.PixelFormat() @@ -121,11 +127,21 @@ func main() { // Create hardware device context var err error - s.hardwareDeviceContext, err = astiav.CreateHardwareDeviceContext(hardwareDeviceType, "", nil) - if err != nil { + if s.hardwareDeviceContext, err = astiav.CreateHardwareDeviceContext(hardwareDeviceType, *hardwareDeviceName, nil); err != nil { log.Fatal(fmt.Errorf("main: creating hardware device context failed: %w", err)) } + + // Update decoder context s.decCodecContext.SetHardwareDeviceContext(s.hardwareDeviceContext) + s.decCodecContext.SetPixelFormatCallback(func(pfs []astiav.PixelFormat) astiav.PixelFormat { + for _, pf := range pfs { + if pf == s.hardwarePixelFormat { + return pf + } + } + log.Fatal(errors.New("main: using hardware pixel format failed")) + return astiav.PixelFormatNone + }) // Open codec context if err := s.decCodecContext.Open(s.decCodec, nil); err != nil { @@ -186,7 +202,7 @@ func main() { } // Do something with decoded frame - log.Printf("new frame: stream %d - pts: %d", pkt.StreamIndex(), finalFrame.Pts()) + log.Printf("new frame: stream %d - pts: %d - transferred: %v", pkt.StreamIndex(), finalFrame.Pts(), hardwareFrame.PixelFormat() == s.hardwarePixelFormat) } } }