package sdk
|
|
import (
|
"analysis/logo"
|
"analysis/work"
|
"context"
|
"plugin"
|
|
"github.com/gogo/protobuf/proto"
|
|
"basic.com/libgowrapper/sdkstruct.git"
|
"basic.com/pubsub/protomsg.git"
|
"basic.com/valib/gogpu.git"
|
)
|
|
// HumanTracker track
|
type HumanTracker struct {
|
gpu int
|
batchSize int
|
flag int
|
|
list *LockList
|
|
handle interface{}
|
fnInit func(int, int, int, func(...interface{})) interface{}
|
fnFree func(interface{})
|
fnRun func(interface{}, []byte, int, int, int) []sdkstruct.FgResult
|
fnProcess func(interface{}, []sdkstruct.SDKImage) ([]sdkstruct.FgResult, error)
|
fnVer func(interface{}) string
|
fnSimilarity func([]float32, []float32) (float64, error)
|
}
|
|
// NewHumanTracker new
|
func NewHumanTracker(gpu, batchSize, flag int) *HumanTracker {
|
soFile := "libhumantrack.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")
|
fnProc, _ := LoadFunc(plug, soFile, "Process")
|
fnVer, _ := LoadFunc(plug, soFile, "GetVersion")
|
fnSim, _ := LoadFunc(plug, soFile, "CFFSimilarity")
|
|
return &HumanTracker{
|
gpu: gpu,
|
batchSize: batchSize,
|
flag: flag,
|
|
list: NewLockList(6),
|
|
handle: nil,
|
fnInit: fnInit.(func(int, int, int, func(...interface{})) interface{}),
|
fnFree: fnFree.(func(interface{})),
|
fnRun: fnRun.(func(interface{}, []byte, int, int, int) []sdkstruct.FgResult),
|
fnProcess: fnProc.(func(interface{}, []sdkstruct.SDKImage) ([]sdkstruct.FgResult, error)),
|
fnVer: fnVer.(func(interface{}) string),
|
fnSimilarity: fnSim.(func([]float32, []float32) (float64, error)),
|
}
|
}
|
|
// Init impl
|
func (t *HumanTracker) Init() bool {
|
if t.batchSize != 1 {
|
logo.Errorln("ONLY SUPPORT BATCH SIZE = 1")
|
return false
|
}
|
gpu := t.gpu
|
if gpu == -1 {
|
gpu = gogpu.ValidGPU(2048)
|
}
|
h := t.fnInit(gpu, t.batchSize, t.flag, logo.Infoln)
|
logo.Infoln("HumanTrack USE GPU: ", gpu)
|
|
if h == nil {
|
logo.Errorln("CREATE HumanTrack DETECTOR ERROR")
|
return false
|
}
|
|
t.handle = h
|
return true
|
}
|
|
func (t *HumanTracker) 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 {
|
ejectResult(nil, rMsg, out)
|
return
|
}
|
|
imgW, imgH := int(i.Width), int(i.Height)
|
|
// var images []sdkstruct.SDKImage
|
// img := sdkstruct.SDKImage{
|
// Data: i.Data,
|
// Width: imgW,
|
// Height: imgH,
|
// Channel: 3,
|
// }
|
// images = append(images, img)
|
res := t.fnRun(t.handle, i.Data, imgW, imgH, 3)
|
if res == nil {
|
logo.Errorln("HUMAN TRACKER RUN ERROR")
|
|
ejectResult(nil, rMsg, out)
|
return
|
}
|
|
hr := convert2ProtoHumanTrackResult(res)
|
result := protomsg.HumanTrackResult{Result: hr[0]}
|
data, err := proto.Marshal(&result)
|
if err != nil {
|
logo.Errorln("HUMAN TRACKER MARSHAL PROTO 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(hr[0]))
|
|
}
|
|
// Run impl
|
func (t *HumanTracker) Run(ctx context.Context, in <-chan work.MsgRS, out chan<- work.MsgRS, typ string) {
|
|
FlowSimple(ctx, in, out, typ, t.list.Push, t.list.Pop, t.track, func() { t.fnFree(t.handle) })
|
}
|
|
// message HumanTrack {
|
// Rect rcHuman = 1;
|
// float confidence = 2;
|
// int32 x = 3;
|
// int32 y = 4;
|
// int32 id = 5;
|
// repeated float feature = 6;
|
// }
|
|
// message HumanTrackResult {
|
// repeated HumanTrack result = 1;
|
// }
|
|
func convert2ProtoHumanTrackResult(obj []sdkstruct.FgResult) [][]*protomsg.HumanTrack {
|
ret := [][]*protomsg.HumanTrack{}
|
for _, v := range obj {
|
res := []*protomsg.HumanTrack{}
|
for i := 0; i < int(v.FgNum); i++ {
|
r := v.Fginfo[i]
|
|
rect := protomsg.Rect{
|
Left: r.Left,
|
Right: r.Right,
|
Top: r.Top,
|
Bottom: r.Bottom,
|
}
|
pr := &protomsg.HumanTrack{
|
RcHuman: &rect,
|
Confidence: r.Confidence,
|
X: r.X,
|
Y: r.Y,
|
Id: r.ID,
|
Feature: r.Feature[:],
|
}
|
res = append(res, pr)
|
}
|
ret = append(ret, res)
|
}
|
return ret
|
}
|
