diff --git a/frame_data.go b/frame_data.go index edd46fd..52fb2dc 100644 --- a/frame_data.go +++ b/frame_data.go @@ -9,6 +9,38 @@ import ( "strings" ) +type frameDataImageFormat int + +const ( + frameDataImageFormatNone frameDataImageFormat = iota + frameDataImageFormatNRGBA + frameDataImageFormatNYCbCrA + frameDataImageFormatYCbCr +) + +func frameDataImageFormatFromPixelFormat(pf PixelFormat) frameDataImageFormat { + // Switch on pixel format + switch pf { + // NRGBA + case PixelFormatRgba: + return frameDataImageFormatNRGBA + // NYCbCrA + case PixelFormatYuva420P, + PixelFormatYuva422P, + PixelFormatYuva444P: + return frameDataImageFormatNYCbCrA + // YCbCr + case PixelFormatYuv410P, + PixelFormatYuv411P, PixelFormatYuvj411P, + PixelFormatYuv420P, PixelFormatYuvj420P, + PixelFormatYuv422P, PixelFormatYuvj422P, + PixelFormatYuv440P, PixelFormatYuvj440P, + PixelFormatYuv444P, PixelFormatYuvj444P: + return frameDataImageFormatYCbCr + } + return frameDataImageFormatNone +} + type FrameData struct { f *Frame } @@ -44,9 +76,9 @@ func (d *FrameData) Bytes(align int) ([]byte, error) { return nil, errors.New("astiav: frame type not implemented") } -func (d *FrameData) planeData(i int, sizeFunc func(linesize int) int) []byte { +func (d *FrameData) planeBytes(i int) []byte { return bytesFromC(func(size *cUlong) *C.uint8_t { - *size = cUlong(sizeFunc(int(d.f.c.linesize[i]))) + *size = cUlong(int(d.f.c.linesize[i]) * d.f.Height()) return d.f.c.data[i] }) } @@ -68,53 +100,99 @@ func (d *FrameData) imageYCbCrSubsampleRatio() image.YCbCrSubsampleRatio { 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) copyPlaneBytes(i int, s *[]uint8) { + b := d.planeBytes(0) + if len(b) > cap(*s) { + *s = make([]uint8, len(b)) } + copy(*s, b) } -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) toImageNRGBA(i *image.NRGBA) { + d.copyPlaneBytes(0, &i.Pix) + if v := d.f.Linesize()[0]; i.Stride != v { + i.Stride = v + } + if w, h := d.f.Width(), d.f.Height(); i.Rect.Dy() != w || i.Rect.Dx() != h { + i.Rect = image.Rect(0, 0, w, h) } } -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) toImageYCbCr(i *image.YCbCr) { + d.copyPlaneBytes(0, &i.Y) + d.copyPlaneBytes(1, &i.Cb) + d.copyPlaneBytes(2, &i.Cr) + if v := d.f.Linesize()[0]; i.YStride != v { + i.YStride = v + } + if v := d.f.Linesize()[1]; i.CStride != v { + i.CStride = v + } + if v := d.imageYCbCrSubsampleRatio(); i.SubsampleRatio != v { + i.SubsampleRatio = v + } + if w, h := d.f.Width(), d.f.Height(); i.Rect.Dy() != w || i.Rect.Dx() != h { + i.Rect = image.Rect(0, 0, w, h) + } +} + +func (d *FrameData) toImageNYCbCrA(i *image.NYCbCrA) { + d.toImageYCbCr(&i.YCbCr) + d.copyPlaneBytes(3, &i.A) + if v := d.f.Linesize()[3]; i.AStride != v { + i.AStride = v } } func (d *FrameData) Image() (image.Image, error) { - // Switch on pixel format - switch d.f.PixelFormat() { + // Switch on image format + switch frameDataImageFormatFromPixelFormat(d.f.PixelFormat()) { // NRGBA - case PixelFormatRgba: - return d.imageNRGBA(), nil + case frameDataImageFormatNRGBA: + i := &image.NRGBA{} + d.toImageNRGBA(i) + return i, nil // NYCbCrA - case PixelFormatYuva420P, - PixelFormatYuva422P, - PixelFormatYuva444P: - return d.imageNYCbCrA(), nil + case frameDataImageFormatNYCbCrA: + i := &image.NYCbCrA{} + d.toImageNYCbCrA(i) + return i, nil // YCbCr - case PixelFormatYuv410P, - PixelFormatYuv411P, PixelFormatYuvj411P, - PixelFormatYuv420P, PixelFormatYuvj420P, - PixelFormatYuv422P, PixelFormatYuvj422P, - PixelFormatYuv440P, PixelFormatYuvj440P, - PixelFormatYuv444P, PixelFormatYuvj444P: - return d.imageYCbCr(), nil + case frameDataImageFormatYCbCr: + i := &image.YCbCr{} + d.toImageYCbCr(i) + return i, nil } return nil, fmt.Errorf("astiav: %s pixel format not handled by the Go standard image package", d.f.PixelFormat()) } + +func (d *FrameData) ToImage(dst image.Image) error { + // Switch on image format + switch frameDataImageFormatFromPixelFormat(d.f.PixelFormat()) { + // NRGBA + case frameDataImageFormatNRGBA: + i, ok := dst.(*image.NRGBA) + if !ok { + return errors.New("astiav: image should be *image.NRGBA") + } + d.toImageNRGBA(i) + return nil + // NYCbCrA + case frameDataImageFormatNYCbCrA: + i, ok := dst.(*image.NYCbCrA) + if !ok { + return errors.New("astiav: image should be *image.NYCbCrA") + } + d.toImageNYCbCrA(i) + return nil + // YCbCr + case frameDataImageFormatYCbCr: + i, ok := dst.(*image.YCbCr) + if !ok { + return errors.New("astiav: image should be *image.YCbCr") + } + d.toImageYCbCr(i) + return nil + } + return fmt.Errorf("astiav: %s pixel format not handled by the Go standard image package", d.f.PixelFormat()) +} diff --git a/frame_data_test.go b/frame_data_test.go index 3732d23..6f3a701 100644 --- a/frame_data_test.go +++ b/frame_data_test.go @@ -1,6 +1,7 @@ package astiav_test import ( + "image" "image/png" "os" "testing" @@ -12,10 +13,12 @@ import ( func TestFrameData(t *testing.T) { for _, v := range []struct { ext string + i image.Image name string }{ { ext: "png", + i: &image.NRGBA{}, name: "image-rgba", }, // TODO Find a way to test yuv and yuva even though result seems to change randomly @@ -39,9 +42,11 @@ func TestFrameData(t *testing.T) { i1, err := fd.Image() require.NoError(t, err) + require.NoError(t, fd.ToImage(v.i)) i2, err := png.Decode(f1) require.NoError(t, err) require.Equal(t, i1, i2) + require.Equal(t, v.i, i2) }() } }