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 |  192 ++++++++++++++++++++++++++++++++++++++++--------
 1 files changed, 160 insertions(+), 32 deletions(-)

diff --git a/work/sdk/humantrack.go b/work/sdk/humantrack.go
index 99e9de0..46ad033 100644
--- a/work/sdk/humantrack.go
+++ b/work/sdk/humantrack.go
@@ -1,6 +1,7 @@
 package sdk
 
 import (
+	"analysis/goconv"
 	"analysis/logo"
 	"analysis/work"
 	"context"
@@ -13,6 +14,13 @@
 	"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
@@ -20,6 +28,12 @@
 	batchSize int
 	flag      int
 	list      *LockList
+
+	mapCameraImageIndex map[string]int
+	recvImageCount      int
+	index               int
+	images              []*imageWithID
+	msgs                []*work.MsgRS
 }
 
 // NewHumanTracker new
@@ -29,6 +43,12 @@
 		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),
 	}
 }
 
@@ -41,10 +61,7 @@
 
 // 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 {
@@ -83,8 +100,8 @@
 		return
 	}
 
-	hr := convert2ProtoHumanTrackResult(res)
-	result := protomsg.HumanTrackResult{Result: hr[0]}
+	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)
@@ -95,13 +112,106 @@
 	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]))
+	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) {
-	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.track, t.Free)
+	FlowSimple(ctx, in, out, typ, t.list.Push, t.list.Pop, t.trackBatch, t.Free)
 }
 
 // message HumanTrack {
@@ -117,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