ace/drivers/video/ppm.go

package video

import (
	"context"
	"image"
	"image/color"
	"io"

	"github.com/nfnt/resize"
	"github.com/pion/mediadevices/pkg/frame"
	"github.com/pion/mediadevices/pkg/io/video"
	"github.com/pion/mediadevices/pkg/prop"
	"github.com/sirupsen/logrus"

	_ "github.com/jbuchbinder/gopnm"
)

const DefaultFPS float32 = 60.0

type PPMStreamDriver struct {
	Height, Width int
	FPS           float32
	PPMImage      <-chan io.ReadCloser
	closed        <-chan struct{}
	cancel        func()
}

func (v *PPMStreamDriver) Open() error {
	defer logrus.Debug("device opened")
	ctx, cancel := context.WithCancel(context.Background())
	v.closed = ctx.Done()
	v.cancel = cancel
	return nil
}

func (v *PPMStreamDriver) Close() error {
	v.cancel()
	return nil
}

func (v *PPMStreamDriver) Properties() []prop.Media {
	return []prop.Media{
		{
			Video: prop.Video{
				Width:       v.Width,
				Height:      v.Height,
				FrameRate:   v.FPS,
				FrameFormat: frame.FormatYUYV,
			},
		},
	}
}

func (v *PPMStreamDriver) VideoRecord(p prop.Media) (video.Reader, error) {
	logrus.Debug(p)
	r := video.ReaderFunc(func() (img image.Image, release func(), err error) {
		canvas := image.NewYCbCr(
			image.Rect(0, 0, p.Width, p.Height),
			image.YCbCrSubsampleRatio420,
		)

		select {
		case <-v.closed:
			return nil, func() {}, io.EOF
		case ppmF, ok := <-v.PPMImage:
			if !ok {
				return nil, func() {}, io.ErrClosedPipe
			}
			defer ppmF.Close()

			img, _, err := image.Decode(ppmF)
			if err != nil {
				return nil, func() {}, err
			}

			// resize image and draw it to canvas
			resized := resize.Resize(uint(p.Width), uint(p.Height), img, resize.Lanczos3)
			for x := 0; x < resized.Bounds().Dx(); x++ {
				for y := 0; y < resized.Bounds().Dy(); y++ {
					r, g, b, _ := resized.At(x, y).RGBA()
					Y, Cb, Cr := color.RGBToYCbCr(uint8(r), uint8(g), uint8(b))

					canvas.Y[canvas.YOffset(x, y)] = Y
					canvas.Cb[canvas.COffset(x, y)] = Cb
					canvas.Cr[canvas.COffset(x, y)] = Cr
				}
			}
		}

		return canvas, func() {}, nil
	})
	return r, nil
}