add bgr and yuv420 to jpeg support

zhangmeng

2023-12-26 0ead91b8788b382d0d1e7f6eda3f1d008ce39602

 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,24 +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) {
   gb := false
   if m == GB28181 {
      gb = true
   } else {
   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 {
@@ -54,64 +84,116 @@
      return nil, err
   }

   conf := goffmpeg.Config{
      Scale:  goffmpeg.ScalePoint,
      Width:  w,
      Height: h,
      GB:     gb,
      CPU:    false,
   }
   gf := goffmpeg.New(conf)
   ret, err = capt2Data(url, m, ext, w, h, maxTry)

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

   gf.Free()
   goffmpeg.FreeFFmpeg()

   return ret, err
}

func capt2Data(gf *goffmpeg.GoFFMPEG, url string, 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
   }

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

   gf.Run(url)
   gf.BuildDecoder()

   tryTime := 0
   // 保证有一个关键帧
   maxTry = maxTry * 1000

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

      data, wid, hei := gf.GetPicDecoder()
      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)
      }
   }
   gf.Free()

   if tryTime > maxTry {
   if tryTime >= maxTry {
      err := fmt.Sprintf("try %d times to capture image, is url correct?\n", tryTime)
      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)
}