vdk/av/transcode/transcode.go

// Package transcoder implements Transcoder based on Muxer/Demuxer and AudioEncoder/AudioDecoder interface.
package transcode

import (
	"fmt"
	"time"

	"git.r-2.top/kunmeng/vdk/av"
	"git.r-2.top/kunmeng/vdk/av/pktque"
)

var Debug bool

type tStream struct {
	codec              av.CodecData
	timeline           *pktque.Timeline
	aencodec, adecodec av.AudioCodecData
	aenc               av.AudioEncoder
	adec               av.AudioDecoder
}

type Options struct {
	// check if transcode is needed, and create the AudioDecoder and AudioEncoder.
	FindAudioDecoderEncoder func(codec av.AudioCodecData, i int) (
		need bool, dec av.AudioDecoder, enc av.AudioEncoder, err error,
	)
}

type Transcoder struct {
	streams []*tStream
}

func NewTranscoder(streams []av.CodecData, options Options) (_self *Transcoder, err error) {
	self := &Transcoder{}
	self.streams = []*tStream{}

	for i, stream := range streams {
		ts := &tStream{codec: stream}
		if stream.Type().IsAudio() {
			if options.FindAudioDecoderEncoder != nil {
				var ok bool
				var enc av.AudioEncoder
				var dec av.AudioDecoder
				ok, dec, enc, err = options.FindAudioDecoderEncoder(stream.(av.AudioCodecData), i)
				if ok {
					if err != nil {
						return
					}
					ts.timeline = &pktque.Timeline{}
					if ts.codec, err = enc.CodecData(); err != nil {
						return
					}
					ts.aencodec = ts.codec.(av.AudioCodecData)
					ts.adecodec = stream.(av.AudioCodecData)
					ts.aenc = enc
					ts.adec = dec
				}
			}
		}
		self.streams = append(self.streams, ts)
	}

	_self = self
	return
}

func (self *tStream) audioDecodeAndEncode(inpkt av.Packet) (outpkts []av.Packet, err error) {
	var dur time.Duration
	var frame av.AudioFrame
	var ok bool
	if ok, frame, err = self.adec.Decode(inpkt.Data); err != nil {
		return
	}
	if !ok {
		return
	}

	if dur, err = self.adecodec.PacketDuration(inpkt.Data); err != nil {
		err = fmt.Errorf("transcode: PacketDuration() failed for input stream #%d", inpkt.Idx)
		return
	}

	if Debug {
		fmt.Println("transcode: push", inpkt.Time, dur)
	}
	self.timeline.Push(inpkt.Time, dur)

	var _outpkts [][]byte
	if _outpkts, err = self.aenc.Encode(frame); err != nil {
		return
	}
	for _, _outpkt := range _outpkts {
		if dur, err = self.aencodec.PacketDuration(_outpkt); err != nil {
			err = fmt.Errorf("transcode: PacketDuration() failed for output stream #%d", inpkt.Idx)
			return
		}
		outpkt := av.Packet{Idx: inpkt.Idx, Data: _outpkt}
		outpkt.Time = self.timeline.Pop(dur)

		if Debug {
			fmt.Println("transcode: pop", outpkt.Time, dur)
		}

		outpkts = append(outpkts, outpkt)
	}

	return
}

// Do the transcode.
//
// In audio transcoding one Packet may transcode into many Packets
// packet time will be adjusted automatically.
func (self *Transcoder) Do(pkt av.Packet) (out []av.Packet, err error) {
	stream := self.streams[pkt.Idx]
	if stream.aenc != nil && stream.adec != nil {
		if out, err = stream.audioDecodeAndEncode(pkt); err != nil {
			return
		}
	} else {
		out = append(out, pkt)
	}
	return
}

// Get CodecDatas after transcoding.
func (self *Transcoder) Streams() (streams []av.CodecData, err error) {
	for _, stream := range self.streams {
		streams = append(streams, stream.codec)
	}
	return
}

// Close transcoder, close related encoder and decoders.
func (self *Transcoder) Close() (err error) {
	for _, stream := range self.streams {
		if stream.aenc != nil {
			stream.aenc.Close()
			stream.aenc = nil
		}
		if stream.adec != nil {
			stream.adec.Close()
			stream.adec = nil
		}
	}
	self.streams = nil
	return
}

// Wrap transcoder and origin Muxer into new Muxer.
// Write to new Muxer will do transcoding automatically.
type Muxer struct {
	av.Muxer   // origin Muxer
	Options    // transcode options
	transcoder *Transcoder
}

func (self *Muxer) WriteHeader(streams []av.CodecData) (err error) {
	if self.transcoder, err = NewTranscoder(streams, self.Options); err != nil {
		return
	}
	var newstreams []av.CodecData
	if newstreams, err = self.transcoder.Streams(); err != nil {
		return
	}
	if err = self.Muxer.WriteHeader(newstreams); err != nil {
		return
	}
	return
}

func (self *Muxer) WritePacket(pkt av.Packet) (err error) {
	var outpkts []av.Packet
	if outpkts, err = self.transcoder.Do(pkt); err != nil {
		return
	}
	for _, pkt := range outpkts {
		if err = self.Muxer.WritePacket(pkt); err != nil {
			return
		}
	}
	return
}

func (self *Muxer) Close() (err error) {
	if self.transcoder != nil {
		return self.transcoder.Close()
	}
	return
}

// Wrap transcoder and origin Demuxer into new Demuxer.
// Read this Demuxer will do transcoding automatically.
type Demuxer struct {
	av.Demuxer
	Options
	transcoder *Transcoder
	outpkts    []av.Packet
}

func (self *Demuxer) prepare() (err error) {
	if self.transcoder == nil {
		var streams []av.CodecData
		if streams, err = self.Demuxer.Streams(); err != nil {
			return
		}
		if self.transcoder, err = NewTranscoder(streams, self.Options); err != nil {
			return
		}
	}
	return
}

func (self *Demuxer) ReadPacket() (pkt av.Packet, err error) {
	if err = self.prepare(); err != nil {
		return
	}
	for {
		if len(self.outpkts) > 0 {
			pkt = self.outpkts[0]
			self.outpkts = self.outpkts[1:]
			return
		}
		var rpkt av.Packet
		if rpkt, err = self.Demuxer.ReadPacket(); err != nil {
			return
		}
		if self.outpkts, err = self.transcoder.Do(rpkt); err != nil {
			return
		}
	}
	return
}

func (self *Demuxer) Streams() (streams []av.CodecData, err error) {
	if err = self.prepare(); err != nil {
		return
	}
	return self.transcoder.Streams()
}

func (self *Demuxer) Close() (err error) {
	if self.transcoder != nil {
		return self.transcoder.Close()
	}
	return
}