package demo
|
|
import (
|
"analysis/valib/gosdk"
|
srv "analysis/work/service"
|
"fmt"
|
"image"
|
"image/color"
|
"math/rand"
|
"strconv"
|
"sync"
|
|
"gocv.io/x/gocv"
|
)
|
|
func min(x, y int) int {
|
if x > y {
|
return y
|
}
|
return x
|
}
|
func max(x, y int) int {
|
if x > y {
|
return x
|
}
|
return y
|
}
|
|
func cvDrawObject(img *gocv.Mat, obj []gosdk.CObjInfo) {
|
clr := color.RGBA{255, 0, 0, 0}
|
clrText := color.RGBA{255, 255, 0, 0}
|
clrTextRect := color.RGBA{0, 255, 255, 0}
|
|
for k, o := range obj {
|
rect := o.RcObj
|
r := gosdk.Rect(rect)
|
objName := gosdk.YoloObjName(int(o.Typ))
|
if len(objName) == 0 {
|
continue
|
}
|
|
text := objName + "-" + strconv.Itoa(k)
|
|
textSize := gocv.GetTextSize(text, gocv.FontHersheyComplexSmall, 1.2, 2)
|
maxWidth := textSize.X
|
if textSize.X > r.Dx()+2 {
|
maxWidth = r.Dx() + 2
|
}
|
textRect := image.Rect(max(r.Min.X-1, 0), max(r.Min.Y-30, 0), min(r.Min.X+maxWidth, r.Min.X+r.Dx()-1), min(r.Min.Y, r.Min.Y+r.Dy()-1))
|
gocv.Rectangle(img, textRect, clrTextRect, -1)
|
|
textPt := image.Pt(r.Min.X, r.Min.Y-10)
|
gocv.PutText(img, text, textPt, gocv.FontHersheyComplexSmall, 1.2, clrText, 2)
|
gocv.Rectangle(img, r, clr, 3)
|
}
|
}
|
|
type Yolo struct {
|
yoloDetInfo []gosdk.CObjInfo
|
startYolo bool
|
m sync.Mutex
|
}
|
|
func NewYolo() *Yolo {
|
return &Yolo{
|
startYolo: false,
|
}
|
}
|
|
// ShowYolo show yolo result
|
func (y *Yolo) ShowYolo(ch <-chan srv.ImageInfo) {
|
y.ShowYoloWithName(ch, "Yolo")
|
}
|
|
// ShowYoloWithName show yolo result
|
func (y *Yolo) ShowYoloWithName(ch <-chan srv.ImageInfo, winName string) {
|
|
cfg := "./data/yolo/cfg/yolov3.cfg"
|
weights := "./data/yolo/yolov3.weights"
|
name := "./data/yolo/data/coco.names"
|
|
yolo := gosdk.InitYolo(cfg, weights, name, 0)
|
|
winName = winName + strconv.Itoa(rand.Intn(10))
|
window := gocv.NewWindow(winName)
|
|
fmt.Println("create window : ", winName)
|
|
pic := gocv.NewMat()
|
|
for {
|
|
i := <-ch
|
|
if !y.startYolo {
|
go func(yolo *gosdk.YoloHandle, i srv.ImageInfo) {
|
y.startYolo = true
|
img := gosdk.SDKImage{Data: i.Data, Width: i.Width, Height: i.Height}
|
obj := gosdk.YoloDetect(yolo, img, 0.4, 0)
|
|
y.m.Lock()
|
y.yoloDetInfo = obj
|
y.m.Unlock()
|
y.startYolo = false
|
}(yolo, i)
|
}
|
|
img, _ := gocv.NewMatFromBytes(i.Height, i.Width, gocv.MatTypeCV8UC3, i.Data)
|
|
var obj []gosdk.CObjInfo
|
y.m.Lock()
|
obj = y.yoloDetInfo
|
y.m.Unlock()
|
cvDrawObject(&img, obj)
|
// show obj
|
gocv.Resize(img, &pic, image.Pt(0, 0), 0.5, 0.5, gocv.InterpolationDefault)
|
window.IMShow(pic)
|
window.WaitKey(1)
|
|
}
|
}
|
