260 lines
9.5 KiB
C++
260 lines
9.5 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#ifndef MMDEPLOY_VISUALIZE_H
|
|
#define MMDEPLOY_VISUALIZE_H
|
|
|
|
#include <algorithm>
|
|
#include <iomanip>
|
|
#include <numeric>
|
|
#include <vector>
|
|
|
|
#include "opencv2/highgui/highgui.hpp"
|
|
#include "opencv2/imgproc/imgproc.hpp"
|
|
#include "palette.h"
|
|
#include "skeleton.h"
|
|
|
|
namespace utils {
|
|
|
|
class Visualize {
|
|
public:
|
|
class Session {
|
|
public:
|
|
explicit Session(Visualize& v, const cv::Mat& frame) : v_(v) {
|
|
if (v_.size_) {
|
|
scale_ = (float)v_.size_ / (float)std::max(frame.cols, frame.rows);
|
|
}
|
|
cv::Mat img;
|
|
if (v.background_ == "black") {
|
|
img = cv::Mat::zeros(frame.size(), CV_8UC3);
|
|
} else {
|
|
img = frame;
|
|
if (img.channels() == 1) {
|
|
cv::cvtColor(img, img, cv::COLOR_GRAY2BGR);
|
|
}
|
|
}
|
|
if (scale_ != 1) {
|
|
cv::resize(img, img, {}, scale_, scale_);
|
|
} else if (img.data == frame.data) {
|
|
img = img.clone();
|
|
}
|
|
img_ = std::move(img);
|
|
}
|
|
|
|
void add_label(int label_id, float score, int index) {
|
|
printf("label: %d, label_id: %d, score: %.4f\n", index, label_id, score);
|
|
auto size = .5f * static_cast<float>(img_.rows + img_.cols);
|
|
offset_ += add_text(to_text(label_id, score), {1, (float)offset_}, size) + 2;
|
|
}
|
|
|
|
int add_text(const std::string& text, const cv::Point2f& origin, float size) {
|
|
static constexpr const int font_face = cv::FONT_HERSHEY_SIMPLEX;
|
|
static constexpr const int thickness = 1;
|
|
static constexpr const auto max_font_scale = .5f;
|
|
static constexpr const auto min_font_scale = .25f;
|
|
float font_scale{};
|
|
if (size < 20) {
|
|
font_scale = min_font_scale;
|
|
} else if (size > 200) {
|
|
font_scale = max_font_scale;
|
|
} else {
|
|
font_scale = min_font_scale + (size - 20) / (200 - 20) * (max_font_scale - min_font_scale);
|
|
}
|
|
int baseline{};
|
|
auto text_size = cv::getTextSize(text, font_face, font_scale, thickness, &baseline);
|
|
cv::Rect rect(origin + cv::Point2f(0, text_size.height + 2 * thickness),
|
|
origin + cv::Point2f(text_size.width, 0));
|
|
rect &= cv::Rect({}, img_.size());
|
|
if (rect.area() > 0) {
|
|
img_(rect) *= .35f;
|
|
cv::putText(img_, text, origin + cv::Point2f(0, text_size.height), font_face, font_scale,
|
|
cv::Scalar::all(255), thickness, cv::LINE_AA);
|
|
}
|
|
return text_size.height;
|
|
}
|
|
|
|
static std::string to_text(int label_id, float score) {
|
|
std::stringstream ss;
|
|
ss << label_id << ": " << std::fixed << std::setprecision(1) << score * 100;
|
|
return ss.str();
|
|
}
|
|
|
|
template <typename Mask>
|
|
void add_det(const mmdeploy_rect_t& rect, int label_id, float score, const Mask* mask,
|
|
int index) {
|
|
printf("bbox %d, left=%.2f, top=%.2f, right=%.2f, bottom=%.2f, label=%d, score=%.4f\n", index,
|
|
rect.left, rect.top, rect.right, rect.bottom, label_id, score);
|
|
if (mask) {
|
|
fprintf(stdout, "mask %d, height=%d, width=%d\n", index, mask->height, mask->width);
|
|
auto x0 = (int)std::max(std::floor(rect.left) - 1, 0.f);
|
|
auto y0 = (int)std::max(std::floor(rect.top) - 1, 0.f);
|
|
add_instance_mask({x0, y0}, rand(), mask->data, mask->height, mask->width);
|
|
}
|
|
add_bbox(rect, label_id, score);
|
|
}
|
|
|
|
void add_instance_mask(const cv::Point& origin, int color_id, const char* mask_data, int mask_h,
|
|
int mask_w, float alpha = .5f) {
|
|
auto color = v_.palette_.data[color_id % v_.palette_.data.size()];
|
|
auto x_end = std::min(origin.x + mask_w, img_.cols);
|
|
auto y_end = std::min(origin.y + mask_h, img_.rows);
|
|
auto img_data = img_.ptr<cv::Vec3b>();
|
|
for (int i = origin.y; i < y_end; ++i) {
|
|
for (int j = origin.x; j < x_end; ++j) {
|
|
if (mask_data[(i - origin.y) * mask_w + (j - origin.x)]) {
|
|
img_data[i * img_.cols + j] = img_data[i * img_.cols + j] * (1 - alpha) + color * alpha;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void add_bbox(mmdeploy_rect_t rect, int label_id, float score) {
|
|
rect.left *= scale_;
|
|
rect.right *= scale_;
|
|
rect.top *= scale_;
|
|
rect.bottom *= scale_;
|
|
if (label_id >= 0 && score > 0) {
|
|
auto area = std::max(0.f, (rect.right - rect.left) * (rect.bottom - rect.top));
|
|
add_text(to_text(label_id, score), {rect.left, rect.top}, std::sqrt(area));
|
|
}
|
|
cv::rectangle(img_, cv::Point2f(rect.left, rect.top), cv::Point2f(rect.right, rect.bottom),
|
|
cv::Scalar(0, 255, 0));
|
|
}
|
|
|
|
void add_text_det(mmdeploy_point_t bbox[4], float score, const char* text, size_t text_size,
|
|
int index) {
|
|
printf("bbox[%d]: (%.2f, %.2f), (%.2f, %.2f), (%.2f, %.2f), (%.2f, %.2f), %.2f\n", index, //
|
|
bbox[0].x, bbox[0].y, //
|
|
bbox[1].x, bbox[1].y, //
|
|
bbox[2].x, bbox[2].y, //
|
|
bbox[3].x, bbox[3].y, score);
|
|
std::vector<cv::Point> poly_points;
|
|
cv::Point2f center{};
|
|
for (int i = 0; i < 4; ++i) {
|
|
poly_points.emplace_back(bbox[i].x * scale_, bbox[i].y * scale_);
|
|
center += cv::Point2f(poly_points.back());
|
|
}
|
|
cv::polylines(img_, poly_points, true, cv::Scalar{0, 255, 0}, 1, cv::LINE_AA);
|
|
if (text) {
|
|
auto area = cv::contourArea(poly_points);
|
|
fprintf(stdout, "text[%d]: %s\n", index, text);
|
|
add_text(std::string(text, text + text_size), center / 4, std::sqrt(area));
|
|
}
|
|
}
|
|
|
|
void add_rotated_det(const float bbox[5], int label_id, float score) {
|
|
float xc = bbox[0] * scale_;
|
|
float yc = bbox[1] * scale_;
|
|
float w = bbox[2] * scale_;
|
|
float h = bbox[3] * scale_;
|
|
float ag = bbox[4];
|
|
float wx = w / 2 * std::cos(ag);
|
|
float wy = w / 2 * std::sin(ag);
|
|
float hx = -h / 2 * std::sin(ag);
|
|
float hy = h / 2 * std::cos(ag);
|
|
cv::Point2f p1{xc - wx - hx, yc - wy - hy};
|
|
cv::Point2f p2{xc + wx - hx, yc + wy - hy};
|
|
cv::Point2f p3{xc + wx + hx, yc + wy + hy};
|
|
cv::Point2f p4{xc - wx + hx, yc - wy + hy};
|
|
cv::Point2f c = .25f * (p1 + p2 + p3 + p4);
|
|
cv::drawContours(
|
|
img_,
|
|
std::vector<std::vector<cv::Point>>{{p1 * scale_, p2 * scale_, p3 * scale_, p4 * scale_}},
|
|
-1, {0, 255, 0}, 2, cv::LINE_AA);
|
|
add_text(to_text(label_id, score), c, std::sqrt(w * h));
|
|
}
|
|
|
|
void add_mask(int height, int width, int n_classes, const int* mask, const float* score) {
|
|
cv::Mat color_mask = cv::Mat::zeros(height, width, CV_8UC3);
|
|
auto n_pix = color_mask.total();
|
|
|
|
// compute top 1 idx if score (CHW) is available
|
|
cv::Mat_<int> top;
|
|
if (!mask && score) {
|
|
top = cv::Mat_<int>::zeros(height, width);
|
|
for (auto c = 1; c < n_classes; ++c) {
|
|
top.forEach([&](int& x, const int* idx) {
|
|
auto offset = idx[0] * width + idx[1];
|
|
if (score[c * n_pix + offset] > score[x * n_pix + offset]) {
|
|
x = c;
|
|
}
|
|
});
|
|
}
|
|
mask = top.ptr<int>();
|
|
}
|
|
|
|
if (mask) {
|
|
// palette look-up
|
|
color_mask.forEach<cv::Vec3b>([&](cv::Vec3b& x, const int* idx) {
|
|
auto& palette = v_.palette_.data;
|
|
x = palette[mask[idx[0] * width + idx[1]] % palette.size()];
|
|
});
|
|
|
|
if (color_mask.size() != img_.size()) {
|
|
cv::resize(color_mask, color_mask, img_.size());
|
|
}
|
|
|
|
// blend mask and background image
|
|
cv::addWeighted(img_, .5, color_mask, .5, 0., img_);
|
|
}
|
|
}
|
|
|
|
void add_pose(const mmdeploy_point_t* pts, const float* scores, int32_t pts_size, double thr) {
|
|
auto& skel = v_.skeleton_;
|
|
if (skel.point_colors.size() != pts_size) {
|
|
std::cout << "error: mismatched number of keypoints: " << skel.point_colors.size() << " vs "
|
|
<< pts_size << ", skip pose visualization.\n";
|
|
return;
|
|
}
|
|
std::vector<int> used(pts_size);
|
|
std::vector<int> is_end_point(pts_size);
|
|
for (size_t i = 0; i < skel.links.size(); ++i) {
|
|
auto u = skel.links[i].first;
|
|
auto v = skel.links[i].second;
|
|
is_end_point[u] = is_end_point[v] = 1;
|
|
if (scores[u] > thr && scores[v] > thr) {
|
|
used[u] = used[v] = 1;
|
|
cv::Point2f p0(pts[u].x, pts[u].y);
|
|
cv::Point2f p1(pts[v].x, pts[v].y);
|
|
cv::line(img_, p0 * scale_, p1 * scale_, skel.palette[skel.link_colors[i]], 1,
|
|
cv::LINE_AA);
|
|
}
|
|
}
|
|
for (size_t i = 0; i < pts_size; ++i) {
|
|
if (!is_end_point[i] && scores[i] > thr || used[i]) {
|
|
cv::Point2f p(pts[i].x, pts[i].y);
|
|
cv::circle(img_, p * scale_, 1, skel.palette[skel.point_colors[i]], 2, cv::LINE_AA);
|
|
}
|
|
}
|
|
}
|
|
|
|
cv::Mat get() { return img_; }
|
|
|
|
private:
|
|
Visualize& v_;
|
|
float scale_{1};
|
|
int offset_{1};
|
|
cv::Mat img_;
|
|
};
|
|
|
|
explicit Visualize(int size = 0) : size_(size) { palette_ = Palette::get(32); }
|
|
|
|
Session get_session(const cv::Mat& frame) { return Session(*this, frame); }
|
|
|
|
void set_skeleton(const Skeleton& skeleton) { skeleton_ = skeleton; }
|
|
|
|
void set_palette(const Palette& palette) { palette_ = palette; }
|
|
|
|
void set_background(const std::string& background) { background_ = background; }
|
|
|
|
private:
|
|
friend Session;
|
|
Skeleton skeleton_;
|
|
Palette palette_;
|
|
std::string background_;
|
|
int size_{};
|
|
};
|
|
|
|
} // namespace utils
|
|
|
|
#endif // MMDEPLOY_VISUALIZE_H
|