241 lines
7.3 KiB
C++
241 lines
7.3 KiB
C++
|
// Tencent is pleased to support the open source community by making ncnn available.
|
||
|
|
//
|
||
|
|
// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
|
||
|
|
//
|
||
|
|
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
|
||
|
|
// in compliance with the License. You may obtain a copy of the License at
|
||
|
|
//
|
||
|
|
// https://opensource.org/licenses/BSD-3-Clause
|
||
|
|
//
|
||
|
|
// Unless required by applicable law or agreed to in writing, software distributed
|
||
|
|
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
|
||
|
|
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
||
|
|
// specific language governing permissions and limitations under the License.
|
||
|
|
|
||
|
|
#include "net.h"
|
||
|
|
#if defined(USE_NCNN_SIMPLEOCV)
|
||
|
|
#include "simpleocv.h"
|
||
|
|
#else
|
||
|
|
#include <opencv2/core/core.hpp>
|
||
|
|
#include <opencv2/highgui/highgui.hpp>
|
||
|
|
#include <opencv2/imgproc/imgproc.hpp>
|
||
|
|
#endif
|
||
|
|
#include <stdlib.h>
|
||
|
|
#include <float.h>
|
||
|
|
#include <stdio.h>
|
||
|
|
#include <vector>
|
||
|
|
|
||
|
|
struct CrowdPoint
|
||
|
|
{
|
||
|
|
cv::Point pt;
|
||
|
|
float prob;
|
||
|
|
};
|
||
|
|
|
||
|
|
static void shift(int w, int h, int stride, std::vector<float> anchor_points, std::vector<float>& shifted_anchor_points)
|
||
|
|
{
|
||
|
|
std::vector<float> x_, y_;
|
||
|
|
for (int i = 0; i < w; i++)
|
||
|
|
{
|
||
|
|
float x = (i + 0.5) * stride;
|
||
|
|
x_.push_back(x);
|
||
|
|
}
|
||
|
|
for (int i = 0; i < h; i++)
|
||
|
|
{
|
||
|
|
float y = (i + 0.5) * stride;
|
||
|
|
y_.push_back(y);
|
||
|
|
}
|
||
|
|
|
||
|
|
std::vector<float> shift_x((size_t)w * h, 0), shift_y((size_t)w * h, 0);
|
||
|
|
for (int i = 0; i < h; i++)
|
||
|
|
{
|
||
|
|
for (int j = 0; j < w; j++)
|
||
|
|
{
|
||
|
|
shift_x[i * w + j] = x_[j];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for (int i = 0; i < h; i++)
|
||
|
|
{
|
||
|
|
for (int j = 0; j < w; j++)
|
||
|
|
{
|
||
|
|
shift_y[i * w + j] = y_[i];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
std::vector<float> shifts((size_t)w * h * 2, 0);
|
||
|
|
for (int i = 0; i < w * h; i++)
|
||
|
|
{
|
||
|
|
shifts[i * 2] = shift_x[i];
|
||
|
|
shifts[i * 2 + 1] = shift_y[i];
|
||
|
|
}
|
||
|
|
|
||
|
|
shifted_anchor_points.resize((size_t)2 * w * h * anchor_points.size() / 2, 0);
|
||
|
|
for (int i = 0; i < w * h; i++)
|
||
|
|
{
|
||
|
|
for (int j = 0; j < anchor_points.size() / 2; j++)
|
||
|
|
{
|
||
|
|
float x = anchor_points[j * 2] + shifts[i * 2];
|
||
|
|
float y = anchor_points[j * 2 + 1] + shifts[i * 2 + 1];
|
||
|
|
shifted_anchor_points[i * anchor_points.size() / 2 * 2 + j * 2] = x;
|
||
|
|
shifted_anchor_points[i * anchor_points.size() / 2 * 2 + j * 2 + 1] = y;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
static void generate_anchor_points(int stride, int row, int line, std::vector<float>& anchor_points)
|
||
|
|
{
|
||
|
|
float row_step = (float)stride / row;
|
||
|
|
float line_step = (float)stride / line;
|
||
|
|
|
||
|
|
std::vector<float> x_, y_;
|
||
|
|
for (int i = 1; i < line + 1; i++)
|
||
|
|
{
|
||
|
|
float x = (i - 0.5) * line_step - stride / 2;
|
||
|
|
x_.push_back(x);
|
||
|
|
}
|
||
|
|
for (int i = 1; i < row + 1; i++)
|
||
|
|
{
|
||
|
|
float y = (i - 0.5) * row_step - stride / 2;
|
||
|
|
y_.push_back(y);
|
||
|
|
}
|
||
|
|
std::vector<float> shift_x((size_t)row * line, 0), shift_y((size_t)row * line, 0);
|
||
|
|
for (int i = 0; i < row; i++)
|
||
|
|
{
|
||
|
|
for (int j = 0; j < line; j++)
|
||
|
|
{
|
||
|
|
shift_x[i * line + j] = x_[j];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for (int i = 0; i < row; i++)
|
||
|
|
{
|
||
|
|
for (int j = 0; j < line; j++)
|
||
|
|
{
|
||
|
|
shift_y[i * line + j] = y_[i];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
anchor_points.resize((size_t)row * line * 2, 0);
|
||
|
|
for (int i = 0; i < row * line; i++)
|
||
|
|
{
|
||
|
|
float x = shift_x[i];
|
||
|
|
float y = shift_y[i];
|
||
|
|
anchor_points[i * 2] = x;
|
||
|
|
anchor_points[i * 2 + 1] = y;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
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)
|
||
|
|
{
|
||
|
|
std::vector<std::pair<int, int> > image_shapes;
|
||
|
|
std::vector<int> strides;
|
||
|
|
for (int i = 0; i < pyramid_levels.size(); i++)
|
||
|
|
{
|
||
|
|
int new_h = std::floor((img_h + std::pow(2, pyramid_levels[i]) - 1) / std::pow(2, pyramid_levels[i]));
|
||
|
|
int new_w = std::floor((img_w + std::pow(2, pyramid_levels[i]) - 1) / std::pow(2, pyramid_levels[i]));
|
||
|
|
image_shapes.push_back(std::make_pair(new_w, new_h));
|
||
|
|
strides.push_back(std::pow(2, pyramid_levels[i]));
|
||
|
|
}
|
||
|
|
|
||
|
|
all_anchor_points.clear();
|
||
|
|
for (int i = 0; i < pyramid_levels.size(); i++)
|
||
|
|
{
|
||
|
|
std::vector<float> anchor_points;
|
||
|
|
generate_anchor_points(std::pow(2, pyramid_levels[i]), row, line, anchor_points);
|
||
|
|
std::vector<float> shifted_anchor_points;
|
||
|
|
shift(image_shapes[i].first, image_shapes[i].second, strides[i], anchor_points, shifted_anchor_points);
|
||
|
|
all_anchor_points.insert(all_anchor_points.end(), shifted_anchor_points.begin(), shifted_anchor_points.end());
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
static int detect_crowd(const cv::Mat& bgr, std::vector<CrowdPoint>& crowd_points)
|
||
|
|
{
|
||
|
|
ncnn::Option opt;
|
||
|
|
opt.num_threads = 4;
|
||
|
|
opt.use_vulkan_compute = false;
|
||
|
|
opt.use_bf16_storage = false;
|
||
|
|
|
||
|
|
ncnn::Net net;
|
||
|
|
net.opt = opt;
|
||
|
|
|
||
|
|
// model is converted from
|
||
|
|
// https://github.com/TencentYoutuResearch/CrowdCounting-P2PNet
|
||
|
|
// the ncnn model https://pan.baidu.com/s/1O1CBgvY6yJkrK8Npxx3VMg pwd: ezhx
|
||
|
|
net.load_param("p2pnet.param");
|
||
|
|
net.load_model("p2pnet.bin");
|
||
|
|
|
||
|
|
int width = bgr.cols;
|
||
|
|
int height = bgr.rows;
|
||
|
|
|
||
|
|
int new_width = width / 128 * 128;
|
||
|
|
int new_height = height / 128 * 128;
|
||
|
|
|
||
|
|
ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, width, height, new_width, new_height);
|
||
|
|
|
||
|
|
std::vector<int> pyramid_levels(1, 3);
|
||
|
|
std::vector<float> all_anchor_points;
|
||
|
|
generate_anchor_points(in.w, in.h, pyramid_levels, 2, 2, all_anchor_points);
|
||
|
|
|
||
|
|
ncnn::Mat anchor_points = ncnn::Mat(2, all_anchor_points.size() / 2, all_anchor_points.data());
|
||
|
|
|
||
|
|
ncnn::Extractor ex = net.create_extractor();
|
||
|
|
const float mean_vals1[3] = {123.675f, 116.28f, 103.53f};
|
||
|
|
const float norm_vals1[3] = {0.01712475f, 0.0175f, 0.01742919f};
|
||
|
|
|
||
|
|
in.substract_mean_normalize(mean_vals1, norm_vals1);
|
||
|
|
|
||
|
|
ex.input("input", in);
|
||
|
|
ex.input("anchor", anchor_points);
|
||
|
|
|
||
|
|
ncnn::Mat score, points;
|
||
|
|
ex.extract("pred_scores", score);
|
||
|
|
ex.extract("pred_points", points);
|
||
|
|
|
||
|
|
for (int i = 0; i < points.h; i++)
|
||
|
|
{
|
||
|
|
float* score_data = score.row(i);
|
||
|
|
float* points_data = points.row(i);
|
||
|
|
CrowdPoint cp;
|
||
|
|
int x = points_data[0] / new_width * width;
|
||
|
|
int y = points_data[1] / new_height * height;
|
||
|
|
cp.pt = cv::Point(x, y);
|
||
|
|
cp.prob = score_data[1];
|
||
|
|
crowd_points.push_back(cp);
|
||
|
|
}
|
||
|
|
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
static void draw_result(const cv::Mat& bgr, const std::vector<CrowdPoint>& crowd_points)
|
||
|
|
{
|
||
|
|
cv::Mat image = bgr.clone();
|
||
|
|
const float threshold = 0.5f;
|
||
|
|
for (int i = 0; i < crowd_points.size(); i++)
|
||
|
|
{
|
||
|
|
if (crowd_points[i].prob > threshold)
|
||
|
|
{
|
||
|
|
cv::circle(image, crowd_points[i].pt, 4, cv::Scalar(0, 0, 255), -1, 8, 0);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
cv::imshow("image", image);
|
||
|
|
cv::waitKey();
|
||
|
|
}
|
||
|
|
int main(int argc, char** argv)
|
||
|
|
{
|
||
|
|
if (argc != 2)
|
||
|
|
{
|
||
|
|
fprintf(stderr, "Usage: %s [imagepath]\n", argv[0]);
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
|
||
|
|
const char* imagepath = argv[1];
|
||
|
|
|
||
|
|
cv::Mat bgr = cv::imread(imagepath, 1);
|
||
|
|
if (bgr.empty())
|
||
|
|
{
|
||
|
|
fprintf(stderr, "cv::imread %s failed\n", imagepath);
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
|
||
|
|
std::vector<CrowdPoint> crowd_points;
|
||
|
|
detect_crowd(bgr, crowd_points);
|
||
|
|
draw_result(bgr, crowd_points);
|
||
|
|
|
||
|
|
return 0;
|
||
|
|
}
|