From 0ead91b8788b382d0d1e7f6eda3f1d008ce39602 Mon Sep 17 00:00:00 2001
From: zhangmeng <775834166@qq.com>
Date: 星期二, 26 十二月 2023 10:34:01 +0800
Subject: [PATCH] add bgr and yuv420 to jpeg support

---
 capture.go |  144 +++++++++++++++++++++++++++++++++++++++--------
 1 files changed, 119 insertions(+), 25 deletions(-)

diff --git a/capture.go b/capture.go
index 050d85c..3164d11 100644
--- a/capture.go
+++ b/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,21 +23,54 @@
 // 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("")
 )
+
+// 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 := fmt.Errorf("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 {
@@ -51,48 +84,59 @@
 		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 = maxTry * 1000
+
+	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 {
+			if ow*oh*3 != len(data) {
+				// data 涓嶆槸 bgr 鏍煎紡锛岄渶瑕佽浆鎹�
+				bgrData = yuv2bgr(data, ow, oh)
+			} else {
+				bgrData = data
+			}
+			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
 		} else {
-			tryTime++
+			tryTime += 40
 			if tryTime > maxTry {
 				break
 			}
-			time.Sleep(time.Second)
+			time.Sleep(40 * time.Millisecond)
 		}
 	}
 
@@ -101,5 +145,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)
 }

--
Gitblit v1.8.0