Added FrameData ToImage()

asticode · Jan 24, 2024 · 0d3468d · 0d3468d
1 parent 28e6b9d
commit 0d3468d
Show file tree

Hide file tree

Showing 2 changed files with 119 additions and 36 deletions.
diff --git 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
@@ -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)
 		}()
 	}
 }