161 lines
3.1 KiB
Go
161 lines
3.1 KiB
Go
package processor
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"github.com/bluenviron/mediacommon/pkg/codecs/mpeg4audio"
|
|
"log"
|
|
"os"
|
|
"sync"
|
|
|
|
"github.com/bluenviron/gortsplib/v4/pkg/format"
|
|
"github.com/bluenviron/mediacommon/pkg/codecs/h264"
|
|
"github.com/bluenviron/mediacommon/pkg/formats/mpegts"
|
|
)
|
|
|
|
var (
|
|
i = 0
|
|
f = 0
|
|
j = 0
|
|
)
|
|
|
|
func multiplyAndDivide(v, m, d int64) int64 {
|
|
secs := v / d
|
|
dec := v % d
|
|
return (secs*m + dec*m/d)
|
|
}
|
|
|
|
// MpegTSMuxer allows to save a H264 / MPEG-4 audio stream into MPEG-TS file.
|
|
type MpegTSMuxer struct {
|
|
FileName string
|
|
H264Format *format.H264
|
|
Mpeg4AudioFormat *format.MPEG4Audio
|
|
|
|
idrFrame [][]byte
|
|
idrChecker bool
|
|
|
|
f *os.File
|
|
b *bufio.Writer
|
|
w *mpegts.Writer
|
|
h264Track *mpegts.Track
|
|
mpeg4AudioTrack *mpegts.Track
|
|
dtsExtractor *h264.DTSExtractor2
|
|
mutex sync.Mutex
|
|
}
|
|
|
|
// Initialize initializes a MpegtsMuxer.
|
|
func (e *MpegTSMuxer) Initialize() error {
|
|
var err error
|
|
e.f, err = os.Create(e.FileName)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
e.b = bufio.NewWriter(e.f)
|
|
|
|
e.h264Track = &mpegts.Track{
|
|
Codec: &mpegts.CodecH264{},
|
|
}
|
|
|
|
e.mpeg4AudioTrack = &mpegts.Track{
|
|
Codec: &mpegts.CodecMPEG4Audio{
|
|
Config: mpeg4audio.Config{},
|
|
},
|
|
}
|
|
|
|
e.w = mpegts.NewWriter(e.b, []*mpegts.Track{e.h264Track, e.mpeg4AudioTrack})
|
|
|
|
e.idrChecker = false
|
|
|
|
return nil
|
|
}
|
|
|
|
// Close closes all the MpegtsMuxer resources.
|
|
func (e *MpegTSMuxer) Close() {
|
|
e.b.Flush()
|
|
e.f.Close()
|
|
}
|
|
|
|
// WriteH264 writes a H264 access unit into MPEG-TS.
|
|
func (e *MpegTSMuxer) WriteH264(pts int64, au [][]byte) error {
|
|
e.mutex.Lock()
|
|
defer e.mutex.Unlock()
|
|
|
|
var filteredAU [][]byte
|
|
|
|
nonIDRPresent := false
|
|
idrPresent := false
|
|
|
|
for _, nalu := range au {
|
|
i++
|
|
typ := h264.NALUType(nalu[0] & 0x1F)
|
|
f++
|
|
switch typ {
|
|
case h264.NALUTypeSPS:
|
|
e.H264Format.SPS = nalu
|
|
continue
|
|
|
|
case h264.NALUTypePPS:
|
|
e.H264Format.PPS = nalu
|
|
continue
|
|
|
|
case h264.NALUTypeAccessUnitDelimiter:
|
|
continue
|
|
|
|
case h264.NALUTypeIDR:
|
|
idrPresent = true
|
|
j++
|
|
e.idrFrame = append([][]byte{e.H264Format.SPS, e.H264Format.PPS}, nalu)
|
|
|
|
case h264.NALUTypeNonIDR:
|
|
nonIDRPresent = true
|
|
}
|
|
|
|
filteredAU = append(filteredAU, nalu)
|
|
}
|
|
|
|
fmt.Println(i, f, j)
|
|
|
|
au = filteredAU
|
|
|
|
if !e.idrChecker {
|
|
au = append(e.idrFrame, au...)
|
|
e.idrChecker = true
|
|
//mux.frameChecker = true
|
|
}
|
|
|
|
if au == nil || ((!nonIDRPresent && !idrPresent) && e.idrFrame == nil) {
|
|
return nil
|
|
}
|
|
|
|
// Add SPS and PPS before access unit that contains an IDR.
|
|
//if idrPresent {
|
|
// au = append([][]byte{e.H264Format.SPS, e.H264Format.PPS}, au...)
|
|
//}
|
|
|
|
if e.dtsExtractor == nil {
|
|
// Skip samples silently until we find one with an IDR.
|
|
if !idrPresent {
|
|
return nil
|
|
}
|
|
e.dtsExtractor = h264.NewDTSExtractor2()
|
|
}
|
|
|
|
log.Printf("pts: %d", pts)
|
|
|
|
dts, err := e.dtsExtractor.Extract(au, pts)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Encode into MPEG-TS.
|
|
return e.w.WriteH2642(e.h264Track, pts, dts, au)
|
|
}
|
|
|
|
// WriteAAC writes MPEG-4 audio access units into MPEG-TS.
|
|
func (e *MpegTSMuxer) WriteAAC(aus [][]byte, pts int64) error {
|
|
e.mutex.Lock()
|
|
defer e.mutex.Unlock()
|
|
|
|
return e.w.WriteMPEG4Audio(e.mpeg4AudioTrack, multiplyAndDivide(pts, 90000, int64(8000)), aus)
|
|
}
|