migrate to github.com/asticode/go-astiav

2025-10-17 22:00:42 +08:00 · 2025-06-13 10:00:56 -04:00
parent 174196f916
commit 9dd255eb80
7 changed files with 266 additions and 125 deletions
--- a/internal/project/video_iterator.go
+++ b/internal/project/video_iterator.go
@@ -1,38 +1,48 @@
 package project

 import (
+	"errors"
 	"fmt"
 	"image"
-	"io"
 	"log"
 	"time"

-	"github.com/nareix/joy4/av"
-	"github.com/nareix/joy4/av/avutil"
-	"github.com/nareix/joy4/cgo/ffmpeg"
-	"github.com/nareix/joy4/format"
-	"github.com/nareix/joy4/format/mp4"
+	"github.com/asticode/go-astiav"
 )

 func init() {
-	format.RegisterAll()
+	astiav.SetLogLevel(astiav.LogLevelError)
 }

 type Iterator struct {
-	err      error
-	demuxer  av.DemuxCloser
-	decoders []*ffmpeg.VideoDecoder
-	rect     image.Rectangle
-	packet   av.Packet
-	frame    int
-	decoded  bool
-	vf       *ffmpeg.VideoFrame
+	err            error
+	formatContext  *astiav.FormatContext
+	codecContext   *astiav.CodecContext
+	videoStreamIdx int
+	rect           image.Rectangle
+	packet         *astiav.Packet
+	frame          int
+	decoded        bool
+	currentFrame   *astiav.Frame
+	img            image.Image
 }

 func (p *Iterator) Close() {
-	if p.demuxer != nil {
-		p.demuxer.Close()
-		p.demuxer = nil
+	if p.formatContext != nil {
+		p.formatContext.Free()
+		p.formatContext = nil
+	}
+	if p.codecContext != nil {
+		p.codecContext.Free()
+		p.codecContext = nil
+	}
+	if p.currentFrame != nil {
+		p.currentFrame.Free()
+		p.currentFrame = nil
+	}
+	if p.packet != nil {
+		p.packet.Free()
+		p.packet = nil
 	}
 }

@@ -40,89 +50,180 @@ func NewIterator(filename string) (iter *Iterator, err error) {
 	if filename == "" {
 		panic("missing filename")
 	}
+
 	iter = &Iterator{frame: -1}
-	iter.demuxer, err = avutil.Open(filename)
-	if err != nil {
-		return nil, err
+
+	// Open input file
+	formatContext := astiav.AllocFormatContext()
+	if formatContext == nil {
+		return nil, fmt.Errorf("allocating format context failed")
+	}
+	iter.formatContext = formatContext
+
+	// Get stream info
+	if err := formatContext.OpenInput(filename, nil, nil); err != nil {
+		iter.Close()
+		return nil, fmt.Errorf("opening input: %w", err)
 	}

-	streams, err := iter.demuxer.Streams()
-	if err != nil {
-		iter.Close()
-		return nil, err
-	}
-	iter.decoders = make([]*ffmpeg.VideoDecoder, len(streams))
-	for i, stream := range streams {
-		// if stream.Type().IsAudio() {
-		// astream := stream.(av.AudioCodecData)
-		// fmt.Println(astream.Type(), astream.SampleRate(), astream.SampleFormat(), astream.ChannelLayout())
-		// } else if stream.Type().IsVideo() {
-		fmt.Printf("stream[%d] = %s (video:%v)\n", i, stream.Type(), stream.Type().IsVideo())
-		if stream.Type().IsVideo() {
-			vstream := stream.(av.VideoCodecData)
-			r := image.Rect(0, 0, vstream.Width(), vstream.Height())
-			if iter.rect.Empty() {
-				iter.rect = r
-			} else if !iter.rect.Eq(r) {
-				return nil, fmt.Errorf("video stream %d(%v) doesn't match expected %v", i, r, iter.rect)
-			}
-			// fmt.Printf("stream[%d] = %s\n", i, vstream.Type())
-			// fmt.Printf("stream[%d] %#v\n", i, vstream)
-			iter.decoders[i], err = ffmpeg.NewVideoDecoder(vstream)
-			if err != nil {
-				log.Fatalf("NewVideoDecoder error: %s", err)
-			}
+	// Find the first video stream
+	iter.videoStreamIdx = -1
+	for i, stream := range formatContext.Streams() {
+		if stream.CodecParameters().MediaType() != astiav.MediaTypeVideo {
+			continue
 		}
+		iter.videoStreamIdx = i
+		width := stream.CodecParameters().Width()
+		height := stream.CodecParameters().Height()
+		iter.rect = image.Rect(0, 0, width, height)
+		fmt.Printf("stream[%d] = video (%dx%d)\n", i, width, height)
+
+		iter.currentFrame = astiav.AllocFrame()
+		// iter.currentFrame.SetWidth(width)
+		// iter.currentFrame.SetHeight(height)
+		// iter.currentFrame.SetPixelFormat(astiav.PixelFormatYuv420P)
+		break
 	}
-	if iter.rect.Empty() {
+
+	if iter.videoStreamIdx == -1 {
+		iter.Close()
 		return nil, fmt.Errorf("no video stream found")
 	}
+
+	// Find decoder
+	stream := formatContext.Streams()[iter.videoStreamIdx]
+	codec := astiav.FindDecoder(stream.CodecParameters().CodecID())
+	if codec == nil {
+		iter.Close()
+		return nil, fmt.Errorf("no codec found for stream %d", iter.videoStreamIdx)
+	}
+	log.Printf("codec found: %s", codec.Name())
+
+	// Create codec context
+	iter.codecContext = astiav.AllocCodecContext(codec)
+	if iter.codecContext == nil {
+		iter.Close()
+		return nil, fmt.Errorf("failed to create codec context")
+	}
+	log.Printf("codec parameters for stream %d: %#v", iter.videoStreamIdx, stream.CodecParameters())
+
+	// Copy codec parameters
+	if err := iter.codecContext.FromCodecParameters(stream.CodecParameters()); err != nil {
+		iter.Close()
+		return nil, fmt.Errorf("copying codec parameters: %w", err)
+	}
+	// iter.codecContext.SetPixelFormat(astiav.PixelFormatYuv420P)
+
+	// Open codec
+	if err := iter.codecContext.Open(codec, nil); err != nil {
+		iter.Close()
+		return nil, fmt.Errorf("opening codec: %w", err)
+	}
+
+	// Allocate packet
+	iter.packet = astiav.AllocPacket()
+
+	log.Printf("codecContext: %s pixelFormat=%#v", iter.codecContext.String(), iter.codecContext.PixelFormat())
+
 	return iter, nil
 }

 func (i *Iterator) Seek(d time.Duration) error {
-	dm := i.demuxer.(*avutil.HandlerDemuxer).Demuxer.(*mp4.Demuxer)
-	log.Printf("should seek %s", d)
-	return dm.SeekToTime(d)
+	// Convert time.Duration to AV timestamp
+	ts := int64(d / time.Microsecond)
+	timeBase := i.formatContext.Streams()[i.videoStreamIdx].TimeBase()
+	timestamp := astiav.RescaleQ(ts, astiav.NewRational(1, 1000000), timeBase)
+
+	log.Printf("should seek %s (timestamp: %d)", d, timestamp)
+
+	// Seek to timestamp
+	seekFlags := astiav.NewSeekFlags(astiav.SeekFlagBackward, astiav.SeekFlagAny)
+	if err := i.formatContext.SeekFrame(i.videoStreamIdx, timestamp, seekFlags); err != nil {
+		return fmt.Errorf("seeking: %w", err)
+	}
+
+	// Reset frame counter
+	i.frame = -1
+
+	// Send empty packet to flush the decoder
+	if err := i.codecContext.SendPacket(nil); err != nil {
+		return fmt.Errorf("flushing decoder: %w", err)
+	}
+
+	return nil
 }

 func (i *Iterator) VideoResolution() string {
 	return fmt.Sprintf("%dx%d", i.rect.Dx(), i.rect.Dy())
 }

-func (i *Iterator) NextWithImage() bool {
-	for i.Next() {
-		i.err = i.DecodeFrame()
-		if i.err != nil {
-			return false
-		}
-		if i.vf == nil {
-			continue
-		}
-		return true
-	}
-	return false
+// func (i *Iterator) NextWithImage() bool {
+// 	for i.Next() {
+// 		i.err = i.DecodeFrame()
+// 		if i.err != nil {
+// 			return false
+// 		}
+// 		if i.currentFrame == nil {
+// 			continue
+// 		}
+// 		return true
+// 	}
+// 	return false
+// }

-}
 func (i *Iterator) Next() bool {
-	var err error
-	var pkt av.Packet
+	i.decoded = false
+
 	for {
-		i.vf = nil
-		i.decoded = false
-		if pkt, err = i.demuxer.ReadPacket(); err != nil {
-			if err == io.EOF {
+		// Read frame
+		err := i.formatContext.ReadFrame(i.packet)
+		if err != nil {
+			if errors.Is(err, astiav.ErrEof) {
 				return false
 			}
-			i.err = err
+			i.err = fmt.Errorf("reading frame: %w", err)
 			return false
 		}
-		// skip packets we don't have a decoder for
-		if i.decoders[pkt.Idx] == nil {
+
+		// Skip if not a video packet
+		if i.packet.StreamIndex() != i.videoStreamIdx {
+			log.Printf("Skipping packet from stream %d (not video)", i.packet.StreamIndex())
+			i.packet.Unref()
 			continue
 		}
-		i.packet = pkt
 		i.frame++
+
+		// Send packet to decoder
+		if i.err = i.codecContext.SendPacket(i.packet); i.err != nil {
+			log.Printf("SendPacket() error: %v", i.err)
+			i.packet.Unref()
+			return false
+		}
+
+		// Get frame from decoder
+		i.err = i.codecContext.ReceiveFrame(i.currentFrame)
+		if errors.Is(i.err, astiav.ErrEagain) {
+			i.packet.Unref()
+			log.Printf("ReceiveFrame() Eagain for frame %d, waiting for more data", i.frame)
+			i.err = nil // No frame available yet, continue to read more packets
+			continue
+		}
+		defer i.packet.Unref()
+		if i.err != nil {
+			// Check if we need more data or reached EOF
+			log.Printf("ReceiveFrame() error: %v", i.err)
+			return false
+		}
+
+		defer i.currentFrame.Unref()
+
+		img, _ := i.currentFrame.Data().GuessImageFormat()
+		i.err = i.currentFrame.Data().ToImage(img)
+		if i.err != nil {
+			log.Printf("ToImage() error: %v", i.err)
+			return false
+		}
+		i.img = img
 		return true
 	}
 }
@@ -131,41 +232,28 @@ func (i *Iterator) Frame() int {
 	return i.frame
 }

-func (i *Iterator) DecodeFrame() error {
-	if i.decoded {
-		return i.err
-	}
-	// decode
-	decoder := i.decoders[i.packet.Idx]
-	var err error
-	if len(i.packet.Data) == 0 {
-		log.Printf("no packet at frame %d", i.frame)
-		return nil
-	}
-	i.vf, err = decoder.Decode(i.packet.Data)
-	if i.vf == nil {
-		log.Printf("no image at frame %d", i.frame)
-		i.frame--
-	}
-	i.decoded = true
-	return err
+func (i *Iterator) Image() image.Image {
+	return i.img
 }

-func (i *Iterator) Image() *image.YCbCr {
-	if i.frame == -1 {
-		if !i.NextWithImage() {
-			panic("no image")
-		}
-	}
-	i.err = i.DecodeFrame()
-	if i.vf == nil {
-		return nil
-	}
-	return &i.vf.Image
+func (i *Iterator) Error() error {
+	return i.err
 }
-func (i *Iterator) Error() error            { return i.err }
-func (i *Iterator) Duration() time.Duration { return i.packet.Time }
+
+func (i *Iterator) Duration() time.Duration {
+	if i.packet == nil || i.packet.Pts() == astiav.NoPtsValue {
+		return 0
+	}
+
+	timeBase := i.formatContext.Streams()[i.videoStreamIdx].TimeBase()
+	durationMicros := astiav.RescaleQ(i.packet.Pts(), timeBase, astiav.NewRational(1, 1000000))
+	return time.Duration(durationMicros) * time.Microsecond
+}
+
 func (i *Iterator) DurationMs() time.Duration {
-	return (i.packet.Time / time.Millisecond) * time.Millisecond
+	return i.Duration().Round(time.Millisecond)
+}
+
+func (i *Iterator) IsKeyFrame() bool {
+	return i.currentFrame.KeyFrame()
 }
-func (i *Iterator) IsKeyFrame() bool { return i.packet.IsKeyFrame }