From 97c8be11558c7aff01ecdbf3d2a6470e2255a013 Mon Sep 17 00:00:00 2001
From: zhangmeng <775834166@qq.com>
Date: 星期五, 13 十二月 2019 12:25:24 +0800
Subject: [PATCH] add humantrack
---
work/sdk/humantrack.go | 190 ++++++++++++++++++++++++++++++++++++++
proc/proc.go | 41 +++++++-
go.mod | 2
proc/single.go | 9 -
work/common.go | 23 ----
5 files changed, 230 insertions(+), 35 deletions(-)
diff --git a/go.mod b/go.mod
index 0227b4f..d808405 100644
--- a/go.mod
+++ b/go.mod
@@ -6,7 +6,7 @@
basic.com/dbapi.git v0.0.0-20190803074805-04f05c8ca762
basic.com/libgowrapper/sdkstruct.git v0.0.0-20191211011351-89daaec8738e
basic.com/pubsub/cache.git v0.0.0-20190718093725-6a413e1d7d48
- basic.com/pubsub/protomsg.git v0.0.0-20191121033547-7af037c469a7
+ basic.com/pubsub/protomsg.git v0.0.0-20191213031824-8f12ebe28c26
basic.com/ruleForSdk.git v0.0.0-20190808095604-936797ef5da3
basic.com/valib/deliver.git v0.0.0-20190830083657-47adbbbb6651
basic.com/valib/godraw.git v0.0.0-20191122082247-26e9987cd183
diff --git a/proc/proc.go b/proc/proc.go
index 28411b4..7cbaa1f 100644
--- a/proc/proc.go
+++ b/proc/proc.go
@@ -6,6 +6,7 @@
"analysis/work"
"analysis/work/sdk"
"context"
+ "os"
)
type infoSDK struct {
@@ -35,6 +36,29 @@
name = "./data/yolo/data/coco.names"
)
+const (
+ FDetect = "FaceDetect"
+ // FExtract = "FaceExtract"
+ // FProperty = "FaceProperty"
+ FCompare = "FaceCompare"
+ FtTract = "FaceTrack"
+ FtDetect = "FaceTrackDetect"
+ FtOnly = "FaceTrackOnly"
+ YDetect = "Yolo"
+ PlateID = "Plate"
+ HumanTrack = "HumanTrack"
+)
+
+var (
+ SDK = []string{
+ YDetect,
+ FDetect,
+ FCompare,
+ PlateID,
+ HumanTrack,
+ }
+)
+
func prepare(name string, gpu int) (string, bool) {
return initSDK(name, gpu)
@@ -44,7 +68,7 @@
sdkType := name
- if s, f := util.FindStringInArray(sdkType, work.SDK); f {
+ if s, f := util.FindStringInArray(sdkType, SDK); f {
sdkType = s
} else {
logo.Errorln("THERE IS NO THIS SDK TYPE: ", sdkType)
@@ -64,23 +88,28 @@
var e sdk.Engine
ret := true
- if sdkType == work.FDetect {
+ if sdkType == FDetect {
e, ret = createFaceDetect(gpu, thrds, util.SFI.DetectNum, util.SFI.DetectAngle, util.SFI.TrackInterval, util.SFI.SampleSize)
- } else if sdkType == work.YDetect {
+ } else if sdkType == YDetect {
e, ret = createYolo(gpu)
- } else if sdkType == work.FtTract {
+ } else if sdkType == FtTract {
e, ret = createFaceTrack(gpu, thrds, util.SFI.DetectNum, util.SFI.DetectAngle, util.SFI.TrackInterval, util.SFI.SampleSize)
- } else if sdkType == work.FCompare {
+ } else if sdkType == FCompare {
e = sdk.NewEFExtract(100)
- } else if sdkType == work.PlateID {
+ } else if sdkType == PlateID {
e = sdk.NewVehicleDetector(util.SCI.LicenseServerPath, util.SCI.Model, util.SCI.MaxImageWidth, util.SCI.MaxImageHeight)
+ } else if sdkType == HumanTrack {
+ e = sdk.NewHumanTracker(gpu, 1, 0)
}
if ret {
mapFunc[sdkType] = e
+ } else {
+ logo.Errorln("CREATE SDK : ", sdkType, " FAILED")
+ os.Exit(0)
}
return sdkType, ret
}
diff --git a/proc/single.go b/proc/single.go
index 6262a28..2284ea2 100644
--- a/proc/single.go
+++ b/proc/single.go
@@ -3,7 +3,6 @@
import (
"analysis/logo"
"analysis/util"
- "analysis/work"
"context"
"strings"
)
@@ -58,12 +57,12 @@
logo.Errorf("NO THIS SDK PROC SDKPROC : ", typ)
}
rSDK := []string{
- work.FDetect,
- work.FCompare,
+ FDetect,
+ FCompare,
}
if typ == util.FeatYolo {
rSDK = rSDK[0:0]
- rSDK = append(rSDK, work.YDetect)
+ rSDK = append(rSDK, YDetect)
}
var res []bool
for _, v := range rSDK {
@@ -95,7 +94,7 @@
func allProc(ctx context.Context, gpu int, shm bool) bool {
var res []bool
- for _, v := range work.SDK {
+ for _, v := range SDK {
_, f := prepare(v, gpu)
res = append(res, f)
}
diff --git a/work/common.go b/work/common.go
index b646ce2..30a0653 100644
--- a/work/common.go
+++ b/work/common.go
@@ -6,29 +6,6 @@
"basic.com/pubsub/protomsg.git"
)
-const (
- FDetect = "FaceDetect"
- // FExtract = "FaceExtract"
- // FProperty = "FaceProperty"
- FCompare = "FaceCompare"
- FtTract = "FaceTrack"
- FtDetect = "FaceTrackDetect"
- FtOnly = "FaceTrackOnly"
- YDetect = "Yolo"
- PlateID = "Plate"
-)
-
-var (
- SDK = []string{
- YDetect,
- FDetect,
- FCompare,
- PlateID,
- // FtTract,
- // FtOnly,
- }
-)
-
const mode = deliver.PushPull
// MsgRS msg recv and snd
diff --git a/work/sdk/humantrack.go b/work/sdk/humantrack.go
new file mode 100644
index 0000000..e84b9ff
--- /dev/null
+++ b/work/sdk/humantrack.go
@@ -0,0 +1,190 @@
+package sdk
+
+import (
+ "analysis/logo"
+ "analysis/work"
+ "context"
+ "plugin"
+ "time"
+
+ "github.com/gogo/protobuf/proto"
+
+ "basic.com/libgowrapper/sdkstruct.git"
+ "basic.com/pubsub/protomsg.git"
+ "basic.com/valib/gogpu.git"
+)
+
+// HumanTracker track
+type HumanTracker struct {
+ gpu int
+ batchSize int
+ flag int
+
+ handle interface{}
+ fnInit func(int, int, int) interface{}
+ fnFree func(interface{})
+ fnRun func(interface{}, []byte, int, int, int) []sdkstruct.FgResult
+ fnProcess func(interface{}, []sdkstruct.SDKImage) ([]sdkstruct.FgResult, error)
+ fnVer func(interface{}) string
+ fnSimilarity func([]float32, []float32) (float64, error)
+}
+
+// NewHumanTracker new
+func NewHumanTracker(gpu, batchSize, flag int) *HumanTracker {
+ 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")
+ fnProc, _ := LoadFunc(plug, soFile, "Process")
+ fnVer, _ := LoadFunc(plug, soFile, "GetVersion")
+ fnSim, _ := LoadFunc(plug, soFile, "CFFSimilarity")
+
+ return &HumanTracker{
+ gpu: gpu,
+ batchSize: batchSize,
+ flag: flag,
+
+ handle: nil,
+ fnInit: fnInit.(func(int, int, int) interface{}),
+ fnFree: fnFree.(func(interface{})),
+ fnRun: fnRun.(func(interface{}, []byte, int, int, int) []sdkstruct.FgResult),
+ fnProcess: fnProc.(func(interface{}, []sdkstruct.SDKImage) ([]sdkstruct.FgResult, error)),
+ fnVer: fnVer.(func(interface{}) string),
+ fnSimilarity: fnSim.(func([]float32, []float32) (float64, error)),
+ }
+}
+
+// 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)
+ }
+ h := t.fnInit(gpu, t.batchSize, t.flag)
+ logo.Infoln("HumanTrack USE GPU: ", gpu)
+
+ if h == nil {
+ logo.Errorln("CREATE HumanTrack DETECTOR ERROR")
+ return false
+ }
+
+ t.handle = h
+ return true
+}
+
+// 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 []sdkstruct.SDKImage
+ // img := sdkstruct.SDKImage{
+ // Data: i.Data,
+ // Width: imgW,
+ // Height: imgH,
+ // Channel: 3,
+ // }
+ // images = append(images, img)
+ res := t.fnRun(t.handle, i.Data, imgW, imgH, 3)
+ if res != 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()
+ }
+
+ }
+ }
+}
+
+// message HumanTrack {
+// Rect rcHuman = 1;
+// float confidence = 2;
+// int32 x = 3;
+// int32 y = 4;
+// int32 id = 5;
+// repeated float feature = 6;
+// }
+
+// message HumanTrackResult {
+// repeated HumanTrack result = 1;
+// }
+
+func convert2ProtoHumanTrackResult(obj []sdkstruct.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)
+ }
+ ret = append(ret, res)
+ }
+ return ret
+}
--
Gitblit v1.8.0