package sdk
|
|
import (
|
"analysis/goconv"
|
"analysis/logo"
|
"analysis/work"
|
"context"
|
|
"github.com/gogo/protobuf/proto"
|
|
"analysis/gohumantrack"
|
|
"basic.com/pubsub/protomsg.git"
|
"basic.com/valib/gogpu.git"
|
)
|
|
type imageWithID struct {
|
img *gohumantrack.ImageHumanTracker
|
fx float64
|
fy float64
|
id string
|
}
|
|
// HumanTracker track
|
type HumanTracker struct {
|
tracker *gohumantrack.HumanTracker
|
gpu int
|
batchSize int
|
flag int
|
list *LockList
|
|
mapCameraImageIndex map[string]int
|
wait2RunMoreThan int
|
index int
|
images []*imageWithID
|
msgs []*work.MsgRS
|
}
|
|
// NewHumanTracker new
|
func NewHumanTracker(gpu, batchSize, flag int) *HumanTracker {
|
return &HumanTracker{
|
gpu: gpu,
|
batchSize: batchSize,
|
flag: flag,
|
list: NewLockList(6),
|
|
wait2RunMoreThan: 0,
|
index: 0,
|
mapCameraImageIndex: make(map[string]int),
|
images: make([]*imageWithID, batchSize),
|
msgs: make([]*work.MsgRS, batchSize),
|
}
|
}
|
|
// Free free
|
func (t *HumanTracker) Free() {
|
if t.tracker != nil {
|
t.tracker.Free()
|
}
|
}
|
|
// Init impl
|
func (t *HumanTracker) Init() bool {
|
|
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[0])
|
result := protomsg.HumanTrackResult{Result: hr}
|
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))
|
|
}
|
|
func (t *HumanTracker) trackBatch(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 imgData []byte
|
rw, rh := 1280, 720
|
fx, fy := 1.0, 1.0
|
if imgW != rw || imgH != rh {
|
imgData = goconv.ResizeBGR(i.Data, imgW, imgH, rw, rh)
|
fx = (float64)(rw) / (float64)(imgW)
|
fy = (float64)(rh) / (float64)(imgH)
|
}
|
img := gohumantrack.ImageHumanTracker{
|
Data: imgData,
|
Width: rw,
|
Height: rh,
|
Channel: 3,
|
}
|
|
// 查找是否有这个摄像机图片的位置,有就替换成最新的,并且保存rMsg
|
if i, ok := t.mapCameraImageIndex[rMsg.Msg.Cid]; ok {
|
if i < t.batchSize {
|
t.images[i] = &imageWithID{&img, fx, fy, rMsg.Msg.Cid}
|
t.msgs[i] = &rMsg
|
}
|
} else {
|
// 没有,如果index小于最大的batch,创建一个新的位置
|
if t.index < t.batchSize {
|
t.images[t.index] = &imageWithID{&img, fx, fy, rMsg.Msg.Cid}
|
t.msgs[t.index] = &rMsg
|
t.mapCameraImageIndex[rMsg.Msg.Cid] = t.index
|
}
|
t.index++
|
}
|
// 等收到次数大于?次后开始运行,因为不一定各自摄像机的图片都会到
|
t.wait2RunMoreThan++
|
if t.wait2RunMoreThan < t.batchSize+t.batchSize/2 {
|
return
|
}
|
|
for k, v := range t.mapCameraImageIndex {
|
logo.Infoln("batch~~~~~~Map index: ", v, " camera: ", k)
|
}
|
|
var pimg []*gohumantrack.ImageHumanTracker
|
for k, v := range t.images[:] {
|
if v != nil {
|
pimg = append(pimg, v.img)
|
logo.Infoln("batch~~~~~~Image index: ", k, " camera: ", v.id, " image address: ", v.img)
|
} else {
|
pimg = append(pimg, nil)
|
}
|
}
|
|
res, err := t.tracker.ProcessImagePointer(pimg, rw, rh, 3)
|
// 从0开始计算收到了多少次图片
|
t.wait2RunMoreThan = 0
|
// 如果index大于最大batch,说明摄像机可能改变了,清空所有结构重新开始
|
if t.index > t.batchSize {
|
t.mapCameraImageIndex = make(map[string]int)
|
}
|
|
if err != nil {
|
logo.Infoln("batch~~~~~~Track Image Count: ", t.index, " Failed: ", err)
|
// 将images设置为nil,重新开始一个流程
|
for i := 0; i < t.batchSize; i++ {
|
t.images[i] = nil
|
t.msgs[i] = nil
|
}
|
|
ejectResult(nil, rMsg, out)
|
return
|
}
|
|
for i := 0; i < t.batchSize; i++ {
|
if t.images[i] == nil {
|
continue
|
}
|
hr := convert2ProtoHumanTrackResultWithScale(res[i], t.images[i].fx, t.images[i].fy)
|
result := protomsg.HumanTrackResult{Result: hr}
|
data, err := proto.Marshal(&result)
|
if err != nil {
|
logo.Errorln("batch~~~~~~HUMAN TRACKER MARSHAL PROTO PLATE IDS ERROR", err)
|
data = nil
|
}
|
msg := *t.msgs[i]
|
ejectResult(data, msg, out)
|
var id, name string
|
if msg.Msg.Tasklab != nil {
|
id, name = msg.Msg.Tasklab.Taskid, msg.Msg.Tasklab.Taskname
|
}
|
logo.Infoln("batch~~~~~~CAMERAID: ", msg.Msg.Cid, " TASKID: ", id, " TASKNAME: ", name, " Human Track COUNT: ", len(hr))
|
|
}
|
// 将images设置为nil,重新开始一个流程
|
for i := 0; i < t.batchSize; i++ {
|
t.images[i] = nil
|
t.msgs[i] = nil
|
}
|
|
}
|
|
// 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)
|
FlowSimple(ctx, in, out, typ, t.list.Push, t.list.Pop, t.trackBatch, 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 {
|
res := []*protomsg.HumanTrack{}
|
for i := 0; i < int(obj.FgNum); i++ {
|
r := obj.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)
|
}
|
return res
|
}
|
|
func convert2ProtoHumanTrackResultWithScale(obj gohumantrack.FgResult, fx, fy float64) []*protomsg.HumanTrack {
|
res := []*protomsg.HumanTrack{}
|
for i := 0; i < int(obj.FgNum); i++ {
|
r := obj.Fginfo[i]
|
rect := protomsg.Rect{
|
Left: (int32)((float64)(r.Left) / fx),
|
Right: (int32)((float64)(r.Right) / fy),
|
Top: (int32)((float64)(r.Top) / fx),
|
Bottom: (int32)((float64)(r.Bottom) / fy),
|
}
|
pr := &protomsg.HumanTrack{
|
RcHuman: &rect,
|
Confidence: r.Confidence,
|
X: (int32)((float64)(r.X) / fx),
|
Y: (int32)((float64)(r.Y) / fy),
|
Id: r.ID,
|
Feature: r.Feature[:],
|
}
|
res = append(res, pr)
|
}
|
return res
|
}
|
