2025-03-26 16:30:50 +05:00

95 lines
2.3 KiB
Go

package formats
import (
"errors"
"fmt"
"github.com/bluenviron/gortsplib/v4"
"github.com/bluenviron/gortsplib/v4/pkg/description"
"github.com/bluenviron/gortsplib/v4/pkg/format"
"github.com/bluenviron/gortsplib/v4/pkg/format/rtph264"
"github.com/bluenviron/mediacommon/v2/pkg/codecs/h264"
"github.com/pion/rtp"
"image"
)
// FindH264Format finds the H264 media and format.
func FindH264Format(desc *description.Session) (*format.H264, *description.Media, error) {
var h264Format *format.H264
h264Media := desc.FindFormat(&h264Format)
if h264Media == nil {
return nil, nil, errors.New("media H264 not found")
}
return h264Format, h264Media, nil
}
// ProcessH264 processes H264 flow and returns PTS and AU.
func ProcessH264(
c *gortsplib.Client,
h264Media *description.Media,
h264RTPDec *rtph264.Decoder,
pkt *rtp.Packet,
) (
int64,
[][]byte,
error) {
// Decode timestamp.
pts, ok := c.PacketPTS2(h264Media, pkt)
if !ok {
return 0, nil, fmt.Errorf("waiting for timestamp\n")
}
// Extract access unit from RTP packets.
au, err := h264RTPDec.Decode(pkt)
if err != nil {
if err != rtph264.ErrNonStartingPacketAndNoPrevious && err != rtph264.ErrMorePacketsNeeded {
return 0, nil, fmt.Errorf("decoding RTP packet error: %w", err)
}
}
return pts, au, nil
}
// ProcessH264RGBA processes H264 flow and returns PTS and IMG.
func ProcessH264RGBA(
c *gortsplib.Client,
h264Media *description.Media,
h264RTPDec *rtph264.Decoder,
h264Dec *H264Decoder,
pkt *rtp.Packet,
firstRandomAccess bool,
) (
int64,
*image.RGBA,
error) {
// Decode timestamp.
pts, ok := c.PacketPTS2(h264Media, pkt)
if !ok {
return 0, nil, fmt.Errorf("waiting for timestamp\n")
}
// Extract access units from RTP packets.
au, err := h264RTPDec.Decode(pkt)
if err != rtph264.ErrNonStartingPacketAndNoPrevious && err != rtph264.ErrMorePacketsNeeded {
return 0, nil, fmt.Errorf("decoding RTP packet error: %w", err)
}
// Wait for a random access unit.
if !firstRandomAccess && !h264.IsRandomAccess(au) {
return 0, nil, nil
}
firstRandomAccess = true
// Convert H264 access units into RGBA frames.
img, err := h264Dec.Decode(au)
if err != nil {
return 0, nil, fmt.Errorf("convert into RGBA frames error\n")
}
// Wait for a frame.
if img == nil {
return 0, nil, fmt.Errorf("frame not found\n")
}
return pts, img, nil
}