Fix IVF timestamps

pion · Jan 2, 2025 · 31d8dbc · 31d8dbc
1 parent 8b3734e
commit 31d8dbc
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 13 deletions.
diff --git a/pkg/media/ivfreader/ivfreader.go b/pkg/media/ivfreader/ivfreader.go
@@ -52,6 +52,7 @@ type IVFFrameHeader struct {
 type IVFReader struct {
 	stream               io.Reader
 	bytesReadSuccesfully int64
+	fileHeader           *IVFFileHeader
 }
 
 // NewWith returns a new IVF reader and IVF file header
@@ -69,6 +70,7 @@ func NewWith(in io.Reader) (*IVFReader, *IVFFileHeader, error) {
 	if err != nil {
 		return nil, nil, err
 	}
+	reader.fileHeader = header
 
 	return reader, header, nil
 }
@@ -80,6 +82,10 @@ func (i *IVFReader) ResetReader(reset func(bytesRead int64) io.Reader) {
 	i.stream = reset(i.bytesReadSuccesfully)
 }
 
+func (i *IVFReader) ptsToTimestamp(pts uint64) uint64 {
+	return pts * uint64(i.fileHeader.TimebaseDenominator) / uint64(i.fileHeader.TimebaseNumerator)
+}
+
 // ParseNextFrame reads from stream and returns IVF frame payload, header,
 // and an error if there is incomplete frame data.
 // Returns all nil values when no more frames are available.
@@ -95,9 +101,10 @@ func (i *IVFReader) ParseNextFrame() ([]byte, *IVFFrameHeader, error) {
 		return nil, nil, err
 	}
 
+	pts := binary.LittleEndian.Uint64(buffer[4:12])
 	header = &IVFFrameHeader{
 		FrameSize: binary.LittleEndian.Uint32(buffer[:4]),
-		Timestamp: binary.LittleEndian.Uint64(buffer[4:12]),
+		Timestamp: i.ptsToTimestamp(pts),
 	}
 
 	payload := make([]byte, header.FrameSize)

diff --git a/pkg/media/ivfwriter/ivfwriter.go b/pkg/media/ivfwriter/ivfwriter.go
@@ -37,6 +37,11 @@ type IVFWriter struct {
 
 	isVP8, isAV1 bool
 
+	timebaseDenominator uint32
+	timebaseNumerator   uint32
+	firstFrameTimestamp uint32
+	clockRate           uint64
+
 	// VP8
 	currentFrame []byte
 
@@ -65,8 +70,11 @@ func NewWith(out io.Writer, opts ...Option) (*IVFWriter, error) {
 	}
 
 	writer := &IVFWriter{
-		ioWriter:     out,
-		seenKeyFrame: false,
+		ioWriter:            out,
+		seenKeyFrame:        false,
+		timebaseDenominator: 30,
+		timebaseNumerator:   1,
+		clockRate:           90000,
 	}
 
 	for _, o := range opts {
@@ -98,21 +106,25 @@ func (i *IVFWriter) writeHeader() error {
 		copy(header[8:], "AV01")
 	}
 
-	binary.LittleEndian.PutUint16(header[12:], 640) // Width in pixels
-	binary.LittleEndian.PutUint16(header[14:], 480) // Height in pixels
-	binary.LittleEndian.PutUint32(header[16:], 30)  // Framerate denominator
-	binary.LittleEndian.PutUint32(header[20:], 1)   // Framerate numerator
-	binary.LittleEndian.PutUint32(header[24:], 900) // Frame count, will be updated on first Close() call
-	binary.LittleEndian.PutUint32(header[28:], 0)   // Unused
+	binary.LittleEndian.PutUint16(header[12:], 640)                   // Width in pixels
+	binary.LittleEndian.PutUint16(header[14:], 480)                   // Height in pixels
+	binary.LittleEndian.PutUint32(header[16:], i.timebaseDenominator) // Framerate denominator
+	binary.LittleEndian.PutUint32(header[20:], i.timebaseNumerator)   // Framerate numerator
+	binary.LittleEndian.PutUint32(header[24:], 900)                   // Frame count, will be updated on first Close() call
+	binary.LittleEndian.PutUint32(header[28:], 0)                     // Unused
 
 	_, err := i.ioWriter.Write(header)
 	return err
 }
 
-func (i *IVFWriter) writeFrame(frame []byte) error {
+func (i *IVFWriter) timestampToPts(timestamp uint64) uint64 {
+	return timestamp * uint64(i.timebaseNumerator) / uint64(i.timebaseDenominator)
+}
+
+func (i *IVFWriter) writeFrame(frame []byte, timestamp uint64) error {
 	frameHeader := make([]byte, 12)
 	binary.LittleEndian.PutUint32(frameHeader[0:], uint32(len(frame))) // Frame length
-	binary.LittleEndian.PutUint64(frameHeader[4:], i.count)            // PTS
+	binary.LittleEndian.PutUint64(frameHeader[4:], i.timestampToPts(timestamp))
 	i.count++
 
 	if _, err := i.ioWriter.Write(frameHeader); err != nil {
@@ -130,6 +142,11 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error {
 		return nil
 	}
 
+	if i.count == 0 {
+		i.firstFrameTimestamp = packet.Header.Timestamp
+	}
+	relativeTstampMs := 1000 * uint64(packet.Header.Timestamp) / i.clockRate
+
 	if i.isVP8 {
 		vp8Packet := codecs.VP8Packet{}
 		if _, err := vp8Packet.Unmarshal(packet.Payload); err != nil {
@@ -153,7 +170,7 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error {
 			return nil
 		}
 
-		if err := i.writeFrame(i.currentFrame); err != nil {
+		if err := i.writeFrame(i.currentFrame, relativeTstampMs); err != nil {
 			return err
 		}
 		i.currentFrame = nil
@@ -169,7 +186,7 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error {
 		}
 
 		for j := range obus {
-			if err := i.writeFrame(obus[j]); err != nil {
+			if err := i.writeFrame(obus[j], relativeTstampMs); err != nil {
 				return err
 			}
 		}