package ffmpeg // #cgo pkg-config: libavformat libavutil libavcodec libswscale // #cgo CFLAGS: -std=c11 // #cgo LDFLAGS: -lm // #include "ffmpeg.h" import "C" import ( "context" "github.com/cshum/imagorvideo/ffmpeg/pointer" "io" "time" "unsafe" ) const ( readPacketFlag = 1 seekPacketFlag = 2 interruptFlag = 3 hasVideo = 1 hasAudio = 2 ) type Metadata struct { Orientation int `json:"orientation"` Duration time.Duration `json:"duration,omitempty"` Width int `json:"width,omitempty"` Height int `json:"height,omitempty"` Title string `json:"title,omitempty"` Artist string `json:"artist,omitempty"` HasVideo bool `json:"has_video"` HasAudio bool `json:"has_audio"` HasAlpha bool `json:"has_alpha"` } type AVContext struct { context context.Context opaque unsafe.Pointer reader io.Reader seeker io.Seeker formatContext *C.AVFormatContext stream *C.AVStream codecContext *C.AVCodecContext thumbContext *C.ThumbContext frame *C.AVFrame durationInFormat bool orientation int size int64 duration time.Duration width, height int title, artist string hasVideo, hasAudio bool hasFrame, hasAlpha bool } func LoadAVContext(ctx context.Context, reader io.Reader, size int64) (*AVContext, error) { av := &AVContext{ context: ctx, reader: reader, size: size, } if seeker, ok := reader.(io.Seeker); ok { av.seeker = seeker } flags := C.int(readPacketFlag | interruptFlag) if av.seeker != nil { flags |= seekPacketFlag } err := createFormatContext(av, flags) if err != nil { return nil, err } if !av.hasVideo { return av, nil } if err = createDecoder(av); err == ErrTooBig || err == ErrDecoderNotFound { return av, err } return av, nil } func (av *AVContext) ExportImage() ([]byte, error) { if !av.hasFrame { return nil, nil } return encodeFrameImage(av) } func (av *AVContext) Close() { freeFormatContext(av) } func (av *AVContext) Metadata() *Metadata { return &Metadata{ Orientation: av.orientation, Duration: av.duration, Width: av.width, Height: av.height, Title: av.title, Artist: av.artist, HasVideo: av.hasVideo, HasAudio: av.hasAudio, HasAlpha: av.hasAlpha, } } func freeFormatContext(av *AVContext) { C.free_format_context(av.formatContext) pointer.Unref(av.opaque) } func createFormatContext(av *AVContext, callbackFlags C.int) error { intErr := C.allocate_format_context(&av.formatContext) if intErr < 0 { return avError(intErr) } av.opaque = pointer.Save(av) intErr = C.create_format_context(av.formatContext, av.opaque, callbackFlags) if intErr < 0 { pointer.Unref(av.opaque) return avError(intErr) } metaData(av) duration(av) err := findStreams(av) if err != nil { freeFormatContext(av) } return err } func metaData(av *AVContext) { var artist, title *C.char C.get_metadata(av.formatContext, &artist, &title) av.artist = C.GoString(artist) av.title = C.GoString(title) } func duration(av *AVContext) { if av.formatContext.duration > 0 { av.durationInFormat = true av.duration = time.Duration(1000 * av.formatContext.duration) } } func fullDuration(av *AVContext) error { if av.durationInFormat { return nil } newDuration := time.Duration(C.find_duration(av.formatContext)) if newDuration < 0 { return avError(newDuration) } if newDuration > av.duration { av.duration = newDuration } return nil } func findStreams(av *AVContext) error { var orientation C.int err := C.find_streams(av.formatContext, &av.stream, &orientation) if err < 0 { return avError(err) } av.hasVideo = err&hasVideo != 0 av.hasAudio = err&hasAudio != 0 if av.hasVideo { av.width = int(av.stream.codecpar.width) av.height = int(av.stream.codecpar.height) av.orientation = int(orientation) } return nil } func createDecoder(av *AVContext) error { err := C.create_codec_context(av.stream, &av.codecContext) if err < 0 { return avError(err) } defer C.avcodec_free_context(&av.codecContext) return createThumbContext(av) } func incrementDuration(av *AVContext, frame *C.AVFrame) { if !av.durationInFormat && frame.pts != C.AV_NOPTS_VALUE { ptsToNano := C.int64_t(1000000000 * av.stream.time_base.num / av.stream.time_base.den) newDuration := time.Duration(frame.pts * ptsToNano) if newDuration > av.duration { av.duration = newDuration } } } func populateHistogram(av *AVContext, frames <-chan *C.AVFrame) <-chan struct{} { done := make(chan struct{}) go func() { var n C.int for frame := range frames { C.populate_histogram(av.thumbContext, n, frame) n++ } av.thumbContext.n = n done <- struct{}{} close(done) }() return done } func createThumbContext(av *AVContext) error { pkt := C.create_packet() var frame *C.AVFrame err := C.obtain_next_frame(av.formatContext, av.codecContext, av.stream.index, &pkt, &frame) if err >= 0 { incrementDuration(av, frame) av.thumbContext = C.create_thumb_context(av.stream, frame) if av.thumbContext == nil { err = C.int(ErrNoMem) } } if err < 0 { if pkt.buf != nil { C.av_packet_unref(&pkt) } if frame != nil { C.av_frame_free(&frame) } return avError(err) } defer C.free_thumb_context(av.thumbContext) frames := make(chan *C.AVFrame, av.thumbContext.max_frames) done := populateHistogram(av, frames) frames <- frame if pkt.buf != nil { C.av_packet_unref(&pkt) } return populateThumbContext(av, frames, done) } func populateThumbContext(av *AVContext, frames chan *C.AVFrame, done <-chan struct{}) error { pkt := C.create_packet() var frame *C.AVFrame var err C.int for i := C.int(1); i < av.thumbContext.max_frames; i++ { err = C.obtain_next_frame(av.formatContext, av.codecContext, av.stream.index, &pkt, &frame) if err < 0 { break } incrementDuration(av, frame) frames <- frame frame = nil } close(frames) if pkt.buf != nil { C.av_packet_unref(&pkt) } if frame != nil { C.av_frame_free(&frame) } <-done if err != 0 && err != C.int(ErrEOF) { return avError(err) } return convertFrameToRGB(av) } func convertFrameToRGB(av *AVContext) error { outputFrame := C.convert_frame_to_rgb(C.process_frames(av.thumbContext), av.thumbContext.alpha) if outputFrame == nil { return ErrNoMem } av.frame = outputFrame av.hasFrame = true av.hasAlpha = av.thumbContext.alpha != 0 return nil } func encodeFrameImage(av *AVContext) ([]byte, error) { pkt := C.create_packet() err := C.encode_frame_to_image(av.formatContext, av.frame, &pkt) if err < 0 { return nil, avError(err) } p := (*[1 << 30]byte)(unsafe.Pointer(pkt.data))[:pkt.size:pkt.size] if pkt.buf != nil { C.av_packet_unref(&pkt) } return p, nil }