udpate

zhangmeng

2019-12-17 00d07f58cd88a1e9dd13deed75ad1af3db6ebcf6

 work/sdk/ydetect.go

@@ -3,45 +3,57 @@
import (
   "analysis/logo"
   "analysis/work"
   "container/list"
   "context"
   "sync"
   "fmt"
   "plugin"

   "basic.com/libgowrapper/sdkstruct.git"
   "basic.com/pubsub/protomsg.git"
   "basic.com/valib/gogpu.git"

   "basic.com/valib/gosdk.git"
   "github.com/gogo/protobuf/proto"
)

type trackInfo struct {
   lastTrackObjs []gosdk.CObjTrackInfo
   lastTrackID   uint64
}

// YoloDetect yolo detect
type YoloDetect struct {
   yolo    *gosdk.YoloHandle
   iGPU    int
   cfg     string
   weights string
   name    string

   tracker map[string]*trackInfo
   list *LockList

   cache *list.List
   cv    *sync.Cond
   cond  bool
   handle interface{}
   fnInit func(string, string, string, int, func(...interface{})) interface{}
   fnFree func(interface{})
   fnRun  func(interface{}, string, []byte, int, int, int, float32, int) ([]sdkstruct.CObjTrackInfo, []sdkstruct.CObjTrackInfo)
}

// NewYDetectWithTrack with track
func NewYDetectWithTrack(gi int, cfg, weights, name string) *YoloDetect {
   soFile := "libyolo.so"

   plug, err := plugin.Open(soFile)
   if err != nil {
      logo.Errorln("Open: ", soFile, " error: ", err)
      return nil
   }

   fnInit, _ := LoadFunc(plug, soFile, "NewSDK")
   fnFree, _ := LoadFunc(plug, soFile, "Free")
   fnRun, _ := LoadFunc(plug, soFile, "Run")

   return &YoloDetect{
      iGPU:    gi,
      cfg:     cfg,
      weights: weights,
      name:    name,
      cache:   list.New(),
      cv:      sync.NewCond(&sync.Mutex{}),
      cond:    false,

      list: NewLockList(6),

      handle: nil,
      fnInit: fnInit.(func(string, string, string, int, func(...interface{})) interface{}),
      fnFree: fnFree.(func(interface{})),
      fnRun:  fnRun.(func(interface{}, string, []byte, int, int, int, float32, int) ([]sdkstruct.CObjTrackInfo, []sdkstruct.CObjTrackInfo)),
   }
}

@@ -52,19 +64,84 @@
   if gpu == -1 {
      gpu = gogpu.ValidGPU(2048)
   }
   yolo := gosdk.InitYolo(y.cfg, y.weights, y.name, gpu)
   logo.Infoln("yolo use gpu: ", gpu)
   h := y.fnInit(y.cfg, y.weights, y.name, gpu, logo.Infoln)
   logo.Infoln("YOLO USE GPU: ", gpu)

   if yolo == nil {
   if h == nil {
      logo.Errorln("CREATE YOLO DETECTOR ERROR")
      return false
   }

   y.yolo = yolo
   y.handle = h
   return true
}

// Run impl interface
func (y *YoloDetect) Run(ctx context.Context, in <-chan work.MsgRS, out chan<- work.MsgRS, typ string) {
   y.detectTrack(ctx, in, out, typ)
   FlowSimple(ctx, in, out, typ, y.list.Push, y.list.Pop, y.track, func() { y.fnFree(y.handle) })
}

func (y *YoloDetect) 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 {
      logo.Errorln("yolo image error: ", i)
      ejectResult(nil, rMsg, out)

      return
   }

   imgW, imgH := int(i.Width), int(i.Height)

   whole, recent := y.fnRun(y.handle, rMsg.Msg.Cid, i.Data, imgW, imgH, 3, 0.4, 0)
   if len(recent) > 0 {
   }

   infos := convert2ProtoYoloTrack(whole, 1.0, 1.0)
   p := protomsg.ParamYoloObj{Infos: infos}

   data, err := proto.Marshal(&p)
   if err != nil {
      fmt.Println("ydetect track marshal proto yolo obj 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, " DETECT YOLO COUNT: ", len(whole))
}

func convert2ProtoYoloTrack(obj []sdkstruct.CObjTrackInfo, fx, fy float64) []*protomsg.ObjInfo {
   ret := []*protomsg.ObjInfo{}

   for _, v := range obj {
      if fx < 1.0 || fy < 1.0 {
         v.ObjInfo.RcObj.Left = (int32)((float64)(v.ObjInfo.RcObj.Left) / fx)
         v.ObjInfo.RcObj.Right = (int32)((float64)(v.ObjInfo.RcObj.Right) / fx)
         v.ObjInfo.RcObj.Top = (int32)((float64)(v.ObjInfo.RcObj.Top) / fy)
         v.ObjInfo.RcObj.Bottom = (int32)((float64)(v.ObjInfo.RcObj.Bottom) / fy)
      }

      rect := protomsg.Rect{
         Left:   v.ObjInfo.RcObj.Left,
         Right:  v.ObjInfo.RcObj.Right,
         Top:    v.ObjInfo.RcObj.Top,
         Bottom: v.ObjInfo.RcObj.Bottom,
      }
      obj := protomsg.ObjInfo{
         RcObj: &rect,
         Typ:   v.ObjInfo.Typ,
         Prob:  v.ObjInfo.Prob,
         ObjID: v.ID,
      }

      ret = append(ret, &obj)
   }
   return ret
}