diff --git a/encoder.go b/encoder.go new file mode 100644 index 0000000..699a031 --- /dev/null +++ b/encoder.go @@ -0,0 +1,98 @@ +package main + +import ( + "errors" + "fmt" + "log" + + "github.com/asticode/go-astiav" + "github.com/asticode/go-astikit" +) + +type EncoderStream struct { + encCodec *astiav.Codec + encCodecContext *astiav.CodecContext + encPkt *astiav.Packet +} + +type Encoder struct { + Stream *EncoderStream + Closer *astikit.Closer +} + +type EncoderParams struct { + Height int + Width int + TimeBase astiav.Rational + AspectRatio astiav.Rational + PixelFormat astiav.PixelFormat +} + +func NewEncoder(params EncoderParams) (ts *Encoder, err error) { + ts = &Encoder{} + s := &EncoderStream{} + c := astikit.NewCloser() + ts.Closer = c + ts.Stream = s + + // Get codec id + codecID := astiav.CodecIDMpeg4 + + // Find encoder + if s.encCodec = astiav.FindEncoder(codecID); s.encCodec == nil { + err = errors.New("main: codec is nil") + return + } + + // Alloc codec context + if s.encCodecContext = astiav.AllocCodecContext(s.encCodec); s.encCodecContext == nil { + err = errors.New("main: codec context is nil") + return + } + c.Add(s.encCodecContext.Free) + + s.encCodecContext.SetHeight(params.Height) + s.encCodecContext.SetPixelFormat(params.PixelFormat) + s.encCodecContext.SetBitRate(1000) + s.encCodecContext.SetSampleAspectRatio(params.AspectRatio) + s.encCodecContext.SetTimeBase(params.TimeBase) + s.encCodecContext.SetWidth(params.Width) + + // Open codec context + if err = s.encCodecContext.Open(s.encCodec, nil); err != nil { + err = fmt.Errorf("encoder: opening codec context failed: %w", err) + return + } + + // Allocate packet + s.encPkt = astiav.AllocPacket() + c.Add(s.encPkt.Free) + + return ts, nil +} + +func (ts *Encoder) Close() { + ts.Closer.Close() +} + +func (ts *Encoder) Encode(f *astiav.Frame) (pkt *astiav.Packet, err error) { + s := ts.Stream + + // Unref packet + s.encPkt.Unref() + + log.Println(len(f.Data())) + + // Send frame + if err = s.encCodecContext.SendFrame(f); err != nil { + err = fmt.Errorf("encoder: sending frame failed: %w", err) + return nil, err + } + + if err := s.encCodecContext.ReceivePacket(s.encPkt); err != nil { + log.Fatal(fmt.Errorf("encoder: receiving frame failed: %w", err)) + return nil, err + } + + return s.encPkt, nil +} diff --git a/go-astiav b/go-astiav index 0f5fd94..f5bf3b5 160000 --- a/go-astiav +++ b/go-astiav @@ -1 +1 @@ -Subproject commit 0f5fd94d5d18518593047bdab84f68f96fe3126d +Subproject commit f5bf3b5e98fa7abe6aeb311eca046b5892901762 diff --git a/main.go b/main.go index e847674..317dc9f 100644 --- a/main.go +++ b/main.go @@ -58,29 +58,52 @@ func ReplaceTrack(prefix string, peerConnection *webrtc.PeerConnection) { log.Println(fmt.Errorf("main: peer connection closed due to error: %w", err)) } } else if connectionState.String() == "connected" { - astiav.SetLogLevel(astiav.LogLevelDebug) + astiav.SetLogLevel(astiav.LogLevelError) astiav.SetLogCallback(func(l astiav.LogLevel, fmt, msg, parent string) { log.Printf("ffmpeg log: %s (level: %d)\n", strings.TrimSpace(msg), l) }) go func() { + encoderParams := EncoderParams{ + Width: visionTrack.Width(), + Height: visionTrack.Height(), + TimeBase: visionTrack.TimeBase(), + AspectRatio: visionTrack.AspectRatio(), + PixelFormat: visionTrack.PixelFormat(), + } + + encoder, err := NewEncoder(encoderParams) + + if err != nil { + log.Fatal(fmt.Errorf("main: creating track failed: %w", err)) + return + } + defer encoder.Close() + go visionTrack.Start() defer visionTrack.Stop() + for visionTrack != nil { + var rtpPacket []byte for frame := range visionTrack.Frame { // Do something with decoded frame fmt.Println(frame.Roll) - // so right now frame is a raw decoded AVFrame // https://github.com/FFmpeg/FFmpeg/blob/n5.0/libavutil/frame.h#L317 - // avframe := frame.Frame + avframe := frame.Frame // we need to: // 1. transcode to h264 with adaptive bitrate using astiav + outFrame, err := encoder.Encode(avframe) + if err != nil { + log.Println(fmt.Errorf("encode error: %w", err)) + + continue + } // 2. create an RTP packet with h264 data inside using pion's rtp packetizer - var rtpPacket []byte - // 3. write the RTP packet to the videoTrack + fmt.Println(len(outFrame.Data())) + // 3. write the RTP packet to the videoTrack if _, err = videoTrack.Write(rtpPacket); err != nil { if errors.Is(err, io.ErrClosedPipe) { // The peerConnection has been closed. diff --git a/run.sh b/run.sh index c80f4c9..b90fd4d 100755 --- a/run.sh +++ b/run.sh @@ -1,4 +1,4 @@ -export CGO_LDFLAGS="-L$PWD/go-astiav/tmp/n4.4.1/lib/" -export CGO_CXXFLAGS="-I$PWD/go-astiav/tmp/n4.4.1/include/" +export CGO_LDFLAGS="-L$PWD/go-astiav/tmp/n4.4.1/lib" +export CGO_CXXFLAGS="-I$PWD/go-astiav/tmp/n4.4.1/include" export PKG_CONFIG_PATH="$PWD/go-astiav/tmp/n4.4.1/lib/pkgconfig" go run . \ No newline at end of file diff --git a/vision_ipc_track.go b/vision_ipc_track.go index a690507..ab6133a 100644 --- a/vision_ipc_track.go +++ b/vision_ipc_track.go @@ -14,7 +14,7 @@ import ( const V4L2_BUF_FLAG_KEYFRAME = uint32(8) -type Stream struct { +type VisionIpcTrackDecoderStream struct { decCodec *astiav.Codec decCodecContext *astiav.CodecContext inputStream *astiav.Stream @@ -22,7 +22,7 @@ type Stream struct { type VisionIpcTrack struct { name string - stream *Stream + stream *VisionIpcTrackDecoderStream context *zmq.Context subscriber *zmq.Socket lastIdx int64 @@ -64,7 +64,7 @@ func NewVisionIpcTrack(name string) (track *VisionIpcTrack, err error) { subscriber.Connect(GetServiceURI(name)) // Create stream - s := &Stream{inputStream: nil} + s := &VisionIpcTrackDecoderStream{inputStream: nil} // Find decoder if s.decCodec = astiav.FindDecoder(astiav.CodecIDHevc); s.decCodec == nil { @@ -76,6 +76,12 @@ func NewVisionIpcTrack(name string) (track *VisionIpcTrack, err error) { return nil, errors.New("main: codec context is nil") } + s.decCodecContext.SetHeight(track.Height()) + s.decCodecContext.SetPixelFormat(track.PixelFormat()) + s.decCodecContext.SetSampleAspectRatio(track.AspectRatio()) + s.decCodecContext.SetTimeBase(track.TimeBase()) + s.decCodecContext.SetWidth(track.Width()) + // Open codec context if err := s.decCodecContext.Open(s.decCodec, nil); err != nil { log.Fatal(fmt.Errorf("main: opening codec context failed: %w", err)) @@ -107,6 +113,30 @@ func NewVisionIpcTrack(name string) (track *VisionIpcTrack, err error) { }, nil } +func (v *VisionIpcTrack) TimeBase() (t astiav.Rational) { + return astiav.NewRational(v.FrameRate(), 1) +} + +func (v *VisionIpcTrack) FrameRate() int { + return 20 +} + +func (v *VisionIpcTrack) Height() int { + return 1208 +} + +func (v *VisionIpcTrack) Width() int { + return 1928 +} + +func (v *VisionIpcTrack) AspectRatio() (t astiav.Rational) { + return astiav.NewRational(151, 241) +} + +func (v *VisionIpcTrack) PixelFormat() (f astiav.PixelFormat) { + return astiav.PixelFormatYuv420P +} + func (v *VisionIpcTrack) Stop() { if !Open { return