zhangmeng
2019-12-19 1118eaddcadb8b7c4a5da084a46c0c4fc18040e1
work/sdk/humantrack.go
@@ -4,14 +4,19 @@
   "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
   id  string
}
// HumanTracker track
type HumanTracker struct {
@@ -19,6 +24,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 +39,168 @@
      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)
   img := gohumantrack.ImageHumanTracker{
      Data:    i.Data,
      Width:   imgW,
      Height:  imgH,
      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, rMsg.Msg.Cid}
         t.msgs[i] = &rMsg
      }
   } else {
      if t.index < t.batchSize {
         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
   }
   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 {
      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 := convert2ProtoHumanTrackResult(res[i])
      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 +216,25 @@
//     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
}