func NewCapture(driver, device string) (*Capture, error) {
id := Capture{index: -1}
id.context = C.avformat_alloc_context()
if id.context == (*C.AVFormatContext)(null) {
return nil, fmt.Errorf("allocate output format context failed")
}
_driver := C.CString(driver)
defer C.free(unsafe.Pointer(_driver))
ifmt := C.av_find_input_format(_driver)
if ifmt == (*C.AVInputFormat)(null) {
return nil, fmt.Errorf("cannot find input driver: %s", driver)
}
dev := C.CString(device)
defer C.free(unsafe.Pointer(dev))
if C.avformat_open_input(&(id.context), dev, ifmt, (**C.AVDictionary)(null)) < 0 {
return nil, fmt.Errorf("cannot open device %s", device)
}
if C.avformat_find_stream_info(id.context, (**C.AVDictionary)(null)) < 0 {
return nil, fmt.Errorf("cannot find stream information")
}
num := int(id.context.nb_streams)
streams := (*[1 << 30]*C.AVStream)(unsafe.Pointer(id.context.streams))
var deCtx *C.AVCodecContext
for i := 0; i < num; i++ {
if streams[i].codec.codec_type == C.AVMEDIA_TYPE_VIDEO {
deCtx = streams[i].codec
id.index = i
break
}
}
if id.index == -1 {
return nil, fmt.Errorf("cannot find video stream")
}
codec := C.avcodec_find_decoder(deCtx.codec_id)
if codec == (*C.AVCodec)(null) {
return nil, fmt.Errorf("cannot find decode codec")
}
id.codec = C.avcodec_alloc_context3(codec)
if C.avcodec_copy_context(id.codec, deCtx) != 0 {
return nil, fmt.Errorf("cannot copy codec context")
}
if C.avcodec_open2(id.codec, codec, (**C.struct_AVDictionary)(null)) < 0 {
return nil, fmt.Errorf("cannot open decode codec")
}
id.sws = C.sws_getContext(id.codec.width,
id.codec.height,
id.codec.pix_fmt,
id.codec.width,
id.codec.height,
C.AV_PIX_FMT_YUV420P, C.SWS_BILINEAR, (*C.struct_SwsFilter)(null), (*C.struct_SwsFilter)(null), (*C.double)(null))
id.frame = C.av_frame_alloc()
return &id, nil
}
func NewPicSwsCtx(srcWidth int, srcHeight int, srcPixFmt int32, dst *CodecCtx, method int) *SwsCtx {
ctx := C.sws_getContext(C.int(srcWidth), C.int(srcHeight), srcPixFmt, C.int(dst.Width()), C.int(dst.Height()), dst.PixFmt(), C.int(method), nil, nil, nil)
if ctx == nil {
return nil
}
return &SwsCtx{swsCtx: ctx}
}
func NewSwsCtx(src *CodecCtx, dst *CodecCtx, method int) *SwsCtx {
ctx := C.sws_getContext(C.int(src.Width()), C.int(src.Height()), src.PixFmt(), C.int(dst.Width()), C.int(dst.Height()), dst.PixFmt(), C.int(method), nil, nil, nil)
if ctx == nil {
return nil
}
return &SwsCtx{swsCtx: ctx}
}
func (g *Generator) NextFrame() (image.Image, int64, error) {
img := image.NewRGBA(image.Rect(0, 0, g.Width, g.Height))
frame := C.av_frame_alloc()
var pkt C.struct_AVPacket
var frameFinished C.int
for C.av_read_frame(g.avfContext, &pkt) == 0 {
if int(pkt.stream_index) != g.vStreamIndex {
C.av_free_packet(&pkt)
continue
}
if C.avcodec_decode_video2(g.avcContext, frame, &frameFinished, &pkt) <= 0 {
C.av_free_packet(&pkt)
return nil, 0, errors.New("can't decode frame")
}
C.av_free_packet(&pkt)
if frameFinished == 0 {
continue
}
ctx := C.sws_getContext(
C.int(g.Width),
C.int(g.Height),
g.avcContext.pix_fmt,
C.int(g.Width),
C.int(g.Height),
C.PIX_FMT_RGBA,
C.SWS_BICUBIC,
nil,
nil,
nil,
)
if ctx == nil {
return nil, 0, errors.New("can't allocate scaling context")
}
srcSlice := (**C.uint8_t)(&frame.data[0])
srcStride := (*C.int)(&frame.linesize[0])
dst := (**C.uint8_t)(unsafe.Pointer(&img.Pix))
dstStride := (*C.int)(unsafe.Pointer(&[1]int{img.Stride}))
C.sws_scale(
ctx,
srcSlice,
srcStride,
C.int(0),
g.avcContext.height,
dst,
dstStride,
)
break
}
timestamp := int64(C.av_frame_get_best_effort_timestamp(frame))
return img, timestamp, nil
}
func NewEncoder(codec uint32, in image.Image, out io.Writer) (*Encoder, error) {
_codec := C.avcodec_find_encoder(codec)
if _codec == nil {
return nil, fmt.Errorf("could not find codec")
}
c := C.avcodec_alloc_context3(_codec)
f := C.avcodec_alloc_frame()
c.bit_rate = 400000
// resolution must be a multiple of two
w, h := C.int(in.Bounds().Dx()), C.int(in.Bounds().Dy())
if w%2 == 1 || h%2 == 1 {
return nil, fmt.Errorf("Bad image dimensions (%d, %d), must be even", w, h)
}
log.Printf("Encoder dimensions: %d, %d", w, h)
c.width = w
c.height = h
c.time_base = C.AVRational{1, 25} // FPS
c.gop_size = 10 // emit one intra frame every ten frames
c.max_b_frames = 1
underlying_im := image.NewYCbCr(in.Bounds(), image.YCbCrSubsampleRatio420)
c.pix_fmt = C.PIX_FMT_YUV420P
f.data[0] = ptr(underlying_im.Y)
f.data[1] = ptr(underlying_im.Cb)
f.data[2] = ptr(underlying_im.Cr)
f.linesize[0] = w
f.linesize[1] = w / 2
f.linesize[2] = w / 2
if C.avcodec_open2(c, _codec, nil) < 0 {
return nil, fmt.Errorf("could not open codec")
}
_swscontext := C.sws_getContext(w, h, C.PIX_FMT_RGB0, w, h, C.PIX_FMT_YUV420P,
C.SWS_BICUBIC, nil, nil, nil)
e := &Encoder{codec, in, underlying_im, out, _codec, c, _swscontext, f, make([]byte, 16*1024)}
return e, nil
}
// Scale a YUVImage and return the new YUVImage
func Scale(src *jpeg.YUVImage, opts ScaleOptions) (*jpeg.YUVImage, error) {
// Figure out what format we're dealing with
var srcFmt, dstFmt int32
var flags C.int
flags = C.SWS_FULL_CHR_H_INT | C.int(opts.Filter) | C.SWS_ACCURATE_RND
components := 3
var dst jpeg.YUVImage
dstFmt = C.PIX_FMT_YUV444P
dst.Format = jpeg.YUV444
switch src.Format {
case jpeg.YUV444:
srcFmt = C.PIX_FMT_YUV444P
flags |= C.SWS_FULL_CHR_H_INP
case jpeg.YUV422:
srcFmt = C.PIX_FMT_YUV422P
case jpeg.YUV440:
srcFmt = C.PIX_FMT_YUV440P
case jpeg.YUV420:
srcFmt = C.PIX_FMT_YUV420P
case jpeg.Grayscale:
srcFmt = C.PIX_FMT_GRAY8
dstFmt = C.PIX_FMT_GRAY8
components = 1
dst.Format = jpeg.Grayscale
}
// swscale can't handle images smaller than this; pad them
paddedDstWidth := opts.DstWidth
paddedSrcWidth := src.Width
padFactor := 1
for paddedDstWidth < 8 || paddedSrcWidth < 4 {
paddedDstWidth *= 2
paddedSrcWidth *= 2
padFactor *= 2
}
// Get the SWS context
sws := C.sws_getContext(C.int(paddedSrcWidth), C.int(src.Height), srcFmt,
C.int(paddedDstWidth), C.int(opts.DstHeight), dstFmt,
flags, nil, nil, nil)
if sws == nil {
return nil, errors.New("sws_getContext failed")
}
defer C.sws_freeContext(sws)
// We only need 3 planes, but libswscale is stupid and checks the alignment
// of all 4 pointers... better give it a dummy one.
var srcYUVPtr [4](*uint8)
var dstYUVPtr [4](*uint8)
var srcStrides [4](C.int)
var dstStrides [4](C.int)
dst.Width = opts.DstWidth
dst.Height = opts.DstHeight
dstStride := pad(paddedDstWidth, jpeg.AlignSize)
dstFinalPaddedWidth := pad(opts.DstWidth, jpeg.AlignSize)
dstPaddedHeight := pad(opts.DstHeight, jpeg.AlignSize)
// Allocate image planes and pointers
for i := 0; i < components; i++ {
dst.Stride[i] = dstStride
dst.Data[i] = make([]byte, dstStride*dstPaddedHeight)
dstYUVPtr[i] = (*uint8)(unsafe.Pointer(&dst.Data[i][0]))
dstStrides[i] = C.int(dstStride)
// apply horizontal padding if image is too small
if padFactor > 1 {
planeWidth := src.PlaneWidth(i)
paddedWidth := planeWidth * padFactor
planeHeight := src.PlaneHeight(i)
paddedStride := pad(paddedWidth, jpeg.AlignSize)
newData := make([]uint8, paddedStride*planeHeight)
for y := 0; y < planeHeight; y++ {
copy(newData[y*paddedStride:], src.Data[i][y*src.Stride[i]:y*src.Stride[i]+planeWidth])
pixel := src.Data[i][y*src.Stride[i]+planeWidth-1]
for x := planeWidth; x < paddedWidth; x++ {
newData[y*paddedStride+x] = pixel
}
}
srcStrides[i] = C.int(paddedStride)
srcYUVPtr[i] = &newData[0]
} else {
srcStrides[i] = C.int(src.Stride[i])
srcYUVPtr[i] = (*uint8)(unsafe.Pointer(&src.Data[i][0]))
}
}
C.sws_scale(sws, (**C.uint8_t)(unsafe.Pointer(&srcYUVPtr[0])), &srcStrides[0], 0, C.int(src.Height),
(**C.uint8_t)(unsafe.Pointer(&dstYUVPtr[0])), &dstStrides[0])
// Replicate the last column and row of pixels as padding, which is typical
// behavior prior to JPEG compression
for i := 0; i < components; i++ {
for y := 0; y < dst.Height; y++ {
pixel := dst.Data[i][y*dstStride+dst.Width-1]
for x := dst.Width; x < dstFinalPaddedWidth; x++ {
dst.Data[i][y*dstStride+x] = pixel
}
}
lastRow := dst.Data[i][dstStride*(dst.Height-1) : dstStride*dst.Height]
for y := dst.Height; y < dstPaddedHeight; y++ {
copy(dst.Data[i][y*dstStride:], lastRow)
}
}
return &dst, nil
}
func sws_scale_getcontext(ctx *SwsContext, srcwidth, srcheight, srcfmt, trgwidth, trgheight, trgfmt, flags int) {
ctx.sws = C.sws_getContext(C.int(srcwidth), C.int(srcheight), int32(srcfmt), C.int(trgwidth), C.int(trgheight), int32(trgfmt), C.int(flags), nil, nil, nil)
if ctx.sws == nil {
println("error!")
}
}
//struct SwsContext * sws_getContext (int srcW, int srcH, enum AVPixelFormat srcFormat, int dstW, int dstH, enum AVPixelFormat dstFormat, int flags, SwsFilter *srcFilter, SwsFilter *dstFilter, const double *param)
//Allocate and return an SwsContext.
func Sws_getContext(sw, sh int, sf AVPixelFormat, dw, dh int, df AVPixelFormat, f int, sfl, dfl *SwsFilter, p *int) *SwsContext {
return (*SwsContext)(C.sws_getContext(C.int(sw), C.int(sh), (C.enum_AVPixelFormat)(sf), C.int(dw), C.int(dh), (C.enum_AVPixelFormat)(df), C.int(f), (*C.struct_SwsFilter)(sfl), (*C.struct_SwsFilter)(dfl), (*C.double)(unsafe.Pointer(p))))
}
// ImageWxH returns a screenshot at the ts milliseconds, scaled to the specified width and height.
func (g *Generator) ImageWxH(ts int64, width, height int) (image.Image, error) {
img := image.NewRGBA(image.Rect(0, 0, width, height))
frame := C.av_frame_alloc()
defer C.av_frame_free(&frame)
frameNum := C.av_rescale(
C.int64_t(ts),
C.int64_t(g.streams[g.vStreamIndex].time_base.den),
C.int64_t(g.streams[g.vStreamIndex].time_base.num),
) / 1000
if C.avformat_seek_file(
g.avfContext,
C.int(g.vStreamIndex),
0,
frameNum,
frameNum,
C.AVSEEK_FLAG_FRAME,
) < 0 {
return nil, errors.New("can't seek to timestamp")
}
C.avcodec_flush_buffers(g.avcContext)
var pkt C.struct_AVPacket
var frameFinished C.int
for C.av_read_frame(g.avfContext, &pkt) == 0 {
if int(pkt.stream_index) != g.vStreamIndex {
C.av_free_packet(&pkt)
continue
}
if C.avcodec_decode_video2(g.avcContext, frame, &frameFinished, &pkt) <= 0 {
C.av_free_packet(&pkt)
return nil, errors.New("can't decode frame")
}
C.av_free_packet(&pkt)
if frameFinished == 0 || pkt.dts < frameNum {
continue
}
ctx := C.sws_getContext(
C.int(g.Width),
C.int(g.Height),
g.avcContext.pix_fmt,
C.int(width),
C.int(height),
C.PIX_FMT_RGBA,
C.SWS_BICUBIC,
nil,
nil,
nil,
)
if ctx == nil {
return nil, errors.New("can't allocate scaling context")
}
srcSlice := (**C.uint8_t)(&frame.data[0])
srcStride := (*C.int)(&frame.linesize[0])
dst := (**C.uint8_t)(unsafe.Pointer(&img.Pix))
dstStride := (*C.int)(unsafe.Pointer(&[1]int{img.Stride}))
C.sws_scale(
ctx,
srcSlice,
srcStride,
0,
g.avcContext.height,
dst,
dstStride,
)
break
}
return img, nil
}
