zhangmeng
2019-12-19 1118eaddcadb8b7c4a5da084a46c0c4fc18040e1
work/sdk/humantrack.go
@@ -13,6 +13,11 @@
   "basic.com/valib/gogpu.git"
)
type imageWithID struct {
   img *gohumantrack.ImageHumanTracker
   id  string
}
// HumanTracker track
type HumanTracker struct {
   tracker   *gohumantrack.HumanTracker
@@ -24,7 +29,7 @@
   mapCameraImageIndex map[string]int
   recvImageCount      int
   index               int
   images              []*gohumantrack.ImageHumanTracker
   images              []*imageWithID
   msgs                []*work.MsgRS
}
@@ -36,10 +41,11 @@
      flag:      flag,
      list:      NewLockList(6),
      recvImageCount: 0,
      index:          0,
      images:         make([]*gohumantrack.ImageHumanTracker, batchSize),
      msgs:           make([]*work.MsgRS, batchSize),
      recvImageCount:      0,
      index:               0,
      mapCameraImageIndex: make(map[string]int),
      images:              make([]*imageWithID, batchSize),
      msgs:                make([]*work.MsgRS, batchSize),
   }
}
@@ -52,10 +58,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 {
@@ -129,37 +132,44 @@
   // mapCameraImageIndex map[string]int
   // images              []gohumantrack.ImageHumanTracker
   t.recvImageCount++
   if t.mapCameraImageIndex == nil {
      t.mapCameraImageIndex = make(map[string]int)
      for i := 0; i < t.batchSize; i++ {
         t.images[i] = nil
      }
      for i := 0; i < t.batchSize; i++ {
         t.msgs[i] = nil
      }
   }
   if i, ok := t.mapCameraImageIndex[rMsg.Msg.Cid]; ok {
      if i < t.batchSize {
         t.images[i] = &img
         t.images[i] = &imageWithID{&img, rMsg.Msg.Cid}
         t.msgs[i] = &rMsg
      }
   } else {
      if t.index < t.batchSize {
         t.images[t.index] = &img
         t.images[t.index] = &imageWithID{&img, 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
   }
   res, err := t.tracker.ProcessImagePointer(t.images[:])
   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)
   t.recvImageCount = 0
   if err != nil {
      t.mapCameraImageIndex = nil
      logo.Infoln("batch~~~~~~Track Image Count: ", t.index, " Failed: ", err)
      ejectResult(nil, rMsg, out)
      return
   }
@@ -172,7 +182,7 @@
      result := protomsg.HumanTrackResult{Result: hr}
      data, err := proto.Marshal(&result)
      if err != nil {
         logo.Errorln("HUMAN TRACKER MARSHAL PROTO PLATE IDS ERROR", err)
         logo.Errorln("batch~~~~~~HUMAN TRACKER MARSHAL PROTO PLATE IDS ERROR", err)
         data = nil
      }
      msg := *t.msgs[i]
@@ -181,7 +191,7 @@
      if msg.Msg.Tasklab != nil {
         id, name = msg.Msg.Tasklab.Taskid, msg.Msg.Tasklab.Taskname
      }
      logo.Infoln("CAMERAID: ", msg.Msg.Cid, " TASKID: ", id, " TASKNAME: ", name, " Human Track COUNT: ", len(hr))
      logo.Infoln("batch~~~~~~CAMERAID: ", msg.Msg.Cid, " TASKID: ", id, " TASKNAME: ", name, " Human Track COUNT: ", len(hr))
   }