package sdk
|
|
import (
|
"analysis/logo"
|
"analysis/work"
|
"context"
|
|
"github.com/gogo/protobuf/proto"
|
|
"analysis/gohumantrack"
|
|
"basic.com/pubsub/protomsg.git"
|
"basic.com/valib/gogpu.git"
|
)
|
|
// HumanTracker track
|
type HumanTracker struct {
|
tracker *gohumantrack.HumanTracker
|
gpu int
|
batchSize int
|
flag int
|
list *LockList
|
}
|
|
// NewHumanTracker new
|
func NewHumanTracker(gpu, batchSize, flag int) *HumanTracker {
|
return &HumanTracker{
|
gpu: gpu,
|
batchSize: batchSize,
|
flag: flag,
|
list: NewLockList(6),
|
}
|
}
|
|
// Free free
|
func (t *HumanTracker) Free() {
|
if t.tracker != nil {
|
t.tracker.Free()
|
}
|
}
|
|
// 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)
|
}
|
logo.Infof("start gpu %d, batch %d, flag %d\n", gpu, t.batchSize, t.flag)
|
t.tracker = gohumantrack.NewHumanTracker(gpu, t.batchSize, t.flag)
|
if t.tracker != nil {
|
logo.Infoln("Start Success")
|
return true
|
}
|
logo.Infoln("Start Failed")
|
return false
|
}
|
|
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 []gohumantrack.ImageHumanTracker
|
img := gohumantrack.ImageHumanTracker{
|
Data: i.Data,
|
Width: imgW,
|
Height: imgH,
|
Channel: 3,
|
}
|
images = append(images, img)
|
res, err := t.tracker.Process(images)
|
if err != nil {
|
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 PLATE IDS 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, " Human Track 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, t.Free)
|
}
|
|
// 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 []gohumantrack.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
|
}
|
