From d8d2bc7b4cf8e4a1f2b2a325d2ca02f0eb409d9c Mon Sep 17 00:00:00 2001
From: zhangmeng <zhangmeng@aiotlink.com>
Date: 星期四, 19 十二月 2019 16:25:24 +0800
Subject: [PATCH] up
---
work/sdk/humantrack.go | 308 ++++++++++++++++++++++++++++++++++++---------------
1 files changed, 217 insertions(+), 91 deletions(-)
diff --git a/work/sdk/humantrack.go b/work/sdk/humantrack.go
index e8a82a9..46ad033 100644
--- a/work/sdk/humantrack.go
+++ b/work/sdk/humantrack.go
@@ -1,17 +1,25 @@
package sdk
import (
+ "analysis/goconv"
"analysis/logo"
"analysis/work"
"context"
- "time"
"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 {
@@ -19,6 +27,13 @@
gpu int
batchSize int
flag int
+ list *LockList
+
+ mapCameraImageIndex map[string]int
+ recvImageCount int
+ index int
+ images []*imageWithID
+ msgs []*work.MsgRS
}
// NewHumanTracker new
@@ -27,83 +42,176 @@
gpu: gpu,
batchSize: batchSize,
flag: flag,
+ list: NewLockList(6),
+
+ recvImageCount: 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 {
- 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
}
- t.tracker = gohumantrack.NewHumanTracker(t.gpu, t.batchSize, t.flag)
- return true
+ 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,
+ }
+
+ // mapCameraImageIndex map[string]int
+ // images []gohumantrack.ImageHumanTracker
+
+ 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 {
+ 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.recvImageCount++
+
+ if t.recvImageCount < 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)
+ t.recvImageCount = 0
+
+ if err != nil {
+ logo.Infoln("batch~~~~~~Track Image Count: ", t.index, " Failed: ", err)
+ 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))
+
+ }
+
}
// Run impl
func (t *HumanTracker) Run(ctx context.Context, in <-chan work.MsgRS, out chan<- work.MsgRS, typ string) {
- tm := time.Now()
- sc := 0
-
- for {
- select {
- case <-ctx.Done():
- return
- default:
- rMsg := <-in
- if !validRemoteMessage(rMsg, typ) {
- ejectResult(nil, rMsg, out)
- continue
- }
-
- i := unpackImage(rMsg, typ)
- if i == nil || i.Data == nil || i.Width <= 0 || i.Height <= 0 {
- ejectResult(nil, rMsg, out)
- continue
- }
-
- 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)
- continue
- }
-
- 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)
-
- /////////////////////////////////////
- sc++
- if sc == 25 {
- logo.Infoln("HUMAN TRACKER RUN 25 FRAME USE TIME: ", time.Since(tm))
- sc = 0
- tm = time.Now()
- }
-
- if time.Since(tm) > time.Second {
- logo.Infof("HUMAN TRACKER RUN %d FRAME USE TIME: %v", sc, time.Since(tm))
- sc = 0
- tm = time.Now()
- }
-
- }
- }
+ // 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 {
@@ -119,30 +227,48 @@
// 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)
+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,
}
- ret = append(ret, res)
+ 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 ret
+ 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
}
--
Gitblit v1.8.0