Files
joy4/av/transcode/transcode.go
2021-01-04 16:03:02 +01:00

336 lines
8.2 KiB
Go

// Package transcoder implements Transcoder based on Muxer/Demuxer and AudioEncoder/AudioDecoder interface.
package transcode
import (
"fmt"
"time"
"github.com/kerberos-io/joy4/av"
"github.com/kerberos-io/joy4/av/pktque"
"github.com/kerberos-io/joy4/cgo/ffmpeg"
)
var Debug bool
type tStream struct {
codec av.CodecData
timeline *pktque.Timeline
aencodec, adecodec av.AudioCodecData
aenc av.AudioEncoder
adec av.AudioDecoder
vencodec, vdecodec av.VideoCodecData
venc *ffmpeg.VideoEncoder
vdec *ffmpeg.VideoDecoder
}
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,
)
FindVideoDecoderEncoder func(codec av.VideoCodecData, i int, percentage int) (
need bool, dec *ffmpeg.VideoDecoder, enc *ffmpeg.VideoEncoder, err error,
)
}
type Transcoder struct {
streams []*tStream
}
func NewTranscoder(streams []av.CodecData, options Options, percentage int) (_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
}
}
} else if stream.Type().IsVideo() {
if options.FindVideoDecoderEncoder != nil {
var ok bool
var enc *ffmpeg.VideoEncoder
var dec *ffmpeg.VideoDecoder
ok, dec, enc, err = options.FindVideoDecoderEncoder(stream.(av.VideoCodecData), i, percentage)
if ok {
if err != nil {
return
}
ts.timeline = &pktque.Timeline{}
if ts.codec, err = enc.CodecData(); err != nil {
return
}
ts.vencodec = ts.codec.(av.VideoCodecData)
ts.vdecodec = stream.(av.VideoCodecData)
ts.venc = enc
ts.vdec = 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 audio 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 audio 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
}
func (self *tStream) videoDecodeAndEncode(inpkt av.Packet) (outpkts []av.Packet, err error) {
var dur time.Duration
var frame *ffmpeg.VideoFrame
if frame, err = self.vdec.Decode(inpkt.Data); err != nil || frame == nil {
return
}
//if dur, err = self.vdecodec.PacketDuration(inpkt.Data); err != nil {
// err = fmt.Errorf("transcode: PacketDuration() failed for input video stream #%d", inpkt.Idx)
// return
//}
if Debug {
fmt.Println("transcode: push", inpkt.Time, dur)
}
self.timeline.Push(inpkt.Time, dur)
var _outpkts []av.Packet
if _outpkts, err = self.venc.Encode(frame); err != nil {
return
}
for _, _outpkt := range _outpkts {
if fpsNum, fpsDen := self.vencodec.Framerate(); fpsNum <= 0 || fpsDen <= 0 {
// FIXME this is a bit hacky
// Read codec data after encoding (because the sps and pps are not ready before the first keyframe is encoded)
var codecData av.VideoCodecData
codecData, err = self.venc.CodecData()
if err != nil {
return
}
self.vencodec = codecData.(av.VideoCodecData)
}
if dur, err = self.vencodec.PacketDuration(_outpkt.Data); err != nil {
err = fmt.Errorf("transcode: PacketDuration() failed for output video stream #%d", inpkt.Idx)
return
}
// TODO probably not needed now that _outpkt is an av.Packet
outpkt := av.Packet{IsKeyFrame: _outpkt.IsKeyFrame, Idx: inpkt.Idx, Data: _outpkt.Data}
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 if stream.venc != nil && stream.vdec != nil {
if out, err = stream.videoDecodeAndEncode(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
}
if stream.venc != nil {
stream.venc.Close()
stream.venc = nil
}
if stream.vdec != nil {
stream.vdec.Close()
stream.vdec = 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, 100); 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, 100); 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
}