zhangmeng
2019-12-13 4ee39b10e09066bbcbf915f4ca251235fe397310
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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
package sdk
 
import (
    "analysis/logo"
    "analysis/work"
    "context"
    "time"
 
    "github.com/gogo/protobuf/proto"
 
    "analysis/gohumantrack"
 
    "basic.com/pubsub/protomsg.git"
    "basic.com/valib/gogpu.git"
)
 
// HumanTracker track
type HumanTracker struct {
    tracker   *gohumantrack.HumanTracker
    gpu       int
    batchSize int
    flag      int
}
 
// NewHumanTracker new
func NewHumanTracker(gpu, batchSize, flag int) *HumanTracker {
    return &HumanTracker{
        gpu:       gpu,
        batchSize: batchSize,
        flag:      flag,
    }
}
 
// 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)
    }
    logo.Infof("start gpu %d, batch %d, flag %d\n", gpu, t.batchSize, t.flag)
    t.tracker = gohumantrack.NewHumanTracker(gpu, t.batchSize, t.flag)
    if t.tracker != nil {
        logo.Infoln("Start Success")
        return true
    }
    logo.Infoln("Start Failed")
    return false
}
 
// 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 []gohumantrack.ImageHumanTracker
            img := gohumantrack.ImageHumanTracker{
                Data:    i.Data,
                Width:   imgW,
                Height:  imgH,
                Channel: 3,
            }
            images = append(images, img)
            res, err := t.tracker.Process(images)
            if err != 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 []gohumantrack.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
}