ncnn

p2pnet.cpp
242 строки · 7.3 Кб
Перенос по словам
1
// Tencent is pleased to support the open source community by making ncnn available.
2
//
3
// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
4
//
5
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
6
// in compliance with the License. You may obtain a copy of the License at
7
//
8
// https://opensource.org/licenses/BSD-3-Clause
9
//
10
// Unless required by applicable law or agreed to in writing, software distributed
11
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
12
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
13
// specific language governing permissions and limitations under the License.
14

15
#include "net.h"
16
#if defined(USE_NCNN_SIMPLEOCV)
17
#include "simpleocv.h"
18
#else
19
#include <opencv2/core/core.hpp>
20
#include <opencv2/highgui/highgui.hpp>
21
#include <opencv2/imgproc/imgproc.hpp>
22
#endif
23
#include <stdlib.h>
24
#include <float.h>
25
#include <stdio.h>
26
#include <vector>
27

28
struct CrowdPoint
29
{
30
    cv::Point pt;
31
    float prob;
32
};
33

34
static void shift(int w, int h, int stride, std::vector<float> anchor_points, std::vector<float>& shifted_anchor_points)
35
{
36
    std::vector<float> x_, y_;
37
    for (int i = 0; i < w; i++)
38
    {
39
        float x = (i + 0.5) * stride;
40
        x_.push_back(x);
41
    }
42
    for (int i = 0; i < h; i++)
43
    {
44
        float y = (i + 0.5) * stride;
45
        y_.push_back(y);
46
    }
47

48
    std::vector<float> shift_x((size_t)w * h, 0), shift_y((size_t)w * h, 0);
49
    for (int i = 0; i < h; i++)
50
    {
51
        for (int j = 0; j < w; j++)
52
        {
53
            shift_x[i * w + j] = x_[j];
54
        }
55
    }
56
    for (int i = 0; i < h; i++)
57
    {
58
        for (int j = 0; j < w; j++)
59
        {
60
            shift_y[i * w + j] = y_[i];
61
        }
62
    }
63

64
    std::vector<float> shifts((size_t)w * h * 2, 0);
65
    for (int i = 0; i < w * h; i++)
66
    {
67
        shifts[i * 2] = shift_x[i];
68
        shifts[i * 2 + 1] = shift_y[i];
69
    }
70

71
    shifted_anchor_points.resize((size_t)2 * w * h * anchor_points.size() / 2, 0);
72
    for (int i = 0; i < w * h; i++)
73
    {
74
        for (int j = 0; j < anchor_points.size() / 2; j++)
75
        {
76
            float x = anchor_points[j * 2] + shifts[i * 2];
77
            float y = anchor_points[j * 2 + 1] + shifts[i * 2 + 1];
78
            shifted_anchor_points[i * anchor_points.size() / 2 * 2 + j * 2] = x;
79
            shifted_anchor_points[i * anchor_points.size() / 2 * 2 + j * 2 + 1] = y;
80
        }
81
    }
82
}
83
static void generate_anchor_points(int stride, int row, int line, std::vector<float>& anchor_points)
84
{
85
    float row_step = (float)stride / row;
86
    float line_step = (float)stride / line;
87

88
    std::vector<float> x_, y_;
89
    for (int i = 1; i < line + 1; i++)
90
    {
91
        float x = (i - 0.5) * line_step - stride / 2;
92
        x_.push_back(x);
93
    }
94
    for (int i = 1; i < row + 1; i++)
95
    {
96
        float y = (i - 0.5) * row_step - stride / 2;
97
        y_.push_back(y);
98
    }
99
    std::vector<float> shift_x((size_t)row * line, 0), shift_y((size_t)row * line, 0);
100
    for (int i = 0; i < row; i++)
101
    {
102
        for (int j = 0; j < line; j++)
103
        {
104
            shift_x[i * line + j] = x_[j];
105
        }
106
    }
107
    for (int i = 0; i < row; i++)
108
    {
109
        for (int j = 0; j < line; j++)
110
        {
111
            shift_y[i * line + j] = y_[i];
112
        }
113
    }
114
    anchor_points.resize((size_t)row * line * 2, 0);
115
    for (int i = 0; i < row * line; i++)
116
    {
117
        float x = shift_x[i];
118
        float y = shift_y[i];
119
        anchor_points[i * 2] = x;
120
        anchor_points[i * 2 + 1] = y;
121
    }
122
}
123
static void generate_anchor_points(int img_w, int img_h, std::vector<int> pyramid_levels, int row, int line, std::vector<float>& all_anchor_points)
124
{
125
    std::vector<std::pair<int, int> > image_shapes;
126
    std::vector<int> strides;
127
    for (int i = 0; i < pyramid_levels.size(); i++)
128
    {
129
        int new_h = std::floor((img_h + std::pow(2, pyramid_levels[i]) - 1) / std::pow(2, pyramid_levels[i]));
130
        int new_w = std::floor((img_w + std::pow(2, pyramid_levels[i]) - 1) / std::pow(2, pyramid_levels[i]));
131
        image_shapes.push_back(std::make_pair(new_w, new_h));
132
        strides.push_back(std::pow(2, pyramid_levels[i]));
133
    }
134

135
    all_anchor_points.clear();
136
    for (int i = 0; i < pyramid_levels.size(); i++)
137
    {
138
        std::vector<float> anchor_points;
139
        generate_anchor_points(std::pow(2, pyramid_levels[i]), row, line, anchor_points);
140
        std::vector<float> shifted_anchor_points;
141
        shift(image_shapes[i].first, image_shapes[i].second, strides[i], anchor_points, shifted_anchor_points);
142
        all_anchor_points.insert(all_anchor_points.end(), shifted_anchor_points.begin(), shifted_anchor_points.end());
143
    }
144
}
145

146
static int detect_crowd(const cv::Mat& bgr, std::vector<CrowdPoint>& crowd_points)
147
{
148
    ncnn::Option opt;
149
    opt.num_threads = 4;
150
    opt.use_vulkan_compute = false;
151
    opt.use_bf16_storage = false;
152

153
    ncnn::Net net;
154
    net.opt = opt;
155

156
    // model is converted from
157
    // https://github.com/TencentYoutuResearch/CrowdCounting-P2PNet
158
    // the ncnn model  https://pan.baidu.com/s/1O1CBgvY6yJkrK8Npxx3VMg pwd: ezhx
159
    if (net.load_param("p2pnet.param"))
160
        exit(-1);
161
    if (net.load_model("p2pnet.bin"))
162
        exit(-1);
163

164
    int width = bgr.cols;
165
    int height = bgr.rows;
166

167
    int new_width = width / 128 * 128;
168
    int new_height = height / 128 * 128;
169

170
    ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, width, height, new_width, new_height);
171

172
    std::vector<int> pyramid_levels(1, 3);
173
    std::vector<float> all_anchor_points;
174
    generate_anchor_points(in.w, in.h, pyramid_levels, 2, 2, all_anchor_points);
175

176
    ncnn::Mat anchor_points = ncnn::Mat(2, all_anchor_points.size() / 2, all_anchor_points.data());
177

178
    ncnn::Extractor ex = net.create_extractor();
179
    const float mean_vals1[3] = {123.675f, 116.28f, 103.53f};
180
    const float norm_vals1[3] = {0.01712475f, 0.0175f, 0.01742919f};
181

182
    in.substract_mean_normalize(mean_vals1, norm_vals1);
183

184
    ex.input("input", in);
185
    ex.input("anchor", anchor_points);
186

187
    ncnn::Mat score, points;
188
    ex.extract("pred_scores", score);
189
    ex.extract("pred_points", points);
190

191
    for (int i = 0; i < points.h; i++)
192
    {
193
        float* score_data = score.row(i);
194
        float* points_data = points.row(i);
195
        CrowdPoint cp;
196
        int x = points_data[0] / new_width * width;
197
        int y = points_data[1] / new_height * height;
198
        cp.pt = cv::Point(x, y);
199
        cp.prob = score_data[1];
200
        crowd_points.push_back(cp);
201
    }
202

203
    return 0;
204
}
205

206
static void draw_result(const cv::Mat& bgr, const std::vector<CrowdPoint>& crowd_points)
207
{
208
    cv::Mat image = bgr.clone();
209
    const float threshold = 0.5f;
210
    for (int i = 0; i < crowd_points.size(); i++)
211
    {
212
        if (crowd_points[i].prob > threshold)
213
        {
214
            cv::circle(image, crowd_points[i].pt, 4, cv::Scalar(0, 0, 255), -1, 8, 0);
215
        }
216
    }
217
    cv::imshow("image", image);
218
    cv::waitKey();
219
}
220
int main(int argc, char** argv)
221
{
222
    if (argc != 2)
223
    {
224
        fprintf(stderr, "Usage: %s [imagepath]\n", argv[0]);
225
        return -1;
226
    }
227

228
    const char* imagepath = argv[1];
229

230
    cv::Mat bgr = cv::imread(imagepath, 1);
231
    if (bgr.empty())
232
    {
233
        fprintf(stderr, "cv::imread %s failed\n", imagepath);
234
        return -1;
235
    }
236

237
    std::vector<CrowdPoint> crowd_points;
238
    detect_crowd(bgr, crowd_points);
239
    draw_result(bgr, crowd_points);
240

241
    return 0;
242
}
243
ncnn

Использование cookies
