/
nn_matching.h
76 lines (59 loc) · 1.81 KB
/
nn_matching.h
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
#ifndef NN_MATCHING_H
#define NN_MATCHING_H
#include <torch/torch.h>
#include <map>
#include <vector>
#include <opencv2/opencv.hpp>
torch::Tensor iou_dist(const std::vector<cv::Rect2f> &dets, const std::vector<cv::Rect2f> &trks);
// save features of the track in GPU
class FeatureBundle {
public:
FeatureBundle() : full(false), next(0), store(torch::empty({budget, feat_dim}).cuda()) {}
void clear() {
next = 0;
full = false;
}
bool empty() const {
return next == 0 && !full;
}
void add(torch::Tensor feat) {
if (next == budget) {
full = true;
next = 0;
}
store[next++] = feat;
}
torch::Tensor get() const {
return full ? store : store.slice(0, 0, next);
}
private:
static const int64_t budget = 100, feat_dim = 512;
torch::Tensor store;
bool full;
int64_t next;
};
template<typename TrackData>
class FeatureMetric {
public:
explicit FeatureMetric(std::vector<TrackData> &data) : data(data) {}
torch::Tensor distance(torch::Tensor features, const std::vector<int> &targets) {
auto dist = torch::empty({int64_t(targets.size()), features.size(0)});
if (features.size(0)) {
for (size_t i = 0; i < targets.size(); ++i) {
dist[i] = nn_cosine_distance(data[targets[i]].feats.get(), features);
}
}
return dist;
}
void update(torch::Tensor feats, const std::vector<int> &targets) {
for (size_t i = 0; i < targets.size(); ++i) {
data[targets[i]].feats.add(feats[i]);
}
}
private:
std::vector<TrackData> &data;
torch::Tensor nn_cosine_distance(torch::Tensor x, torch::Tensor y) {
return std::get<0>(torch::min(1 - torch::matmul(x, y.t()), 0)).cpu();
}
};
#endif //NN_MATCHING_H