update

zhangmeng

2019-11-26 bf954c0fa627c53f6e105a9d5bb8a34b88304ce7

update

 capture.go

@@ -3,11 +3,11 @@
import (
    "errors"
    "fmt"
    "image"
    "time"

    "basic.com/valib/godraw.git"
    "basic.com/valib/goffmpeg.git"
    "gocv.io/x/gocv"
    "github.com/disintegration/imaging"
)

// StreamMode stream mode
@@ -23,13 +23,17 @@
// PicExt file extention
type PicExt string

// PNGFileExt FileExt = ".png"
// JPEGFileExt FileExt = ".jpg"
// GIFFileExt FileExt = ".gif"

const (
    // PNGFileExt png
    PNGFileExt PicExt = PicExt(gocv.PNGFileExt)
    PNGFileExt PicExt = PicExt(".png")
    // JPEGFileExt jpeg
    JPEGFileExt PicExt = PicExt(gocv.JPEGFileExt)
    JPEGFileExt PicExt = PicExt(".jpg")
    // GIFFileExt gif
    GIFFileExt PicExt = PicExt(gocv.GIFFileExt)
    GIFFileExt PicExt = PicExt(".gif")
    // InvalidFileExt invalid
    InvalidFileExt PicExt = PicExt("")
)
@@ -51,39 +55,49 @@
        return nil, err
    }

    ret, err = capt2Data(url, m, gocv.FileExt(ext), w, h, maxTry)
    ret, err = capt2Data(url, m, ext, w, h, maxTry)

    goffmpeg.FreeFFmpeg()

    return ret, err
}

func capt2Data(url string, m StreamMode, ext gocv.FileExt, w, h, maxTry int) ([]byte, error) {
func capt2Data(url string, m StreamMode, ext PicExt, w, h, maxTry int) ([]byte, error) {

    gb := false
    if m == GB28181 {
        gb = true
    }

    tryTime := 0
    gf := goffmpeg.New(gb, false)
    defer gf.Free()

    var picData []byte
    var err2 error
    gf.Run(url)
    gf.BuildDecoder()

    tryTime := 0
    // 保证有一个关键帧
    maxTry += 50
    if maxTry > 60 {
        maxTry = 60
    }

    var bgrData, jpgData []byte
    var err error
    for {

        data, wid, hei := goffmpeg.Decode(url, gb)
        if wid > 0 && hei > 0 {
            img, err := gocv.NewMatFromBytes(hei, wid, gocv.MatTypeCV8UC3, data)
        data, ow, oh, _ := gf.GetYUV()

        if ow > 0 && oh > 0 {

            bgrData = yuv2bgr(data, ow, oh)
            if w > 0 && h > 0 {
                bgrData = bgresize(bgrData, ow, oh, w, h)
                ow, oh = w, h
            }

            jpgData, err = godraw.ToJpeg(bgrData, ow, oh, nil)
            if err != nil {
                continue
            }
            pic := img
            if w >= 0 && h >= 0 {
                pic = gocv.NewMatWithSize(h, w, gocv.MatTypeCV8UC3)
                gocv.Resize(img, &pic, image.Pt(w, h), 0, 0, gocv.InterpolationDefault)
            }
            picData, err2 = gocv.IMEncode(ext, pic)
            if err2 != nil {
                continue
            }
            break
@@ -92,7 +106,7 @@
            if tryTime > maxTry {
                break
            }
            time.Sleep(time.Second)
            time.Sleep(800 * time.Millisecond)
        }
    }

@@ -101,5 +115,55 @@
        return nil, errors.New(err)
    }

    return picData, nil
    return jpgData, nil
}

func yuv2bgr(yuv []byte, w, h int) []byte {

    data := make([]byte, 0, w*h*3)
    start := w * h
    for i := 0; i < h; i++ {
        for j := 0; j < w; j++ {

            index := i/2*w + j - (j & 0x01)

            y := int32(yuv[j+i*w])
            v := int32(yuv[start+index])
            u := int32(yuv[start+index+1])

            r := y + (140*(v-128))/100
            g := y - (34*(u-128)+71*(v-128))/100
            b := y + (177*(u-128))/100

            if r > 255 {
                r = 255
            }
            if r < 0 {
                r = 0
            }
            if g > 255 {
                g = 255
            }
            if g < 0 {
                g = 0
            }
            if b > 255 {
                b = 255
            }
            if b < 0 {
                b = 0
            }
            data = append(data, byte(r), byte(g), byte(b))
        }
    }
    return data
}

func bgresize(bgr []byte, w, h, rw, rh int) []byte {
    img, err := godraw.ToImage(bgr, w, h)
    if err != nil {
        return nil
    }
    dstImg := imaging.Resize(img, rw, rh, imaging.NearestNeighbor)
    return godraw.Image2BGR(dstImg)
}