81 lines
2.8 KiB
C++
81 lines
2.8 KiB
C++
|
// Copyright (c) OpenMMLab. All rights reserved.
|
||
|
|
|
||
|
|
#include <numeric>
|
||
|
|
|
||
|
|
#include "archive/value_archive.h"
|
||
|
|
#include "core/module.h"
|
||
|
|
#include "core/registry.h"
|
||
|
|
#include "core/utils/formatter.h"
|
||
|
|
#include "core/value.h"
|
||
|
|
#include "experimental/module_adapter.h"
|
||
|
|
#include "preprocess/cpu/opencv_utils.h"
|
||
|
|
|
||
|
|
namespace mmdeploy {
|
||
|
|
|
||
|
|
// warp rotated rect
|
||
|
|
class WarpBoxes {
|
||
|
|
public:
|
||
|
|
Result<Value> operator()(const Value& img, const Value& det) {
|
||
|
|
Value patches = ValueType::kArray;
|
||
|
|
if (det.is_object() && det.contains("boxes")) {
|
||
|
|
auto boxes = from_value<std::vector<std::vector<cv::Point>>>(det["boxes"]);
|
||
|
|
auto ori_img = mmdeploy::cpu::Mat2CVMat(img["ori_img"].get<mmdeploy::Mat>()).clone();
|
||
|
|
for (int i = 0; i < boxes.size(); ++i) {
|
||
|
|
auto patch = warp(ori_img, boxes[i]);
|
||
|
|
patches.push_back(make_pointer({{"ori_img", cpu::CVMat2Mat(patch, PixelFormat::kBGR)}}));
|
||
|
|
// cv::imwrite(std::to_string(i) + ".png", patch);
|
||
|
|
}
|
||
|
|
} else { // whole image as a bbox
|
||
|
|
patches.push_back({{"ori_img", img["ori_img"].get<mmdeploy::Mat>()}});
|
||
|
|
}
|
||
|
|
return patches;
|
||
|
|
}
|
||
|
|
|
||
|
|
// assuming rect
|
||
|
|
static cv::Mat warp(const cv::Mat& img, const std::vector<cv::Point>& _pts) {
|
||
|
|
auto pts = sort_vertex(_pts);
|
||
|
|
std::vector<cv::Point2f> src(begin(pts), end(pts));
|
||
|
|
auto e0 = norm(pts[0] - pts[1]);
|
||
|
|
auto e1 = norm(pts[1] - pts[2]);
|
||
|
|
auto w = static_cast<float>(std::max(e0, e1));
|
||
|
|
auto h = static_cast<float>(std::min(e0, e1));
|
||
|
|
std::vector<cv::Point2f> dst{{0, 0}, {w, 0}, {w, h}, {0, h}};
|
||
|
|
auto m = getAffineTransform(src.data(), dst.data());
|
||
|
|
cv::Mat warped;
|
||
|
|
warpAffine(img, warped, m, {static_cast<int>(w), static_cast<int>(h)});
|
||
|
|
return warped;
|
||
|
|
}
|
||
|
|
|
||
|
|
static std::vector<cv::Point> sort_vertex(std::vector<cv::Point> ps) {
|
||
|
|
auto pivot = *min_element(begin(ps), end(ps), [](auto r, auto p) {
|
||
|
|
return (r.y != p.y) ? (r.y < p.y) : (r.x < p.x);
|
||
|
|
});
|
||
|
|
// TODO: resolve tie with distance
|
||
|
|
sort(begin(ps), end(ps), [&](auto a, auto b) {
|
||
|
|
if (a == pivot) return b != pivot;
|
||
|
|
return (a - pivot).cross(b - pivot) > 0;
|
||
|
|
});
|
||
|
|
auto tl = accumulate(begin(ps) + 1, end(ps), ps[0], [](auto r, auto p) {
|
||
|
|
return cv::Point{std::min(r.x, p.x), std::min(r.y, p.y)};
|
||
|
|
});
|
||
|
|
auto cmp = [&](auto r, auto p) {
|
||
|
|
cv::Point2f u{r - tl}, v{p - tl};
|
||
|
|
return u.dot(u) < v.dot(v);
|
||
|
|
};
|
||
|
|
auto tl_idx = min_element(begin(ps), end(ps), cmp) - begin(ps);
|
||
|
|
rotate(begin(ps), begin(ps) + tl_idx, end(ps));
|
||
|
|
return ps;
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
class WarpBoxesCreator : public Creator<Module> {
|
||
|
|
public:
|
||
|
|
const char* GetName() const override { return "WarpBoxes"; }
|
||
|
|
int GetVersion() const override { return 0; }
|
||
|
|
std::unique_ptr<Module> Create(const Value& value) override { return CreateTask(WarpBoxes{}); }
|
||
|
|
};
|
||
|
|
|
||
|
|
REGISTER_MODULE(Module, WarpBoxesCreator);
|
||
|
|
|
||
|
|
} // namespace mmdeploy
|