Replaced FrameData.Image() with FrameData.GuessImageFormat() + now handling premultiplied image formats

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 {

frame_data.go

@@ -9,43 +9,62 @@ 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 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
 }
 
-type FrameData struct {
+var _ FrameDataFrame = (*frameDataFrame)(nil)
+
+type frameDataFrame struct {
 	f *Frame
 }
 
-func newFrameData(f *Frame) *FrameData {
+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 FrameDataFrame
+}
+
+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
 }

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 {
-		ext  string
-		i    image.Image
-		name string
+		err bool
+		i   image.Image
+		pfs []astiav.PixelFormat
 	}{
 		{
-			ext:  "png",
-			i:    &image.NRGBA{},
-			name: "image-rgba",
+			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},
 		},
-		// 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()
+		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)
+			}
+		}
+	}
 
-			b1, err := fd.Bytes(1)
-			require.NoError(t, err)
+	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)
 
-			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)
-		}()
+	for _, v := range []struct {
+		e           image.Image
+		err         bool
+		i           image.Image
+		linesizes   []int
+		pixelFormat astiav.PixelFormat
+		planesBytes [][]byte
+	}{
+		{
+			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}},
+		},
+	} {
+		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)
+		}
 	}
 }