zhangmeng
2019-12-17 39fbaaff56ec2b6d2135fccbdfe337a8591eff03
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package sdk
 
import (
    "analysis/logo"
    "analysis/work"
    "context"
 
    "basic.com/pubsub/protomsg.git"
    "basic.com/valib/gosdk.git"
    "github.com/gogo/protobuf/proto"
)
 
func (y *YoloDetect) detectTrack(ctx context.Context, in <-chan work.MsgRS, out chan<- work.MsgRS, typ string) {
    y.tracker = make(map[string]*trackInfo)
    FlowSimple(ctx, in, out, typ, y.list.Push, y.list.Pop, y.track, func() { gosdk.Free() })
}
 
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
    }
 
    // conv to bgr24 and resize
    imgW, imgH := int(i.Width), int(i.Height)
 
    img := gosdk.SDKImage{Data: i.Data, Width: imgW, Height: imgH}
 
    v, ok := y.tracker[rMsg.Msg.Cid]
    if !ok {
        i := &trackInfo{}
        y.tracker[rMsg.Msg.Cid] = i
        v = i
    }
    allOrig, _ := gosdk.YoloDetectTrack2(y.yolo, v.lastTrackObjs, &v.lastTrackID, img, 0.4, 0)
    y.tracker[rMsg.Msg.Cid].lastTrackObjs = allOrig
    y.tracker[rMsg.Msg.Cid].lastTrackID = v.lastTrackID
 
    /// filter rules
    var allF []gosdk.CObjTrackInfo
    sdkid := rMsg.Msg.Tasklab.Sdkinfos[rMsg.Msg.Tasklab.Index].Ipcid
    for _, v := range allOrig {
 
        size := (v.ObjInfo.RcObj.Right - v.ObjInfo.RcObj.Left) * (v.ObjInfo.RcObj.Bottom - v.ObjInfo.RcObj.Top)
        if filter(rMsg.Msg.Tasklab.Taskid, sdkid, v.ObjInfo.Prob, 0, int(size)) {
            allF = append(allF, v)
        }
    }
    /// filter rules
    allO := allF
    // allO := allOrig
 
    var data []byte
    var err error
    if len(allO) > 0 {
 
        infos := convert2ProtoYoloTrack(allO, 1.0, 1.0)
        p := protomsg.ParamYoloObj{Infos: infos}
        // logo.Infoln("CID: ", rMsg.Msg.Cid, " TASK: ", rMsg.Msg.Tasklab.Taskid, " YOLO TRACK OBJS: ", len(infos))
 
        data, err = proto.Marshal(&p)
        if err != nil {
            logo.Errorln("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(allO))
}
 
func convert2ProtoYoloTrack(obj []gosdk.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
}