Rework timestamp sync (resolves #17)

Now syncs A/V timestamps with HLS timestamps, fixing A/V desync for some older episodes
xypwn · Jun 26, 2024 · 2538745 · 2538745
1 parent 8f57cd4
commit 2538745
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 87 deletions.
diff --git a/pkg/southpark/av.go b/pkg/southpark/av.go
@@ -1,104 +1,120 @@
 package southpark
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"io"
+	"os"
 
 	"github.com/yapingcat/gomedia/go-codec"
 	"github.com/yapingcat/gomedia/go-mp4"
 	"github.com/yapingcat/gomedia/go-mpeg2"
 )
 
-func ConvertTSAndAACToMP4(tsInput io.Reader, aacInput io.Reader, mp4Output io.WriteSeeker) error {
-	var nextAACFrame func() ([]byte, uint64, uint64, bool)
-	{
-		const samplesPerFrame = 1024
-		const sampleRate = 48000
-		const millisecondsPerFrame = float64(samplesPerFrame) / float64(sampleRate) * 1000
-
-		aacData, err := io.ReadAll(aacInput)
-		if err != nil {
-			return fmt.Errorf("read aac: %w", err)
-		}
-
-		var aacFrameBuf [][]byte
-		codec.SplitAACFrame(aacData, func(frame []byte) {
-			aacFrameBuf = append(aacFrameBuf, frame)
-		})
-
-		pts := float64(0) // in ms
-		dts := float64(0) // in ms
-		nextAACFrame = func() ([]byte, uint64, uint64, bool) {
-			if len(aacFrameBuf) == 0 {
-				return nil, uint64(pts), uint64(dts), false
-			}
-
-			defer func() {
-				pts += millisecondsPerFrame
-				dts += millisecondsPerFrame
-			}()
+type SegmentFile struct {
+	Filename string
+	Duration float64
+}
 
-			res := aacFrameBuf[0]
-			aacFrameBuf = aacFrameBuf[1:]
-			return res, uint64(pts), uint64(dts), true
-		}
-	}
+func ConvertTSAndAACToMP4(tsInput []SegmentFile, aacInput []SegmentFile, mp4Output io.WriteSeeker, onProgress func(progress float64)) error {
+	//var _vdts uint64
+	//var _vpts uint64
 
 	muxer, err := mp4.CreateMp4Muxer(mp4Output)
 	if err != nil {
 		return fmt.Errorf("create mp4 muxer: %w", err)
 	}
 
-	var writeErr error
+	var onFrameErr error
+
+	const aacSamplesPerFrame = 1024
+	const aacSampleRate = 48000
+	const aacMillisecondsPerFrame = float64(aacSamplesPerFrame) / float64(aacSampleRate) * 1000
+	apts := float64(0)
+	adts := float64(0)
+	var aacFrameBuf [][]byte
 
 	// https://github.com/yapingcat/gomedia/blob/main/example/example_convert_ts_to_mp4.go
 	hasAudio := false
 	hasVideo := false
 	var atid uint32 = 0
 	var vtid uint32 = 0
-	prevADTS := uint64(0)
+	prevADTS := float64(0)
 	tsDemuxer := mpeg2.NewTSDemuxer()
 	tsDemuxer.OnFrame = func(cid mpeg2.TS_STREAM_TYPE, vframe []byte, vpts uint64 /* in ms */, vdts uint64 /* in ms */) {
 		if cid == mpeg2.TS_STREAM_H264 {
-			if !hasVideo {
-				vtid = muxer.AddVideoTrack(mp4.MP4_CODEC_H264)
-				hasVideo = true
-			}
-			if err := muxer.Write(vtid, vframe, uint64(vpts), uint64(vdts)); err != nil {
-				writeErr = err
-			}
-
-			for vdts > prevADTS {
+			//for uint64(float64(vdts) * 1.000642) > prevADTS {
+			for float64(vdts) > prevADTS {
 				if !hasAudio {
 					atid = muxer.AddAudioTrack(mp4.MP4_CODEC_AAC)
 					hasAudio = true
 				}
-				if aframe, apts, adts, ok := nextAACFrame(); ok {
-					if err := muxer.Write(atid, aframe, uint64(apts), uint64(adts)); err != nil {
-						writeErr = err
+				if len(aacFrameBuf) > 0 {
+					if err := muxer.Write(atid, aacFrameBuf[0], uint64(apts), uint64(adts)); err != nil {
+						onFrameErr = err
 					}
 					prevADTS = adts
 				} else {
 					break
 				}
+				aacFrameBuf = aacFrameBuf[1:]
+				apts += aacMillisecondsPerFrame
+				adts += aacMillisecondsPerFrame
 			}
+
+			if !hasVideo {
+				vtid = muxer.AddVideoTrack(mp4.MP4_CODEC_H264)
+				hasVideo = true
+			}
+			if err := muxer.Write(vtid, vframe, uint64(vpts), uint64(vdts)); err != nil {
+				onFrameErr = err
+			}
+			//_vpts = vpts
+			//_vdts = vdts
 		}
 	}
 
-	if err := tsDemuxer.Input(tsInput); err != nil {
-		if errors.Is(err, io.ErrUnexpectedEOF) {
-			// File is incomplete, ignore
-		} else {
-			return fmt.Errorf("MPEG-TS demuxer: %w", err)
+	if len(aacInput) != len(tsInput) {
+		return fmt.Errorf("number of AAC segments (%v) and TS segments (%v) doesn't match", len(aacInput), len(tsInput))
+	}
+
+	var aHLSTime float64
+	for i := range tsInput {
+		adts = aHLSTime * 1000
+		apts = aHLSTime * 1000
+		{
+			data, err := os.ReadFile(aacInput[i].Filename)
+			if err != nil {
+				return err
+			}
+			codec.SplitAACFrame(data, func(frame []byte) {
+				aacFrameBuf = append(aacFrameBuf, frame)
+			})
+		}
+		{
+			data, err := os.ReadFile(tsInput[i].Filename)
+			if err != nil {
+				return err
+			}
+			if err := tsDemuxer.Input(bytes.NewReader(data)); err != nil {
+				if errors.Is(err, io.ErrUnexpectedEOF) {
+					// File is incomplete, ignore
+				} else {
+					return fmt.Errorf("MPEG-TS demuxer: %w", err)
+				}
+			}
 		}
+		onProgress(float64(i) / float64(len(tsInput)))
+		aHLSTime += aacInput[i].Duration
 	}
 
 	muxer.WriteTrailer()
 
-	if writeErr != nil {
-		// Currently only propagates last error
-		return fmt.Errorf("mp4 muxer: %w", err)
+	//fmt.Printf("**** %f %v %v %f %f\n", adts, _vpts, _vdts, float64(_vdts) - adts, float64(_vdts) / adts)
+
+	if onFrameErr != nil {
+		return fmt.Errorf("mp4 muxer: %w", onFrameErr)
 	}
 
 	return nil

diff --git a/pkg/southpark/download.go b/pkg/southpark/download.go
@@ -4,11 +4,8 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"io"
 	"os"
 	"path"
-
-	"github.com/xypwn/southpark-downloader-ui/pkg/ioutils"
 )
 
 type DownloaderStatus int
@@ -137,41 +134,28 @@ func (d *Downloader) Do() error {
 		}
 		defer outputFileMP4.Close()
 
-		segReader := func(startSeg, endSeg int, onError func(error)) io.Reader {
-			baseReader, w := io.Pipe()
-			r := ioutils.NewCtxReader(d.ctx, baseReader)
-
-			go func() {
-				for i := startSeg; i < endSeg; i++ {
-					data, err := os.ReadFile(getSegFileName(i))
-					if err != nil {
-						onError(err)
-						break
-					}
-					w.Write(data)
-					d.OnStatusChanged(DownloaderStatusPostprocessingVideo, float64(i)/float64(endSeg))
-				}
-				w.Close()
-			}()
-
-			return r
+		var tsSegs []SegmentFile
+		for i, seg := range stream.Video.Segments {
+			tsSegs = append(tsSegs, SegmentFile{
+				Filename: getSegFileName(i),
+				Duration: seg.Duration,
+			})
 		}
 
-		var videoConvertErr error
-		tsReader := segReader(0, len(stream.Video.Segments), func(err error) { videoConvertErr = err })
+		var aacSegs []SegmentFile
+		for i, seg := range stream.Audio.Segments {
+			aacSegs = append(aacSegs, SegmentFile{
+				Filename: getSegFileName(len(stream.Video.Segments)+i),
+				Duration: seg.Duration,
+			})
+		}
 
-		var audioConvertErr error
-		aacReader := segReader(len(stream.Video.Segments), len(stream.Video.Segments)+len(stream.Audio.Segments), func(err error) { audioConvertErr = err })
 
-		if err := ConvertTSAndAACToMP4(tsReader, aacReader, outputFileMP4); err != nil {
+		if err := ConvertTSAndAACToMP4(tsSegs, aacSegs, outputFileMP4, func(progress float64) {
+			d.OnStatusChanged(DownloaderStatusPostprocessingVideo, progress)
+		}); err != nil {
 			return fmt.Errorf("convert MPEG-TS and AAC to MP4: %w", err)
 		}
-		if videoConvertErr != nil {
-			return fmt.Errorf("convert MPEG-TS to MP4: %w", videoConvertErr)
-		}
-		if audioConvertErr != nil {
-			return fmt.Errorf("convert AAC to MP4: %w", audioConvertErr)
-		}
 	}
 
 	if d.outputSubtitlePath != "" {