package sdk
|
|
import (
|
"analysis/logo"
|
"analysis/work"
|
"container/list"
|
"context"
|
"fmt"
|
"plugin"
|
"sync"
|
"time"
|
|
"basic.com/libgowrapper/sdkstruct.git"
|
"basic.com/pubsub/protomsg.git"
|
"basic.com/valib/gogpu.git"
|
"github.com/gogo/protobuf/proto"
|
)
|
|
// YoloDetect yolo detect
|
type YoloDetect struct {
|
iGPU int
|
cfg string
|
weights string
|
name string
|
|
cache *list.List
|
cv *sync.Cond
|
cond bool
|
|
handle interface{}
|
fnInit func(string, string, string, int) interface{}
|
fnFree func(interface{})
|
fnRun func(interface{}, string, []byte, int, int, int, float32, int) ([]sdkstruct.CObjTrackInfo, []sdkstruct.CObjTrackInfo)
|
}
|
|
// NewYDetectWithTrack with track
|
func NewYDetectWithTrack(gi int, cfg, weights, name string) *YoloDetect {
|
soFile := "libyolo.so"
|
|
plug, err := plugin.Open(soFile)
|
if err != nil {
|
logo.Errorln("Open: ", soFile, " error: ", err)
|
return nil
|
}
|
|
fnInit, _ := LoadFunc(plug, soFile, "NewSDK")
|
fnFree, _ := LoadFunc(plug, soFile, "Free")
|
fnRun, _ := LoadFunc(plug, soFile, "Run")
|
|
return &YoloDetect{
|
iGPU: gi,
|
cfg: cfg,
|
weights: weights,
|
name: name,
|
cache: list.New(),
|
cv: sync.NewCond(&sync.Mutex{}),
|
cond: false,
|
|
handle: nil,
|
fnInit: fnInit.(func(string, string, string, int) interface{}),
|
fnFree: fnFree.(func(interface{})),
|
fnRun: fnRun.(func(interface{}, string, []byte, int, int, int, float32, int) ([]sdkstruct.CObjTrackInfo, []sdkstruct.CObjTrackInfo)),
|
}
|
}
|
|
// Init impl interface
|
func (y *YoloDetect) Init() bool {
|
|
gpu := y.iGPU
|
if gpu == -1 {
|
gpu = gogpu.ValidGPU(2048)
|
}
|
h := y.fnInit(y.cfg, y.weights, y.name, gpu)
|
logo.Infoln("YOLO USE GPU: ", gpu)
|
|
if h == nil {
|
logo.Errorln("CREATE YOLO DETECTOR ERROR")
|
return false
|
}
|
|
y.handle = h
|
return true
|
}
|
|
// Run impl interface
|
func (y *YoloDetect) Run(ctx context.Context, in <-chan work.MsgRS, out chan<- work.MsgRS, typ string) {
|
go y.work(ctx, out, typ)
|
|
for {
|
select {
|
case <-ctx.Done():
|
y.fnFree(y.handle)
|
return
|
default:
|
rMsg := <-in
|
if !validRemoteMessage(rMsg, typ) {
|
logo.Errorln("yolo validremotemessage invalid")
|
ejectResult(nil, rMsg, out)
|
continue
|
}
|
y.push(rMsg)
|
}
|
}
|
|
}
|
|
func (y *YoloDetect) track(rMsg work.MsgRS, out chan<- work.MsgRS, typ string) {
|
|
/////////////////////////////////////////////
|
i := unpackImage(rMsg, typ)
|
if i == nil || i.Data == nil || i.Width <= 0 || i.Height <= 0 {
|
logo.Errorln("yolo image error: ", i)
|
ejectResult(nil, rMsg, out)
|
|
return
|
}
|
|
imgW, imgH := int(i.Width), int(i.Height)
|
|
whole, recent := y.fnRun(y.handle, rMsg.Msg.Cid, i.Data, imgW, imgH, 3, 0.4, 0)
|
if len(recent) > 0 {
|
}
|
|
infos := convert2ProtoYoloTrack(whole, 1.0, 1.0)
|
p := protomsg.ParamYoloObj{Infos: infos}
|
|
data, err := proto.Marshal(&p)
|
if err != nil {
|
fmt.Println("ydetect track marshal proto yolo obj error", err)
|
data = nil
|
}
|
|
ejectResult(data, rMsg, out)
|
|
var id, name string
|
if rMsg.Msg.Tasklab != nil {
|
id, name = rMsg.Msg.Tasklab.Taskid, rMsg.Msg.Tasklab.Taskname
|
}
|
|
logo.Infoln("CAMERAID: ", rMsg.Msg.Cid, " TASKID: ", id, " TASKNAME: ", name, " DETECT YOLO COUNT: ", len(whole))
|
}
|
|
func convert2ProtoYoloTrack(obj []sdkstruct.CObjTrackInfo, fx, fy float64) []*protomsg.ObjInfo {
|
ret := []*protomsg.ObjInfo{}
|
|
for _, v := range obj {
|
if fx < 1.0 || fy < 1.0 {
|
v.ObjInfo.RcObj.Left = (int32)((float64)(v.ObjInfo.RcObj.Left) / fx)
|
v.ObjInfo.RcObj.Right = (int32)((float64)(v.ObjInfo.RcObj.Right) / fx)
|
v.ObjInfo.RcObj.Top = (int32)((float64)(v.ObjInfo.RcObj.Top) / fy)
|
v.ObjInfo.RcObj.Bottom = (int32)((float64)(v.ObjInfo.RcObj.Bottom) / fy)
|
}
|
|
rect := protomsg.Rect{
|
Left: v.ObjInfo.RcObj.Left,
|
Right: v.ObjInfo.RcObj.Right,
|
Top: v.ObjInfo.RcObj.Top,
|
Bottom: v.ObjInfo.RcObj.Bottom,
|
}
|
obj := protomsg.ObjInfo{
|
RcObj: &rect,
|
Typ: v.ObjInfo.Typ,
|
Prob: v.ObjInfo.Prob,
|
ObjID: v.ID,
|
}
|
|
ret = append(ret, &obj)
|
}
|
return ret
|
}
|
|
func (y *YoloDetect) work(ctx context.Context, out chan<- work.MsgRS, typ string) {
|
tm := time.Now()
|
sc := 0
|
|
for {
|
select {
|
case <-ctx.Done():
|
return
|
default:
|
|
y.cv.L.Lock()
|
|
for !y.cond {
|
y.cv.Wait()
|
}
|
|
rMsg := y.cache.Front().Value.(work.MsgRS)
|
|
y.track(rMsg, out, typ)
|
y.cache.Remove(y.cache.Front())
|
y.cond = false
|
y.cv.L.Unlock()
|
|
sc++
|
if sc == 25 {
|
logo.Infoln("YOLO RUN 25 FRAME USE TIME: ", time.Since(tm))
|
sc = 0
|
tm = time.Now()
|
}
|
if time.Since(tm) > time.Second {
|
logo.Infof("YOLO RUN %d FRAME USE TIME: %v", sc, time.Since(tm))
|
sc = 0
|
tm = time.Now()
|
}
|
}
|
}
|
|
}
|
|
func (y *YoloDetect) push(data work.MsgRS) {
|
y.cv.L.Lock()
|
y.cache.PushBack(data)
|
if y.cache.Len() > 12 {
|
for i := 0; i < y.cache.Len(); {
|
y.cache.Remove(y.cache.Front())
|
i = i + 2
|
}
|
}
|
// logo.Infof("push to cache count : %d\n", t.cache.Len())
|
y.cond = true
|
y.cv.Signal()
|
y.cv.L.Unlock()
|
}
|
