update

zhangmeng

2019-12-10 f1c7582db3ce820ff1dc1cefa10fb3723ede1265

update

 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
}