Files
go-astiav/frame_data.go
Quentin Renard 63a083bcad Added FrameData
2023-09-02 17:26:25 +02:00

121 lines
3.2 KiB
Go

package astiav
//#include <stdint.h>
import "C"
import (
"errors"
"fmt"
"image"
"strings"
)
type FrameData struct {
f *Frame
}
func newFrameData(f *Frame) *FrameData {
return &FrameData{f: f}
}
func (d *FrameData) Bytes(align int) ([]byte, error) {
switch {
// Video
case d.f.Height() > 0 && d.f.Width() > 0:
// Get buffer size
s, err := d.f.ImageBufferSize(align)
if err != nil {
return nil, fmt.Errorf("astiav: getting image buffer size failed: %w", err)
}
// Invalid buffer size
if s == 0 {
return nil, errors.New("astiav: invalid image buffer size")
}
// Create buffer
b := make([]byte, s)
// Copy image to buffer
if _, err = d.f.ImageCopyToBuffer(b, align); err != nil {
return nil, fmt.Errorf("astiav: copying image to buffer failed: %w", err)
}
return b, nil
}
return nil, errors.New("astiav: frame type not implemented")
}
func (d *FrameData) planeData(i int, sizeFunc func(linesize int) int) []byte {
return bytesFromC(func(size *cUlong) *C.uint8_t {
*size = cUlong(sizeFunc(int(d.f.c.linesize[i])))
return d.f.c.data[i]
})
}
func (d *FrameData) imageYCbCrSubsampleRatio() image.YCbCrSubsampleRatio {
name := d.f.PixelFormat().Name()
for s, r := range map[string]image.YCbCrSubsampleRatio{
"410": image.YCbCrSubsampleRatio410,
"411": image.YCbCrSubsampleRatio411,
"420": image.YCbCrSubsampleRatio420,
"422": image.YCbCrSubsampleRatio422,
"440": image.YCbCrSubsampleRatio440,
"444": image.YCbCrSubsampleRatio444,
} {
if strings.Contains(name, s) {
return r
}
}
return image.YCbCrSubsampleRatio444
}
func (d *FrameData) imageNRGBA() *image.NRGBA {
return &image.NRGBA{
Pix: d.planeData(0, func(linesize int) int { return linesize * d.f.Height() }),
Stride: d.f.Linesize()[0],
Rect: image.Rect(0, 0, d.f.Width(), d.f.Height()),
}
}
func (d *FrameData) imageYCbCr() *image.YCbCr {
return &image.YCbCr{
Y: d.planeData(0, func(linesize int) int { return linesize * d.f.Height() }),
Cb: d.planeData(1, func(linesize int) int { return linesize * d.f.Height() }),
Cr: d.planeData(2, func(linesize int) int { return linesize * d.f.Height() }),
YStride: d.f.Linesize()[0],
CStride: d.f.Linesize()[1],
SubsampleRatio: d.imageYCbCrSubsampleRatio(),
Rect: image.Rect(0, 0, d.f.Width(), d.f.Height()),
}
}
func (d *FrameData) imageNYCbCrA() *image.NYCbCrA {
return &image.NYCbCrA{
YCbCr: *d.imageYCbCr(),
A: d.planeData(3, func(linesize int) int { return linesize * d.f.Height() }),
AStride: d.f.Linesize()[3],
}
}
func (d *FrameData) Image() (image.Image, error) {
// Switch on pixel format
switch d.f.PixelFormat() {
// NRGBA
case PixelFormatRgba:
return d.imageNRGBA(), nil
// NYCbCrA
case PixelFormatYuva420P,
PixelFormatYuva422P,
PixelFormatYuva444P:
return d.imageNYCbCrA(), nil
// YCbCr
case PixelFormatYuv410P,
PixelFormatYuv411P, PixelFormatYuvj411P,
PixelFormatYuv420P, PixelFormatYuvj420P,
PixelFormatYuv422P, PixelFormatYuvj422P,
PixelFormatYuv440P, PixelFormatYuvj440P,
PixelFormatYuv444P, PixelFormatYuvj444P:
return d.imageYCbCr(), nil
}
return nil, fmt.Errorf("astiav: %s pixel format not handled by the Go standard image package", d.f.PixelFormat())
}