package service
|
|
import (
|
"errors"
|
"fmt"
|
"time"
|
|
"basic.com/valib/godraw.git"
|
"basic.com/valib/goffmpeg.git"
|
"github.com/disintegration/imaging"
|
)
|
|
// StreamMode stream mode
|
type StreamMode int
|
|
const (
|
// Rtsp mode
|
Rtsp = StreamMode(iota)
|
// GB28181 mode
|
GB28181
|
)
|
|
// PicExt file extention
|
type PicExt string
|
|
// PNGFileExt FileExt = ".png"
|
// JPEGFileExt FileExt = ".jpg"
|
// GIFFileExt FileExt = ".gif"
|
|
const (
|
// PNGFileExt png
|
PNGFileExt PicExt = PicExt(".png")
|
// JPEGFileExt jpeg
|
JPEGFileExt PicExt = PicExt(".jpg")
|
// GIFFileExt gif
|
GIFFileExt PicExt = PicExt(".gif")
|
// InvalidFileExt invalid
|
InvalidFileExt PicExt = PicExt("")
|
)
|
|
// capture pic
|
func captureGB28181(soFile string, url string, maxTry int) ([]byte, error) {
|
|
var ret []byte
|
var err error
|
|
err = goffmpeg.InitFFmpeg(soFile)
|
if err != nil {
|
return nil, err
|
}
|
|
reterr := errors.New(fmt.Sprintf("gb28181 try %d times to capture image, is url correct?", maxTry))
|
for i := 0; i < maxTry; i++ {
|
ret = goffmpeg.GetGBJpg(url, maxTry)
|
if len(ret) > 0 {
|
reterr = nil
|
break
|
}
|
}
|
|
goffmpeg.FreeFFmpeg()
|
|
return ret, reterr
|
}
|
|
// Capture pic
|
func Capture(soFile string, m StreamMode, url string, ext PicExt, w, h, maxTry int) ([]byte, error) {
|
if m != GB28181 && m != Rtsp {
|
return nil, errors.New("there is no this mode, try capture.Rtsp/capture.GB28181")
|
}
|
|
if m == GB28181 {
|
return captureGB28181(soFile, url, maxTry)
|
}
|
|
if ext == InvalidFileExt {
|
ext = JPEGFileExt
|
}
|
var ret []byte
|
var err error
|
|
err = goffmpeg.InitFFmpeg(soFile)
|
if err != nil {
|
return nil, err
|
}
|
|
ret, err = capt2Data(url, m, ext, w, h, maxTry)
|
|
goffmpeg.FreeFFmpeg()
|
|
return ret, err
|
}
|
|
func capt2Data(url string, m StreamMode, ext PicExt, w, h, maxTry int) ([]byte, error) {
|
|
gb := false
|
if m == GB28181 {
|
gb = true
|
}
|
|
gf := goffmpeg.New(gb, false)
|
defer gf.Free()
|
|
gf.Run(url)
|
gf.BuildDecoder()
|
|
tryTime := 0
|
// 保证有一个关键帧
|
maxTry += 50
|
if maxTry > 60 {
|
maxTry = 60
|
}
|
|
var bgrData, jpgData []byte
|
var err error
|
for {
|
|
data, ow, oh, _ := gf.GetYUV()
|
if ow > 0 && oh > 0 {
|
if ow*oh*3 != len(data) {
|
// data 不是 bgr 格式，需要转换
|
bgrData = yuv2bgr(data, ow, oh)
|
} else {
|
bgrData = data
|
}
|
|
if w > 0 && h > 0 {
|
bgrData = bgresize(data, ow, oh, w, h)
|
ow, oh = w, h
|
}
|
|
jpgData, err = godraw.ToJpeg(bgrData, ow, oh, nil)
|
if err != nil {
|
continue
|
}
|
break
|
} else {
|
tryTime++
|
if tryTime > maxTry {
|
break
|
}
|
time.Sleep(800 * time.Millisecond)
|
}
|
}
|
|
if tryTime >= maxTry {
|
err := fmt.Sprintf("try %d times to capture image, is url correct?\n", tryTime)
|
return nil, errors.New(err)
|
}
|
|
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)
|
}
|
