diff --git a/frame.go b/frame.go index 06a6cad..16e9e01 100644 --- a/frame.go +++ b/frame.go @@ -59,7 +59,7 @@ func (f *Frame) SetColorRange(r ColorRange) { } func (f *Frame) Data() *FrameData { - return newFrameData(f) + return NewFrameData(newFrameDataFrame(f)) } func (f *Frame) Height() int { diff --git a/frame_data.go b/frame_data.go index 52fb2dc..0b9c047 100644 --- a/frame_data.go +++ b/frame_data.go @@ -9,43 +9,62 @@ import ( "strings" ) -type frameDataImageFormat int +type FrameDataFrame interface { + Height() int + ImageBufferSize(align int) (int, error) + ImageCopyToBuffer(b []byte, align int) (int, error) + Linesize(i int) int + PixelFormat() PixelFormat + PlaneBytes(i int) []byte + Width() int +} -const ( - frameDataImageFormatNone frameDataImageFormat = iota - frameDataImageFormatNRGBA - frameDataImageFormatNYCbCrA - frameDataImageFormatYCbCr -) +var _ FrameDataFrame = (*frameDataFrame)(nil) -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 frameDataFrame struct { + f *Frame +} + +func newFrameDataFrame(f *Frame) *frameDataFrame { + return &frameDataFrame{f: f} +} + +func (f *frameDataFrame) Height() int { + return f.f.Height() +} + +func (f *frameDataFrame) ImageBufferSize(align int) (int, error) { + return f.f.ImageBufferSize(align) +} + +func (f *frameDataFrame) ImageCopyToBuffer(b []byte, align int) (int, error) { + return f.f.ImageCopyToBuffer(b, align) +} + +func (f *frameDataFrame) Linesize(i int) int { + return f.f.Linesize()[i] +} + +func (f *frameDataFrame) PixelFormat() PixelFormat { + return f.f.PixelFormat() +} + +func (f *frameDataFrame) PlaneBytes(i int) []byte { + return bytesFromC(func(size *cUlong) *C.uint8_t { + *size = cUlong(int(f.f.c.linesize[i]) * f.f.Height()) + return f.f.c.data[i] + }) +} + +func (f *frameDataFrame) Width() int { + return f.f.Width() } type FrameData struct { - f *Frame + f FrameDataFrame } -func newFrameData(f *Frame) *FrameData { +func NewFrameData(f FrameDataFrame) *FrameData { return &FrameData{f: f} } @@ -76,11 +95,33 @@ func (d *FrameData) Bytes(align int) ([]byte, error) { return nil, errors.New("astiav: frame type not implemented") } -func (d *FrameData) planeBytes(i int) []byte { - return bytesFromC(func(size *cUlong) *C.uint8_t { - *size = cUlong(int(d.f.c.linesize[i]) * d.f.Height()) - return d.f.c.data[i] - }) +// Always returns non-premultiplied formats when dealing with alpha channels, however this might not +// always be accurate. In this case, use your own format in .ToImage() +func (d *FrameData) GuessImageFormat() (image.Image, error) { + switch d.f.PixelFormat() { + case PixelFormatGray8: + return &image.Gray{}, nil + case PixelFormatGray16Be: + return &image.Gray16{}, nil + case PixelFormatRgb0, PixelFormat0Rgb, PixelFormatRgb4, PixelFormatRgb8: + return &image.RGBA{}, nil + case PixelFormatRgba: + return &image.NRGBA{}, nil + case PixelFormatRgba64Be: + return &image.NRGBA64{}, nil + case PixelFormatYuva420P, + PixelFormatYuva422P, + PixelFormatYuva444P: + return &image.NYCbCrA{}, nil + case PixelFormatYuv410P, + PixelFormatYuv411P, PixelFormatYuvj411P, + PixelFormatYuv420P, PixelFormatYuvj420P, + PixelFormatYuv422P, PixelFormatYuvj422P, + PixelFormatYuv440P, PixelFormatYuvj440P, + PixelFormatYuv444P, PixelFormatYuvj444P: + return &image.YCbCr{}, nil + } + return nil, fmt.Errorf("astiav: pixel format %s not handled by Go", d.f.PixelFormat()) } func (d *FrameData) imageYCbCrSubsampleRatio() image.YCbCrSubsampleRatio { @@ -101,98 +142,74 @@ func (d *FrameData) imageYCbCrSubsampleRatio() image.YCbCrSubsampleRatio { } func (d *FrameData) copyPlaneBytes(i int, s *[]uint8) { - b := d.planeBytes(0) + b := d.f.PlaneBytes(i) if len(b) > cap(*s) { *s = make([]uint8, len(b)) } copy(*s, b) } -func (d *FrameData) toImageNRGBA(i *image.NRGBA) { - d.copyPlaneBytes(0, &i.Pix) - if v := d.f.Linesize()[0]; i.Stride != v { - i.Stride = v +func (d *FrameData) toImagePix(pix *[]uint8, stride *int, rect *image.Rectangle) { + d.copyPlaneBytes(0, pix) + if v := d.f.Linesize(0); *stride != v { + *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) + if w, h := d.f.Width(), d.f.Height(); rect.Dy() != w || rect.Dx() != h { + *rect = image.Rect(0, 0, w, h) } } -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 +func (d *FrameData) toImageYCbCr(y, cb, cr *[]uint8, yStride, cStride *int, subsampleRatio *image.YCbCrSubsampleRatio, rect *image.Rectangle) { + d.copyPlaneBytes(0, y) + d.copyPlaneBytes(1, cb) + d.copyPlaneBytes(2, cr) + if v := d.f.Linesize(0); *yStride != v { + *yStride = v } - if v := d.f.Linesize()[1]; i.CStride != v { - i.CStride = v + if v := d.f.Linesize(1); *cStride != v { + *cStride = v } - if v := d.imageYCbCrSubsampleRatio(); i.SubsampleRatio != v { - i.SubsampleRatio = v + if v := d.imageYCbCrSubsampleRatio(); *subsampleRatio != v { + *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) + if w, h := d.f.Width(), d.f.Height(); rect.Dy() != w || rect.Dx() != h { + *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) toImageYCbCrA(y, cb, cr, a *[]uint8, yStride, cStride, aStride *int, subsampleRatio *image.YCbCrSubsampleRatio, rect *image.Rectangle) { + d.toImageYCbCr(y, cb, cr, yStride, cStride, subsampleRatio, rect) + d.copyPlaneBytes(3, a) + if v := d.f.Linesize(3); *aStride != v { + *aStride = v } } -func (d *FrameData) Image() (image.Image, error) { - // Switch on image format - switch frameDataImageFormatFromPixelFormat(d.f.PixelFormat()) { - // NRGBA - case frameDataImageFormatNRGBA: - i := &image.NRGBA{} - d.toImageNRGBA(i) - return i, nil - // NYCbCrA - case frameDataImageFormatNYCbCrA: - i := &image.NYCbCrA{} - d.toImageNYCbCrA(i) - return i, nil - // YCbCr - 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 + if v, ok := dst.(*image.Alpha); ok { + d.toImagePix(&v.Pix, &v.Stride, &v.Rect) + } else if v, ok := dst.(*image.Alpha16); ok { + d.toImagePix(&v.Pix, &v.Stride, &v.Rect) + } else if v, ok := dst.(*image.CMYK); ok { + d.toImagePix(&v.Pix, &v.Stride, &v.Rect) + } else if v, ok := dst.(*image.Gray); ok { + d.toImagePix(&v.Pix, &v.Stride, &v.Rect) + } else if v, ok := dst.(*image.Gray16); ok { + d.toImagePix(&v.Pix, &v.Stride, &v.Rect) + } else if v, ok := dst.(*image.NRGBA); ok { + d.toImagePix(&v.Pix, &v.Stride, &v.Rect) + } else if v, ok := dst.(*image.NRGBA64); ok { + d.toImagePix(&v.Pix, &v.Stride, &v.Rect) + } else if v, ok := dst.(*image.NYCbCrA); ok { + d.toImageYCbCrA(&v.Y, &v.Cb, &v.Cr, &v.A, &v.YStride, &v.CStride, &v.AStride, &v.SubsampleRatio, &v.Rect) + } else if v, ok := dst.(*image.RGBA); ok { + d.toImagePix(&v.Pix, &v.Stride, &v.Rect) + } else if v, ok := dst.(*image.RGBA64); ok { + d.toImagePix(&v.Pix, &v.Stride, &v.Rect) + } else if v, ok := dst.(*image.YCbCr); ok { + d.toImageYCbCr(&v.Y, &v.Cb, &v.Cr, &v.YStride, &v.CStride, &v.SubsampleRatio, &v.Rect) + } else { + return errors.New("astiav: image format is not handled") } - return fmt.Errorf("astiav: %s pixel format not handled by the Go standard image package", d.f.PixelFormat()) + return nil } diff --git a/frame_data_test.go b/frame_data_test.go index 6f3a701..6b33d71 100644 --- a/frame_data_test.go +++ b/frame_data_test.go @@ -2,51 +2,285 @@ package astiav_test import ( "image" - "image/png" - "os" "testing" "github.com/asticode/go-astiav" "github.com/stretchr/testify/require" ) +type frameDataFrame struct { + height int + imageBytes []byte + linesizes []int + pixelFormat astiav.PixelFormat + planesBytes [][]byte + width int +} + +var _ astiav.FrameDataFrame = (*frameDataFrame)(nil) + +func (f *frameDataFrame) Height() int { + return f.height +} + +func (f *frameDataFrame) ImageBufferSize(align int) (int, error) { + return len(f.imageBytes), nil +} + +func (f *frameDataFrame) ImageCopyToBuffer(b []byte, align int) (int, error) { + copy(b, f.imageBytes) + return len(f.imageBytes), nil +} + +func (f *frameDataFrame) Linesize(i int) int { + return f.linesizes[i] +} + +func (f *frameDataFrame) PixelFormat() astiav.PixelFormat { + return f.pixelFormat +} + +func (f *frameDataFrame) PlaneBytes(i int) []byte { + return f.planesBytes[i] +} + +func (f *frameDataFrame) Width() int { + return f.width +} + func TestFrameData(t *testing.T) { + f := &frameDataFrame{} + fd := astiav.NewFrameData(f) + + for _, v := range []struct { + err bool + i image.Image + pfs []astiav.PixelFormat + }{ + { + i: &image.Gray{}, + pfs: []astiav.PixelFormat{astiav.PixelFormatGray8}, + }, + { + i: &image.Gray16{}, + pfs: []astiav.PixelFormat{astiav.PixelFormatGray16Be}, + }, + { + i: &image.RGBA{}, + pfs: []astiav.PixelFormat{ + astiav.PixelFormatRgb0, + astiav.PixelFormat0Rgb, + astiav.PixelFormatRgb4, + astiav.PixelFormatRgb8, + }, + }, + { + i: &image.NRGBA{}, + pfs: []astiav.PixelFormat{astiav.PixelFormatRgba}, + }, + { + i: &image.NRGBA64{}, + pfs: []astiav.PixelFormat{astiav.PixelFormatRgba64Be}, + }, + { + i: &image.NYCbCrA{}, + pfs: []astiav.PixelFormat{ + astiav.PixelFormatYuva420P, + astiav.PixelFormatYuva422P, + astiav.PixelFormatYuva444P, + }, + }, + { + i: &image.YCbCr{}, + pfs: []astiav.PixelFormat{ + astiav.PixelFormatYuv410P, + astiav.PixelFormatYuv411P, + astiav.PixelFormatYuvj411P, + astiav.PixelFormatYuv420P, + astiav.PixelFormatYuvj420P, + astiav.PixelFormatYuv422P, + astiav.PixelFormatYuvj422P, + astiav.PixelFormatYuv440P, + astiav.PixelFormatYuvj440P, + astiav.PixelFormatYuv444P, + astiav.PixelFormatYuvj444P, + }, + }, + { + err: true, + pfs: []astiav.PixelFormat{astiav.PixelFormatAbgr}, + }, + } { + for _, pf := range v.pfs { + f.pixelFormat = pf + i, err := fd.GuessImageFormat() + if v.err { + require.Error(t, err) + } else { + require.IsType(t, v.i, i) + } + } + } + + f.imageBytes = []byte{0, 1, 2, 3} + _, err := fd.Bytes(0) + require.Error(t, err) + f.height = 1 + f.width = 2 + b, err := fd.Bytes(0) + require.NoError(t, err) + require.Equal(t, f.imageBytes, b) + for _, v := range []struct { - ext string - i image.Image - name string + e image.Image + err bool + i image.Image + linesizes []int + pixelFormat astiav.PixelFormat + planesBytes [][]byte }{ { - ext: "png", - i: &image.NRGBA{}, - name: "image-rgba", + e: &image.Alpha{ + Pix: []byte{0, 1, 2, 3}, + Stride: 1, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.Alpha{}, + linesizes: []int{1}, + pixelFormat: astiav.PixelFormatRgba, + planesBytes: [][]byte{{0, 1, 2, 3}}, + }, + { + e: &image.Alpha16{ + Pix: []byte{0, 1, 2, 3}, + Stride: 1, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.Alpha16{}, + linesizes: []int{1}, + pixelFormat: astiav.PixelFormatRgba, + planesBytes: [][]byte{{0, 1, 2, 3}}, + }, + { + e: &image.CMYK{ + Pix: []byte{0, 1, 2, 3}, + Stride: 1, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.CMYK{}, + linesizes: []int{1}, + pixelFormat: astiav.PixelFormatRgba, + planesBytes: [][]byte{{0, 1, 2, 3}}, + }, + { + e: &image.Gray{ + Pix: []byte{0, 1, 2, 3}, + Stride: 1, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.Gray{}, + linesizes: []int{1}, + pixelFormat: astiav.PixelFormatRgba, + planesBytes: [][]byte{{0, 1, 2, 3}}, + }, + { + e: &image.Gray16{ + Pix: []byte{0, 1, 2, 3}, + Stride: 1, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.Gray16{}, + linesizes: []int{1}, + pixelFormat: astiav.PixelFormatRgba, + planesBytes: [][]byte{{0, 1, 2, 3}}, + }, + { + e: &image.NRGBA{ + Pix: []byte{0, 1, 2, 3}, + Stride: 1, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.NRGBA{}, + linesizes: []int{1}, + pixelFormat: astiav.PixelFormatRgba, + planesBytes: [][]byte{{0, 1, 2, 3}}, + }, + { + e: &image.NRGBA64{ + Pix: []byte{0, 1, 2, 3}, + Stride: 1, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.NRGBA64{}, + linesizes: []int{1}, + pixelFormat: astiav.PixelFormatRgba, + planesBytes: [][]byte{{0, 1, 2, 3}}, + }, + { + e: &image.NYCbCrA{ + A: []byte{6, 7}, + AStride: 4, + YCbCr: image.YCbCr{ + Y: []byte{0, 1}, + Cb: []byte{2, 3}, + Cr: []byte{4, 5}, + YStride: 1, + CStride: 2, + SubsampleRatio: image.YCbCrSubsampleRatio444, + Rect: image.Rect(0, 0, 2, 1), + }, + }, + i: &image.NYCbCrA{}, + linesizes: []int{1, 2, 3, 4}, + pixelFormat: astiav.PixelFormatYuv444P, + planesBytes: [][]byte{{0, 1}, {2, 3}, {4, 5}, {6, 7}}, + }, + { + e: &image.RGBA{ + Pix: []byte{0, 1, 2, 3}, + Stride: 1, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.RGBA{}, + linesizes: []int{1}, + pixelFormat: astiav.PixelFormatRgba, + planesBytes: [][]byte{{0, 1, 2, 3}}, + }, + { + e: &image.RGBA64{ + Pix: []byte{0, 1, 2, 3}, + Stride: 1, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.RGBA64{}, + linesizes: []int{1}, + pixelFormat: astiav.PixelFormatRgba, + planesBytes: [][]byte{{0, 1, 2, 3}}, + }, + { + e: &image.YCbCr{ + Y: []byte{0, 1}, + Cb: []byte{2, 3}, + Cr: []byte{4, 5}, + YStride: 1, + CStride: 2, + SubsampleRatio: image.YCbCrSubsampleRatio420, + Rect: image.Rect(0, 0, 2, 1), + }, + i: &image.YCbCr{}, + linesizes: []int{1, 2, 3}, + pixelFormat: astiav.PixelFormatYuv420P, + planesBytes: [][]byte{{0, 1}, {2, 3}, {4, 5}}, }, - // TODO Find a way to test yuv and yuva even though result seems to change randomly } { - // We use a closure to ease closing files - func() { - f, err := globalHelper.inputLastFrame(v.name+"."+v.ext, astiav.MediaTypeVideo) - require.NoError(t, err) - fd := f.Data() - - b1, err := fd.Bytes(1) - require.NoError(t, err) - - b2, err := os.ReadFile("testdata/" + v.name + "-bytes") - require.NoError(t, err) - require.Equal(t, b1, b2) - - f1, err := os.Open("testdata/" + v.name + "." + v.ext) - require.NoError(t, err) - defer f1.Close() - - 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) - }() + f.linesizes = v.linesizes + f.pixelFormat = v.pixelFormat + f.planesBytes = v.planesBytes + err = fd.ToImage(v.i) + if v.err { + require.Error(t, err) + } else { + require.Equal(t, v.e, v.i) + } } } diff --git a/testdata/image-rgba-bytes b/testdata/image-rgba-bytes deleted file mode 100644 index 55092c8..0000000 Binary files a/testdata/image-rgba-bytes and /dev/null differ diff --git a/testdata/image-rgba.png b/testdata/image-rgba.png deleted file mode 100644 index d15e9af..0000000 Binary files a/testdata/image-rgba.png and /dev/null differ