package sdk
|
|
import (
|
"analysis/logo"
|
"analysis/work"
|
"context"
|
"fmt"
|
"plugin"
|
|
"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
|
|
list *LockList
|
|
handle interface{}
|
fnInit func(string, string, string, int, func(...interface{})) 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,
|
|
list: NewLockList(6),
|
|
handle: nil,
|
fnInit: fnInit.(func(string, string, string, int, func(...interface{})) 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)
|
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) {
|
FlowSimple(ctx, in, out, typ, y.list.Push, y.list.Pop, y.track, func() { y.fnFree(y.handle) })
|
}
|
|
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
|
}
|
