commit 4ce46594fd4534653d3fe5eab5200e19fcdb2fc1 Author: Daniel Sullivan Date: Tue Aug 13 16:54:29 2024 +0900 Initial diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1933786 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/test/ diff --git a/ffmpeg/consts.go b/ffmpeg/consts.go new file mode 100644 index 0000000..c2040a3 --- /dev/null +++ b/ffmpeg/consts.go @@ -0,0 +1,3 @@ +package ffmpeg + +const AV_DICT_IGNORE_SUFFIX = 2 diff --git a/ffmpeg/ffmpeg.go b/ffmpeg/ffmpeg.go new file mode 100644 index 0000000..8667aa9 --- /dev/null +++ b/ffmpeg/ffmpeg.go @@ -0,0 +1,100 @@ +package ffmpeg + +import ( + "fmt" + "github.com/asticode/go-astiav" + "sync" + "time" +) + +type VideoStream struct { + Width int + Height int + Framerate float64 + Codec string +} + +type AudioStream struct { + Channels int + SampleRate int + Codec string +} + +type File struct { + inputContext *astiav.FormatContext + Metadata map[string]string + Duration time.Duration + videoStreams []VideoStream + audioStreams []AudioStream + accessLock sync.Mutex +} + +func Open(path string) (*File, error) { + inputContext := astiav.AllocFormatContext() + if err := inputContext.OpenInput(path, nil, nil); err != nil { + return nil, fmt.Errorf("could not open input: %w", err) + } + + f := &File{ + inputContext: inputContext, + Metadata: make(map[string]string), + } + + // Retrieve the Metadata + metadata := inputContext.Metadata() + + if metadata != nil { + var previousMetadata *astiav.DictionaryEntry + for { + previousMetadata = metadata.Get("", previousMetadata, AV_DICT_IGNORE_SUFFIX) + if previousMetadata == nil { + break + } + f.Metadata[previousMetadata.Key()] = previousMetadata.Value() + fmt.Printf("%s: %s\n", previousMetadata.Key(), previousMetadata.Value()) + } + } + + // Retrieve the Duration + duration := inputContext.Duration() + f.Duration = time.Duration(duration) * time.Microsecond + + // Iterate over the streams to get codec information + for i := 0; i < inputContext.NbStreams(); i++ { + stream := inputContext.Streams()[i] + codecParams := stream.CodecParameters() + + switch codecParams.CodecType() { + case astiav.MediaTypeVideo: + rational := stream.AvgFrameRate() + vs := VideoStream{ + Width: codecParams.Width(), + Height: codecParams.Height(), + Framerate: float64(rational.Num()) / float64(rational.Den()), + Codec: astiav.FindDecoder(codecParams.CodecID()).Name(), + } + f.videoStreams = append(f.videoStreams, vs) + case astiav.MediaTypeAudio: + as := AudioStream{ + Channels: codecParams.Channels(), + SampleRate: codecParams.SampleRate(), + Codec: astiav.FindDecoder(codecParams.CodecID()).Name(), + } + f.audioStreams = append(f.audioStreams, as) + } + } + + return f, nil +} + +func (f *File) GetKeyframes() int64 { + f.accessLock.Lock() + defer f.accessLock.Unlock() + + var keyframeIndexes []int64 + +} + +func (f *File) Close() { + f.inputContext.CloseInput() +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..58ea913 --- /dev/null +++ b/go.mod @@ -0,0 +1,12 @@ +module media_manager + +go 1.22 + +replace ( + github.com/asticode/go-astiav => ../go-astiav +) + +require ( + github.com/asticode/go-astiav v0.13.2-0.20240505170917-b205dafea242 // indirect + github.com/asticode/go-astikit v0.42.0 // indirect +) diff --git a/info.go b/info.go new file mode 100644 index 0000000..7f582cf --- /dev/null +++ b/info.go @@ -0,0 +1,95 @@ +package main + +import ( + "errors" + "fmt" + "log" + "media_manager/ffmpeg" + "time" + + "github.com/asticode/go-astiav" +) + +func info() { + // Initialize the library + astiav.SetLogLevel(astiav.LogLevelInfo) + + // Open the media file + filePath := "test/test.mkv" + formatCtx := astiav.AllocFormatContext() + if err := formatCtx.OpenInput(filePath, nil, nil); err != nil { + log.Fatalf("could not open input: %v", err) + } + defer formatCtx.CloseInput() + + // Retrieve the metadata + metadata := formatCtx.Metadata() + + if metadata != nil { + var previousMetadata *astiav.DictionaryEntry + for { + previousMetadata = metadata.Get("", previousMetadata, ffmpeg.AV_DICT_IGNORE_SUFFIX) + if previousMetadata == nil { + break + } + fmt.Printf("%s: %s\n", previousMetadata.Key(), previousMetadata.Value()) + } + } else { + fmt.Println("No metadata found") + } + + // Retrieve the duration + duration := formatCtx.Duration() + fmt.Printf("Duration: %v\n", time.Duration(duration)*time.Microsecond) + + // Iterate over the streams to get codec information + for i := 0; i < formatCtx.NbStreams(); i++ { + stream := formatCtx.Streams()[i] + codecParams := stream.CodecParameters() + codec := astiav.FindDecoder(codecParams.CodecID()) + fmt.Printf("Stream %d: Codec: %s\n", i, codec.Name()) + + // Get resolution information + if codecParams.CodecType() == astiav.MediaTypeVideo { + rational := stream.AvgFrameRate() + framerate := float64(rational.Num()) / float64(rational.Den()) + fmt.Printf("Stream %d: Framerate: %.2f fps\n", i, framerate) + width := codecParams.Width() + height := codecParams.Height() + fmt.Printf("Stream %d: Resolution: %dx%d\n", i, width, height) + } + + // Get codec details + fmt.Printf("Stream %d: Codec Level: %d\n", i, codecParams.Level()) + fmt.Printf("Stream %d: Codec Profile: %d\n", i, codecParams.Profile()) + } + + var keyframeTimestamps []int64 + + // Alloc packet + pkt := astiav.AllocPacket() + defer pkt.Free() + + // Loop through packets + for { + // Read frame + if err := formatCtx.ReadFrame(pkt); err != nil { + if errors.Is(err, astiav.ErrEof) { + break + } + log.Fatal(fmt.Errorf("main: reading frame failed: %w", err)) + } + + // Check if the packet is a keyframe + if pkt.Flags().Has(astiav.PacketFlagKey) { + keyframeTimestamps = append(keyframeTimestamps, pkt.Pts()) + } + + pkt.Unref() + } + + // Print keyframe timestamps + for _, ts := range keyframeTimestamps { + fmt.Printf("Keyframe at PTS: %d\n", ts) + } +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..4b15acf --- /dev/null +++ b/main.go @@ -0,0 +1,180 @@ +package main + +import "C" +import ( + "errors" + "fmt" + "github.com/asticode/go-astiav" + "log" + "os" + "strings" +) + +var output = "test/test.mp4" + +func main() { + // Handle ffmpeg logs + astiav.SetLogLevel(astiav.LogLevelDebug) + astiav.SetLogCallback(func(c astiav.Classer, l astiav.LogLevel, fmt, msg string) { + var cs string + log.Printf("ffmpeg log: %s%s - level: %d\n", strings.TrimSpace(msg), cs, l) + }) + + info() + + // Alloc packet + pkt := astiav.AllocPacket() + defer pkt.Free() + + // Alloc input format context + inputFormatContext := astiav.AllocFormatContext() + if inputFormatContext == nil { + log.Fatal(errors.New("main: input format context is nil")) + } + defer inputFormatContext.Free() + + // Open input + if err := inputFormatContext.OpenInput("test/test.mkv", nil, nil); err != nil { + log.Fatal(fmt.Errorf("main: opening input failed: %w", err)) + } + defer inputFormatContext.CloseInput() + + // Find stream info + if err := inputFormatContext.FindStreamInfo(nil); err != nil { + log.Fatal(fmt.Errorf("main: finding stream info failed: %w", err)) + } + + outputFormat := astiav.FindOutputFormat("mp4") + + if outputFormat == nil { + log.Fatal(errors.New("main: could not find output format")) + } + + // Alloc output format context + outputFormatContext, err := astiav.AllocOutputFormatContext(outputFormat, "", "") + if err != nil { + log.Fatal(fmt.Errorf("main: allocating output format context failed: %w", err)) + } + if outputFormatContext == nil { + log.Fatal(errors.New("main: output format context is nil")) + } + defer outputFormatContext.Free() + + // Loop through streams + inputStreams := make(map[int]*astiav.Stream) // Indexed by input stream index + outputStreams := make(map[int]*astiav.Stream) // Indexed by input stream index + for _, is := range inputFormatContext.Streams() { + // Only process audio or video + if is.CodecParameters().MediaType() != astiav.MediaTypeAudio && + is.CodecParameters().MediaType() != astiav.MediaTypeVideo { + continue + } + + // Add input stream + inputStreams[is.Index()] = is + + // Add stream to output format context + os := outputFormatContext.NewStream(nil) + if os == nil { + log.Fatal(errors.New("main: output stream is nil")) + } + + // Copy codec parameters + if err = is.CodecParameters().Copy(os.CodecParameters()); err != nil { + log.Fatal(fmt.Errorf("main: copying codec parameters failed: %w", err)) + } + + // Reset codec tag + os.CodecParameters().SetCodecTag(0) + + // Add output stream + outputStreams[is.Index()] = os + } + + f, err := os.OpenFile("test/test.mp4", os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + log.Fatal(fmt.Errorf("main: opening output failed: %w", err)) + } + defer f.Close() + + // If this is a file, we need to use an io context + if !outputFormatContext.OutputFormat().Flags().Has(astiav.IOFormatFlagNofile) { + ioContext, err := astiav.AllocIOContext(4092, true, f.Read, f.Seek, f.Write) + if err != nil { + log.Fatal(fmt.Errorf("main: opening io context failed: %w", err)) + } + defer ioContext.Free() //nolint:errcheck + + // Update output format context + outputFormatContext.SetPb(ioContext) + } + + // Set the start and end time for the desired range (in seconds) + startTime := 10 + endTime := 20 + + // Seek to the start time + if err := inputFormatContext.SeekFrame(-1, int64(startTime*astiav.TimeBase), astiav.NewSeekFlags(astiav.SeekFlagBackward)); err != nil { + log.Fatal(fmt.Errorf("main: seeking to start time failed: %w", err)) + } + + // Write header + if err = outputFormatContext.WriteHeader(nil); err != nil { + log.Fatal(fmt.Errorf("main: writing header failed: %w", err)) + } + + var startPts int64 + for { + // Read frame + if err = inputFormatContext.ReadFrame(pkt); err != nil { + if errors.Is(err, astiav.ErrEof) { + break + } + log.Fatal(fmt.Errorf("main: reading frame failed: %w", err)) + } + + // Get input stream + inputStream, ok := inputStreams[pkt.StreamIndex()] + if !ok { + pkt.Unref() + continue + } + + // Get output stream + outputStream, ok := outputStreams[pkt.StreamIndex()] + if !ok { + pkt.Unref() + continue + } + + // Check if the packet is within the desired time range + if pkt.Pts() != astiav.NoPtsValue && pkt.Pts()*int64(inputStream.TimeBase().Num()) > int64(endTime*inputStream.TimeBase().Den()) { + break + } + + if startPts == 0 { + startPts = pkt.Pts() + } + + pkt.SetDts(pkt.Dts() - startPts) + pkt.SetPts(pkt.Pts() - startPts) + + // Update packet + pkt.SetStreamIndex(outputStream.Index()) + pkt.RescaleTs(inputStream.TimeBase(), outputStream.TimeBase()) + pkt.SetPos(-1) + + // Write frame + if err = outputFormatContext.WriteInterleavedFrame(pkt); err != nil { + log.Fatal(fmt.Errorf("main: writing interleaved frame failed: %w", err)) + } + } + + // Write trailer + if err = outputFormatContext.WriteTrailer(); err != nil { + log.Fatal(fmt.Errorf("main: writing trailer failed: %w", err)) + } + + // Success + log.Println("success") +}