From f1c7582db3ce820ff1dc1cefa10fb3723ede1265 Mon Sep 17 00:00:00 2001
From: zhangmeng <775834166@qq.com>
Date: 星期二, 10 十二月 2019 11:14:49 +0800
Subject: [PATCH] update
---
goyolo.go | 112 ++++++++++++++++++++++++++++++++++++++++++++------------
1 files changed, 88 insertions(+), 24 deletions(-)
diff --git a/goyolo.go b/goyolo.go
index 1366e93..0c2687b 100644
--- a/goyolo.go
+++ b/goyolo.go
@@ -10,7 +10,13 @@
#include "cyolo.h"
*/
import "C"
-import "unsafe"
+import (
+ "fmt"
+ "unsafe"
+
+ "basic.com/pubsub/protomsg.git"
+ "github.com/gogo/protobuf/proto"
+)
// CPOINT pt
type CPOINT struct {
@@ -48,11 +54,15 @@
ID uint64
}
+type trackInfo struct {
+ lastTrackObjs []CObjTrackInfo
+ lastTrackID uint64
+}
+
// YoloHandle wrap C
type YoloHandle struct {
- handle unsafe.Pointer
- LastYoloObjs []CObjTrackInfo //yolo璺熻釜鐨勪笂涓�甯т俊鎭�
- LastTrackID uint64 //yolo 琚娇鐢ㄧ殑ID
+ handle unsafe.Pointer
+ tracker map[string]*trackInfo
}
// RatioInterTrack 璺熻釜鍒ゆ柇閲嶅彔闃堝��
@@ -79,7 +89,10 @@
g := C.int(gi)
p := C.init(c, w, n, g)
- return &YoloHandle{handle: p}
+ return &YoloHandle{
+ handle: p,
+ tracker: make(map[string]*trackInfo),
+ }
}
// Free free
@@ -104,15 +117,7 @@
}
// YoloDetect yolo detect
-func (y *YoloHandle) YoloDetect(img SDKImage, thrsh float32, umns int) []CObjInfo {
-
- data := img.Data
- w := img.Width
- h := img.Height
- c := img.Channel
- if c == 0 {
- c = 3
- }
+func (y *YoloHandle) YoloDetect(data []byte, w, h, c int, thrsh float32, umns int) []CObjInfo {
var count C.int
var cobjinfo unsafe.Pointer
@@ -176,14 +181,14 @@
return ratio
}
-// YoloDetectTrack yolo detect (鍙瘑鍒汉)
-func (y *YoloHandle) YoloDetectTrack(img SDKImage, thrsh float32, umns int) (allObjs []CObjTrackInfo, newObjs []CObjTrackInfo) {
+// YoloDetectTrack2 yolo detect (鍙瘑鍒汉)
+func (y *YoloHandle) YoloDetectTrack2(LastYoloObjs []CObjTrackInfo, LastTrackID *uint64, data []byte, w, h, c int, thrsh float32, umns int) (allObjs []CObjTrackInfo, newObjs []CObjTrackInfo) {
var tmp CObjTrackInfo
//LastYoloObjs
- detectObjs := y.YoloDetect(img, thrsh, umns)
+ detectObjs := y.YoloDetect(data, w, h, c, thrsh, umns)
- for _, vLast := range y.LastYoloObjs {
+ for _, vLast := range LastYoloObjs {
for i := 0; i < len(detectObjs); i++ {
//fmt.Println("vNew.Typ:", vNew.Typ)
if vLast.ObjInfo.Typ == detectObjs[i].Typ { //鍚屼竴绫诲埆锛屾瘮濡傞兘鏄汉浣�
@@ -203,19 +208,78 @@
}
//澶勭悊鏂板嚭鐜扮殑鐩爣
+ id := *LastTrackID
if len(detectObjs) > 0 {
for _, vAdd := range detectObjs {
tmp.ObjInfo = vAdd
- tmp.ID = y.LastTrackID
- y.LastTrackID++
+ tmp.ID = id
+ id++
allObjs = append(allObjs, tmp)
newObjs = append(newObjs, tmp)
}
}
-
- //鍒锋柊涓婁竴甯х殑璺熻釜鐩爣
- y.LastYoloObjs = allObjs
-
+ *LastTrackID = id
return allObjs, newObjs
}
+
+func convert2ProtoYoloTrack(obj []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
+}
+
+// YoloDetectTrack yolo detect (鍙瘑鍒汉)
+func (y *YoloHandle) YoloDetectTrack(id string, data []byte, w, h, c int, thrsh float32, umns int) ([]byte, []byte) {
+
+ v, ok := y.tracker[id]
+ if !ok {
+ i := &trackInfo{}
+ y.tracker[id] = i
+ v = i
+ }
+ whole, recent := y.YoloDetectTrack2(v.lastTrackObjs, &v.lastTrackID, data, w, h, c, thrsh, umns)
+ y.tracker[id].lastTrackObjs = whole
+ y.tracker[id].lastTrackID = v.lastTrackID
+
+ var dWhole, dRecent []byte
+ var err error
+ if len(whole) > 0 {
+
+ infos := convert2ProtoYoloTrack(whole, 1.0, 1.0)
+ p := protomsg.ParamYoloObj{Infos: infos}
+
+ dWhole, err = proto.Marshal(&p)
+ if err != nil {
+ fmt.Println("ydetect track marshal proto yolo obj error", err)
+ dWhole = nil
+ }
+ }
+ if len(recent) > 0 {
+ dRecent = nil
+ }
+ return dWhole, dRecent
+}
--
Gitblit v1.8.0