mirrors
/
mmdeploy
mirror of https://github.com/open-mmlab/mmdeploy.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
mmdeploy/csrc/codebase/mmpose/keypoints_from_heatmap.cpp

391 lines
15 KiB
C++
Raw Blame History

// Copyright (c) OpenMMLab. All rights reserved.
#include <cctype>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include "core/device.h"
#include "core/registry.h"
#include "core/serialization.h"
#include "core/tensor.h"
#include "core/utils/device_utils.h"
#include "core/utils/formatter.h"
#include "core/value.h"
#include "experimental/module_adapter.h"
#include "mmpose.h"
#include "opencv_utils.h"
namespace mmdeploy::mmpose {
using std::string;
using std::vector;
template <class F>
struct _LoopBody : public cv::ParallelLoopBody {
F f_;
_LoopBody(F f) : f_(std::move(f)) {}
void operator()(const cv::Range& range) const override { f_(range); }
};
std::string to_lower(const std::string& s) {
std::string t = s;
std::transform(t.begin(), t.end(), t.begin(), [](unsigned char c) { return std::tolower(c); });
return t;
}
class TopdownHeatmapBaseHeadDecode : public MMPose {
public:
explicit TopdownHeatmapBaseHeadDecode(const Value& config) : MMPose(config) {
if (config.contains("params")) {
auto& params = config["params"];
flip_test_ = params.value("flip_test", flip_test_);
use_udp_ = params.value("use_udp", use_udp_);
target_type_ = params.value("target_type", target_type_);
valid_radius_factor_ = params.value("valid_radius_factor", valid_radius_factor_);
unbiased_decoding_ = params.value("unbiased_decoding", unbiased_decoding_);
post_process_ = params.value("post_process", post_process_);
shift_heatmap_ = params.value("shift_heatmap", shift_heatmap_);
modulate_kernel_ = params.value("modulate_kernel", modulate_kernel_);
}
}
Result<Value> operator()(const Value& _data, const Value& _prob) {
MMDEPLOY_DEBUG("preprocess_result: {}", _data);
MMDEPLOY_DEBUG("inference_result: {}", _prob);
Device cpu_device{"cpu"};
OUTCOME_TRY(auto heatmap,
MakeAvailableOnDevice(_prob["output"].get<Tensor>(), cpu_device, stream()));
OUTCOME_TRY(stream().Wait());
if (!(heatmap.shape().size() == 4 && heatmap.data_type() == DataType::kFLOAT)) {
MMDEPLOY_ERROR("unsupported `output` tensor, shape: {}, dtype: {}", heatmap.shape(),
(int)heatmap.data_type());
return Status(eNotSupported);
}
auto& img_metas = _data["img_metas"];
vector<float> center;
vector<float> scale;
from_value(img_metas["center"], center);
from_value(img_metas["scale"], scale);
Tensor pred =
keypoints_from_heatmap(heatmap, center, scale, unbiased_decoding_, post_process_,
modulate_kernel_, valid_radius_factor_, use_udp_, target_type_);
return GetOutput(pred);
}
Value GetOutput(Tensor& pred) {
PoseDetectorOutput output;
int K = pred.shape(1);
float* data = pred.data<float>();
for (int i = 0; i < K; i++) {
float x = *(data + 0);
float y = *(data + 1);
float s = *(data + 2);
output.key_points.push_back({{x, y}, s});
data += 3;
}
return to_value(std::move(output));
}
Tensor keypoints_from_heatmap(const Tensor& _heatmap, const vector<float>& center,
const vector<float>& scale, bool unbiased_decoding,
const string& post_process, int modulate_kernel,
float valid_radius_factor, bool use_udp,
const string& target_type) {
Tensor heatmap(_heatmap.desc());
heatmap.CopyFrom(_heatmap, stream()).value();
stream().Wait().value();
int K = heatmap.shape(1);
int H = heatmap.shape(2);
int W = heatmap.shape(3);
if (post_process == "megvii") {
heatmap = gaussian_blur(heatmap, modulate_kernel);
}
Tensor pred;
if (use_udp) {
if (to_lower(target_type) == to_lower(string("GaussianHeatMap"))) {
pred = get_max_pred(heatmap);
post_dark_udp(pred, heatmap, modulate_kernel);
} else if (to_lower(target_type) == to_lower(string("CombinedTarget"))) {
// output channel = 3 * channel_cfg['num_output_channels']
assert(K % 3 == 0);
cv::parallel_for_(cv::Range(0, K), _LoopBody{[&](const cv::Range& r) {
for (int i = r.start; i < r.end; i++) {
int kt = (i % 3 == 0) ? 2 * modulate_kernel + 1 : modulate_kernel;
float* data = heatmap.data<float>() + i * H * W;
cv::Mat work = cv::Mat(H, W, CV_32FC(1), data);
cv::GaussianBlur(work, work, {kt, kt}, 0); // inplace
}
}});
float valid_radius = valid_radius_factor_ * H;
TensorDesc desc = {Device{"cpu"}, DataType::kFLOAT, {1, K / 3, H, W}};
Tensor offset_x(desc);
Tensor offset_y(desc);
Tensor heatmap_(desc);
{
// split heatmap
float* src = heatmap.data<float>();
float* dst0 = heatmap_.data<float>();
float* dst1 = offset_x.data<float>();
float* dst2 = offset_y.data<float>();
for (int i = 0; i < K / 3; i++) {
std::copy_n(src, H * W, dst0);
std::transform(src + H * W, src + 2 * H * W, dst1,
[=](float& x) { return x * valid_radius; });
std::transform(src + 2 * H * W, src + 3 * H * W, dst2,
[=](float& x) { return x * valid_radius; });
src += 3 * H * W;
dst0 += H * W;
dst1 += H * W;
dst2 += H * W;
}
}
pred = get_max_pred(heatmap_);
for (int i = 0; i < K / 3; i++) {
float* data = pred.data<float>() + i * 3;
int index = *(data + 0) + *(data + 1) * W + H * W * i;
float* offx = offset_x.data<float>() + index;
float* offy = offset_y.data<float>() + index;
*(data + 0) += *offx;
*(data + 1) += *offy;
}
}
} else {
pred = get_max_pred(heatmap);
if (post_process == "unbiased") {
heatmap = gaussian_blur(heatmap, modulate_kernel);
float* data = heatmap.data<float>();
std::for_each(data, data + K * H * W, [](float& v) {
double _v = std::max((double)v, 1e-10);
v = std::log(_v);
});
for (int i = 0; i < K; i++) {
taylor(heatmap, pred, i);
}
} else if (post_process != "null") {
for (int i = 0; i < K; i++) {
float* data = heatmap.data<float>() + i * W * H;
auto _data = [&](int y, int x) { return *(data + y * W + x); };
int px = *(pred.data<float>() + i * 3 + 0);
int py = *(pred.data<float>() + i * 3 + 1);
if (1 < px && px < W - 1 && 1 < py && py < H - 1) {
float v1 = _data(py, px + 1) - _data(py, px - 1);
float v2 = _data(py + 1, px) - _data(py - 1, px);
*(pred.data<float>() + i * 3 + 0) += (v1 > 0) ? 0.25 : ((v1 < 0) ? -0.25 : 0);
*(pred.data<float>() + i * 3 + 1) += (v2 > 0) ? 0.25 : ((v2 < 0) ? -0.25 : 0);
if (post_process_ == "megvii") {
*(pred.data<float>() + i * 3 + 0) += 0.5;
*(pred.data<float>() + i * 3 + 1) += 0.5;
}
}
}
}
}
K = pred.shape(1); // changed if target_type is CombinedTarget
// Transform back to the image
for (int i = 0; i < K; i++) {
transform_pred(pred, i, center, scale, {W, H}, use_udp);
}
if (post_process_ == "megvii") {
for (int i = 0; i < K; i++) {
float* data = pred.data<float>() + i * 3 + 2;
*data = *data / 255.0 + 0.5;
}
}
return pred;
}
void post_dark_udp(Tensor& pred, Tensor& heatmap, int kernel) {
int K = heatmap.shape(1);
int H = heatmap.shape(2);
int W = heatmap.shape(3);
cv::parallel_for_(cv::Range(0, K), _LoopBody{[&](const cv::Range& r) {
for (int i = r.start; i < r.end; i++) {
float* data = heatmap.data<float>() + i * H * W;
cv::Mat work = cv::Mat(H, W, CV_32FC(1), data);
cv::GaussianBlur(work, work, {kernel, kernel}, 0); // inplace
}
}});
std::for_each(heatmap.data<float>(), heatmap.data<float>() + K * H * W, [](float& x) {
x = std::max(0.001f, std::min(50.f, x));
x = std::log(x);
});
auto _heatmap_data = [&](int index, int c) -> float {
int y = index / (W + 2);
int x = index % (W + 2);
y = std::max(0, y - 1);
x = std::max(0, x - 1);
return *(heatmap.data<float>() + c * H * W + y * W + x);
};
for (int i = 0; i < K; i++) {
float* data = pred.data<float>() + i * 3;
int index = *(data + 0) + 1 + (*(data + 1) + 1) * (W + 2);
float i_ = _heatmap_data(index, i);
float ix1 = _heatmap_data(index + 1, i);
float iy1 = _heatmap_data(index + W + 2, i);
float ix1y1 = _heatmap_data(index + W + 3, i);
float ix1_y1_ = _heatmap_data(index - W - 3, i);
float ix1_ = _heatmap_data(index - 1, i);
float iy1_ = _heatmap_data(index - 2 - W, i);
float dx = 0.5 * (ix1 - ix1_);
float dy = 0.5 * (iy1 - iy1_);
float dxx = ix1 - 2 * i_ + ix1_;
float dyy = iy1 - 2 * i_ + iy1_;
float dxy = 0.5 * (ix1y1 - ix1 - iy1 + i_ + i_ - ix1_ - iy1_ + ix1_y1_);
vector<float> _data0 = {dx, dy};
vector<float> _data1 = {dxx, dxy, dxy, dyy};
cv::Mat derivative = cv::Mat(2, 1, CV_32FC1, _data0.data());
cv::Mat hessian = cv::Mat(2, 2, CV_32FC1, _data1.data());
cv::Mat hessianinv = hessian.inv();
cv::Mat offset = -hessianinv * derivative;
*(data + 0) += offset.at<float>(0, 0);
*(data + 1) += offset.at<float>(1, 0);
}
}
void transform_pred(Tensor& pred, int k, const vector<float>& center, const vector<float>& _scale,
const vector<int>& output_size, bool use_udp = false) {
auto scale = _scale;
scale[0] *= 200;
scale[1] *= 200;
float scale_x, scale_y;
if (use_udp) {
scale_x = scale[0] / (output_size[0] - 1.0);
scale_y = scale[1] / (output_size[1] - 1.0);
} else {
scale_x = scale[0] / output_size[0];
scale_y = scale[1] / output_size[1];
}
float* data = pred.data<float>() + k * 3;
*(data + 0) = *(data + 0) * scale_x + center[0] - scale[0] * 0.5;
*(data + 1) = *(data + 1) * scale_y + center[1] - scale[1] * 0.5;
}
void taylor(const Tensor& heatmap, Tensor& pred, int k) {
int K = heatmap.shape(1);
int H = heatmap.shape(2);
int W = heatmap.shape(3);
int px = *(pred.data<float>() + k * 3 + 0);
int py = *(pred.data<float>() + k * 3 + 1);
if (1 < px && px < W - 2 && 1 < py && py < H - 2) {
float* data = const_cast<float*>(heatmap.data<float>() + k * H * W);
auto get_data = [&](int r, int c) { return *(data + r * W + c); };
float dx = 0.5 * (get_data(py, px + 1) - get_data(py, px - 1));
float dy = 0.5 * (get_data(py + 1, px) - get_data(py - 1, px));
float dxx = 0.25 * (get_data(py, px + 2) - 2 * get_data(py, px) + get_data(py, px - 2));
float dxy = 0.25 * (get_data(py + 1, px + 1) - get_data(py - 1, px + 1) -
get_data(py + 1, px - 1) + get_data(py - 1, px - 1));
float dyy = 0.25 * (get_data(py + 2, px) - 2 * get_data(py, px) + get_data(py - 2, px));
vector<float> _data0 = {dx, dy};
vector<float> _data1 = {dxx, dxy, dxy, dyy};
cv::Mat derivative = cv::Mat(2, 1, CV_32FC1, _data0.data());
cv::Mat hessian = cv::Mat(2, 2, CV_32FC1, _data1.data());
if (std::fabs(dxx * dyy - dxy * dxy) > 1e-6) {
cv::Mat hessianinv = hessian.inv();
cv::Mat offset = -hessianinv * derivative;
*(pred.data<float>() + k * 3 + 0) += offset.at<float>(0, 0);
*(pred.data<float>() + k * 3 + 1) += offset.at<float>(1, 0);
}
}
}
Tensor gaussian_blur(const Tensor& _heatmap, int kernel) {
assert(kernel % 2 == 1);
auto desc = _heatmap.desc();
Tensor heatmap(desc);
int K = _heatmap.shape(1);
int H = _heatmap.shape(2);
int W = _heatmap.shape(3);
int num_points = H * W;
int border = (kernel - 1) / 2;
for (int i = 0; i < K; i++) {
int offset = i * H * W;
float* data = const_cast<float*>(_heatmap.data<float>()) + offset;
float origin_max = *std::max_element(data, data + num_points);
cv::Mat work = cv::Mat(H + 2 * border, W + 2 * border, CV_32FC1, cv::Scalar{});
cv::Mat curr = cv::Mat(H, W, CV_32FC1, data);
cv::Rect roi = {border, border, W, H};
curr.copyTo(work(roi));
cv::GaussianBlur(work, work, {kernel, kernel}, 0);
cv::Mat valid = work(roi).clone();
float cur_max = *std::max_element((float*)valid.data, (float*)valid.data + num_points);
float* dst = heatmap.data<float>() + offset;
std::transform((float*)valid.data, (float*)valid.data + num_points, dst,
[&](float v) { return v * origin_max / cur_max; });
}
return heatmap;
}
Tensor get_max_pred(const Tensor& heatmap) {
int K = heatmap.shape(1);
int H = heatmap.shape(2);
int W = heatmap.shape(3);
int num_points = H * W;
TensorDesc pred_desc = {Device{"cpu"}, DataType::kFLOAT, {1, K, 3}};
Tensor pred(pred_desc);
cv::parallel_for_(cv::Range(0, K), _LoopBody{[&](const cv::Range& r) {
for (int i = r.start; i < r.end; i++) {
float* src_data = const_cast<float*>(heatmap.data<float>()) + i * H * W;
cv::Mat mat = cv::Mat(H, W, CV_32FC1, src_data);
double min_val, max_val;
cv::Point min_loc, max_loc;
cv::minMaxLoc(mat, &min_val, &max_val, &min_loc, &max_loc);
float* dst_data = pred.data<float>() + i * 3;
*(dst_data + 0) = -1;
*(dst_data + 1) = -1;
*(dst_data + 2) = max_val;
if (max_val > 0.0) {
*(dst_data + 0) = max_loc.x;
*(dst_data + 1) = max_loc.y;
}
}
}});
return pred;
}
private:
bool flip_test_{true};
bool shift_heatmap_{true};
string post_process_ = {"default"};
int modulate_kernel_{11};
bool unbiased_decoding_{false};
float valid_radius_factor_{0.0546875f};
bool use_udp_{false};
string target_type_{"GaussianHeatmap"};
};
REGISTER_CODEBASE_COMPONENT(MMPose, TopdownHeatmapBaseHeadDecode);
// decode process is same
using TopdownHeatmapSimpleHeadDecode = TopdownHeatmapBaseHeadDecode;
REGISTER_CODEBASE_COMPONENT(MMPose, TopdownHeatmapSimpleHeadDecode);
using TopdownHeatmapMultiStageHeadDecode = TopdownHeatmapBaseHeadDecode;
REGISTER_CODEBASE_COMPONENT(MMPose, TopdownHeatmapMultiStageHeadDecode);
using ViPNASHeatmapSimpleHeadDecode = TopdownHeatmapBaseHeadDecode;
REGISTER_CODEBASE_COMPONENT(MMPose, ViPNASHeatmapSimpleHeadDecode);
using TopdownHeatmapMSMUHeadDecode = TopdownHeatmapBaseHeadDecode;
REGISTER_CODEBASE_COMPONENT(MMPose, TopdownHeatmapMSMUHeadDecode);
} // namespace mmdeploy::mmpose
Reference in New Issue View Git Blame Copy Permalink