Add streaming to rawvideo format

AlexxIT · May 25, 2024 · 6f34cf0 · 6f34cf0
1 parent f8bc25d
commit 6f34cf0
Show file tree

Hide file tree

Showing 5 changed files with 207 additions and 104 deletions.
diff --git a/internal/mjpeg/init.go b/internal/mjpeg/init.go
@@ -17,13 +17,15 @@ import (
 	"github.com/AlexxIT/go2rtc/pkg/magic"
 	"github.com/AlexxIT/go2rtc/pkg/mjpeg"
 	"github.com/AlexxIT/go2rtc/pkg/tcp"
+	"github.com/AlexxIT/go2rtc/pkg/y4m"
 	"github.com/rs/zerolog/log"
 )
 
 func Init() {
 	api.HandleFunc("api/frame.jpeg", handlerKeyframe)
 	api.HandleFunc("api/stream.mjpeg", handlerStream)
 	api.HandleFunc("api/stream.ascii", handlerStream)
+	api.HandleFunc("api/stream.y4m", apiStreamY4M)
 
 	ws.HandleFunc("mjpeg", handlerWS)
 }
@@ -166,3 +168,25 @@ func handlerWS(tr *ws.Transport, _ *ws.Message) error {
 
 	return nil
 }
+
+func apiStreamY4M(w http.ResponseWriter, r *http.Request) {
+	src := r.URL.Query().Get("src")
+	stream := streams.Get(src)
+	if stream == nil {
+		http.Error(w, api.StreamNotFound, http.StatusNotFound)
+		return
+	}
+
+	cons := y4m.NewConsumer()
+	cons.RemoteAddr = tcp.RemoteAddr(r)
+	cons.UserAgent = r.UserAgent()
+
+	if err := stream.AddConsumer(cons); err != nil {
+		log.Error().Err(err).Caller().Send()
+		return
+	}
+
+	_, _ = cons.WriteTo(w)
+
+	stream.RemoveConsumer(cons)
+}
diff --git a/pkg/y4m/README.md b/pkg/y4m/README.md
@@ -0,0 +1,5 @@
+## Useful links
+
+- https://learn.microsoft.com/en-us/windows/win32/medfound/recommended-8-bit-yuv-formats-for-video-rendering
+- https://developer.mozilla.org/en-US/docs/Web/Media/Formats/Video_concepts
+- https://fourcc.org/yuv.php#YV12
diff --git a/pkg/y4m/consumer.go b/pkg/y4m/consumer.go
@@ -0,0 +1,67 @@
+package y4m
+
+import (
+	"fmt"
+	"io"
+
+	"github.com/AlexxIT/go2rtc/pkg/core"
+	"github.com/pion/rtp"
+)
+
+type Consumer struct {
+	core.SuperConsumer
+	wr *core.WriteBuffer
+}
+
+func NewConsumer() *Consumer {
+	return &Consumer{
+		core.SuperConsumer{
+			Type: "YUV4MPEG2 passive consumer",
+			Medias: []*core.Media{
+				{
+					Kind:      core.KindVideo,
+					Direction: core.DirectionSendonly,
+					Codecs: []*core.Codec{
+						{Name: core.CodecRAW},
+					},
+				},
+			},
+		},
+		core.NewWriteBuffer(nil),
+	}
+}
+
+func (c *Consumer) AddTrack(media *core.Media, _ *core.Codec, track *core.Receiver) error {
+	sender := core.NewSender(media, track.Codec)
+	sender.Handler = func(packet *rtp.Packet) {
+		if n, err := c.wr.Write([]byte(frameHdr)); err == nil {
+			c.Send += n
+		}
+		if n, err := c.wr.Write(packet.Payload); err == nil {
+			c.Send += n
+		}
+	}
+
+	hdr := fmt.Sprintf(
+		"YUV4MPEG2 W%s H%s C%s\n",
+		core.Between(track.Codec.FmtpLine, "width=", ";"),
+		core.Between(track.Codec.FmtpLine, "height=", ";"),
+		core.Between(track.Codec.FmtpLine, "colorspace=", ";"),
+	)
+	if _, err := c.wr.Write([]byte(hdr)); err != nil {
+		return err
+	}
+
+	sender.HandleRTP(track)
+	c.Senders = append(c.Senders, sender)
+	return nil
+}
+
+func (c *Consumer) WriteTo(wr io.Writer) (int64, error) {
+	return c.wr.WriteTo(wr)
+}
+
+func (c *Consumer) Stop() error {
+	_ = c.SuperConsumer.Close()
+	return c.wr.Close()
+}
diff --git a/pkg/y4m/producer.go b/pkg/y4m/producer.go
@@ -0,0 +1,110 @@
+package y4m
+
+import (
+	"bufio"
+	"bytes"
+	"errors"
+	"io"
+
+	"github.com/AlexxIT/go2rtc/pkg/core"
+	"github.com/pion/rtp"
+)
+
+func Open(r io.Reader) (*Producer, error) {
+	rd := bufio.NewReaderSize(r, core.BufferSize)
+	b, err := rd.ReadBytes('\n')
+	if err != nil {
+		return nil, err
+	}
+
+	b = b[:len(b)-1] // remove \n
+
+	sdp := string(b)
+	var fmtp string
+
+	for b != nil {
+		// YUV4MPEG2 W1280 H720 F24:1 Ip A1:1 C420mpeg2 XYSCSS=420MPEG2
+		// https://manned.org/yuv4mpeg.5
+		// https://github.com/FFmpeg/FFmpeg/blob/master/libavformat/yuv4mpegenc.c
+		key := b[0]
+		var value string
+		if i := bytes.IndexByte(b, ' '); i > 0 {
+			value = string(b[1:i])
+			b = b[i+1:]
+		} else {
+			value = string(b[1:])
+			b = nil
+		}
+
+		switch key {
+		case 'W':
+			fmtp = "width=" + value
+		case 'H':
+			fmtp += ";height=" + value
+		case 'C':
+			fmtp += ";colorspace=" + value
+		}
+	}
+
+	if GetSize(fmtp) == 0 {
+		return nil, errors.New("y4m: unsupported format: " + sdp)
+	}
+
+	prod := &Producer{rd: rd, cl: r.(io.Closer)}
+	prod.Type = "YUV4MPEG2 producer"
+	prod.SDP = sdp
+	prod.Medias = []*core.Media{
+		{
+			Kind:      core.KindVideo,
+			Direction: core.DirectionRecvonly,
+			Codecs: []*core.Codec{
+				{
+					Name:        core.CodecRAW,
+					ClockRate:   90000,
+					FmtpLine:    fmtp,
+					PayloadType: core.PayloadTypeRAW,
+				},
+			},
+		},
+	}
+
+	return prod, nil
+}
+
+type Producer struct {
+	core.SuperProducer
+	rd *bufio.Reader
+	cl io.Closer
+}
+
+func (c *Producer) Start() error {
+	size := GetSize(c.Medias[0].Codecs[0].FmtpLine)
+
+	for {
+		if _, err := c.rd.Discard(len(frameHdr)); err != nil {
+			return err
+		}
+
+		frame := make([]byte, size)
+		if _, err := io.ReadFull(c.rd, frame); err != nil {
+			return err
+		}
+
+		c.Recv += size
+
+		if len(c.Receivers) == 0 {
+			continue
+		}
+
+		pkt := &rtp.Packet{
+			Header:  rtp.Header{Timestamp: core.Now90000()},
+			Payload: frame,
+		}
+		c.Receivers[0].WriteRTP(pkt)
+	}
+}
+
+func (c *Producer) Stop() error {
+	_ = c.SuperProducer.Close()
+	return c.cl.Close()
+}
diff --git a/pkg/y4m/y4m.go b/pkg/y4m/y4m.go
@@ -1,117 +1,14 @@
 package y4m
 
 import (
-	"bufio"
-	"bytes"
-	"errors"
 	"image"
-	"io"
 
 	"github.com/AlexxIT/go2rtc/pkg/core"
-	"github.com/pion/rtp"
 )
 
 const FourCC = "YUV4"
 
-func Open(r io.Reader) (*Producer, error) {
-	rd := bufio.NewReaderSize(r, core.BufferSize)
-	b, err := rd.ReadBytes('\n')
-	if err != nil {
-		return nil, err
-	}
-
-	b = b[:len(b)-1] // remove \n
-
-	sdp := string(b)
-	var fmtp string
-
-	for b != nil {
-		// YUV4MPEG2 W1280 H720 F24:1 Ip A1:1 C420mpeg2 XYSCSS=420MPEG2
-		// https://manned.org/yuv4mpeg.5
-		// https://github.com/FFmpeg/FFmpeg/blob/master/libavformat/yuv4mpegenc.c
-		key := b[0]
-		var value string
-		if i := bytes.IndexByte(b, ' '); i > 0 {
-			value = string(b[1:i])
-			b = b[i+1:]
-		} else {
-			value = string(b[1:])
-			b = nil
-		}
-
-		switch key {
-		case 'W':
-			fmtp = "width=" + value
-		case 'H':
-			fmtp += ";height=" + value
-		case 'C':
-			fmtp += ";colorspace=" + value
-		}
-	}
-
-	if GetSize(fmtp) == 0 {
-		return nil, errors.New("y4m: unsupported format: " + sdp)
-	}
-
-	prod := &Producer{rd: rd, cl: r.(io.Closer)}
-	prod.Type = "YUV4MPEG2 producer"
-	prod.SDP = sdp
-	prod.Medias = []*core.Media{
-		{
-			Kind:      core.KindVideo,
-			Direction: core.DirectionRecvonly,
-			Codecs: []*core.Codec{
-				{
-					Name:        core.CodecRAW,
-					ClockRate:   90000,
-					FmtpLine:    fmtp,
-					PayloadType: core.PayloadTypeRAW,
-				},
-			},
-		},
-	}
-
-	return prod, nil
-}
-
-type Producer struct {
-	core.SuperProducer
-	rd *bufio.Reader
-	cl io.Closer
-}
-
-func (c *Producer) Start() error {
-	size := GetSize(c.Medias[0].Codecs[0].FmtpLine)
-
-	for {
-		// FRAME\n
-		if _, err := c.rd.Discard(6); err != nil {
-			return err
-		}
-
-		frame := make([]byte, size)
-		if _, err := io.ReadFull(c.rd, frame); err != nil {
-			return err
-		}
-
-		c.Recv += size
-
-		if len(c.Receivers) == 0 {
-			continue
-		}
-
-		pkt := &rtp.Packet{
-			Header:  rtp.Header{Timestamp: core.Now90000()},
-			Payload: frame,
-		}
-		c.Receivers[0].WriteRTP(pkt)
-	}
-}
-
-func (c *Producer) Stop() error {
-	_ = c.SuperProducer.Close()
-	return c.cl.Close()
-}
+const frameHdr = "FRAME\n"
 
 func GetSize(fmtp string) int {
 	w := core.Atoi(core.Between(fmtp, "width=", ";"))