Add auto transcoding to JPEG snapshot

pkg/h264/avc.go

@@ -26,6 +26,19 @@ func AnnexB2AVC(b []byte) []byte {
 	return b
 }
 
+func AVCtoAnnexB(b []byte) []byte {
+	b = bytes.Clone(b)
+	for i := 0; i < len(b); {
+		size := int(binary.BigEndian.Uint32(b[i:]))
+		b[i] = 0
+		b[i+1] = 0
+		b[i+2] = 0
+		b[i+3] = 1
+		i += 4 + size
+	}
+	return b
+}
+
 const forbiddenZeroBit = 0x80
 const nalUnitType = 0x1F
 

pkg/pipe/keyframe.go

@@ -0,0 +1,91 @@
+package pipe
+
+import (
+	"github.com/AlexxIT/go2rtc/pkg/core"
+	"github.com/AlexxIT/go2rtc/pkg/h264"
+	"github.com/AlexxIT/go2rtc/pkg/h265"
+	"github.com/AlexxIT/go2rtc/pkg/mjpeg"
+	"github.com/pion/rtp"
+)
+
+type Keyframe struct {
+	core.Listener
+
+	UserAgent  string
+	RemoteAddr string
+
+	medias []*core.Media
+	sender *core.Sender
+}
+
+func (k *Keyframe) GetMedias() []*core.Media {
+	if k.medias == nil {
+		k.medias = append(k.medias, &core.Media{
+			Kind:      core.KindVideo,
+			Direction: core.DirectionSendonly,
+			Codecs: []*core.Codec{
+				{Name: core.CodecH264},
+				{Name: core.CodecH265},
+				{Name: core.CodecJPEG},
+			},
+		})
+	}
+	return k.medias
+}
+
+func (k *Keyframe) AddTrack(media *core.Media, _ *core.Codec, track *core.Receiver) error {
+	var handler core.HandlerFunc
+
+	switch track.Codec.Name {
+	case core.CodecH264:
+		handler = func(packet *rtp.Packet) {
+			if !h264.IsKeyframe(packet.Payload) {
+				return
+			}
+			b := h264.AVCtoAnnexB(packet.Payload)
+			k.Fire(b)
+		}
+
+		if track.Codec.IsRTP() {
+			handler = h264.RTPDepay(track.Codec, handler)
+		}
+	case core.CodecH265:
+		handler = func(packet *rtp.Packet) {
+			if !h265.IsKeyframe(packet.Payload) {
+				return
+			}
+			k.Fire(packet.Payload)
+		}
+
+		if track.Codec.IsRTP() {
+			handler = h265.RTPDepay(track.Codec, handler)
+		}
+	case core.CodecJPEG:
+		handler = func(packet *rtp.Packet) {
+			k.Fire(packet.Payload)
+		}
+
+		if track.Codec.IsRTP() {
+			handler = mjpeg.RTPDepay(handler)
+		}
+	}
+
+	k.sender = core.NewSender(media, track.Codec)
+	k.sender.Handler = handler
+	k.sender.HandleRTP(track)
+	return nil
+}
+
+func (k *Keyframe) CodecName() string {
+	if k.sender != nil {
+		return k.sender.Codec.Name
+	}
+	return ""
+}
+
+func (k *Keyframe) Stop() error {
+	if k.sender != nil {
+		k.sender.Close()
+	}
+	return nil
+}