64 lines
2.0 KiB
C++
64 lines
2.0 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "codebase/mmocr/panet.h"
|
|
|
|
#include "opencv2/imgproc.hpp"
|
|
|
|
namespace mmdeploy::mmocr {
|
|
|
|
class PaHeadCpuImpl : public PaHeadImpl {
|
|
public:
|
|
Result<void> Process(Tensor text_pred, //
|
|
Tensor kernel_pred, //
|
|
Tensor embed_pred, //
|
|
float min_text_confidence, //
|
|
float min_kernel_confidence, //
|
|
cv::Mat_<float>& text_score, //
|
|
cv::Mat_<uint8_t>& text, //
|
|
cv::Mat_<uint8_t>& kernel, //
|
|
cv::Mat_<int>& label, //
|
|
cv::Mat_<float>& embed, //
|
|
int& region_num) override {
|
|
OUTCOME_TRY(stream_.Wait());
|
|
|
|
auto height = static_cast<int>(text_pred.shape(1));
|
|
auto width = static_cast<int>(text_pred.shape(2));
|
|
|
|
text_score = cv::Mat_<float>(height, width, text_pred.data<float>());
|
|
sigmoid(text_score);
|
|
|
|
text = text_score > min_text_confidence;
|
|
|
|
cv::Mat_<float> kernel_score(height, width, kernel_pred.data<float>());
|
|
sigmoid(kernel_score);
|
|
|
|
kernel = kernel_score > min_kernel_confidence & text;
|
|
|
|
auto n_embed_channels = static_cast<int>(embed_pred.shape(0));
|
|
embed = cv::Mat_<float>(n_embed_channels, height * width, embed_pred.data<float>());
|
|
cv::transpose(embed, embed);
|
|
|
|
region_num = cv::connectedComponents(kernel, label, 4, CV_32S);
|
|
|
|
return success();
|
|
}
|
|
|
|
static void sigmoid(cv::Mat_<float>& score) {
|
|
cv::exp(-score, score);
|
|
score = 1 / (1 + score);
|
|
}
|
|
};
|
|
|
|
class PaHeadCpuImplCreator : public ::mmdeploy::Creator<PaHeadImpl> {
|
|
public:
|
|
const char* GetName() const override { return "cpu"; }
|
|
int GetVersion() const override { return 0; }
|
|
std::unique_ptr<PaHeadImpl> Create(const Value&) override {
|
|
return std::make_unique<PaHeadCpuImpl>();
|
|
}
|
|
};
|
|
|
|
REGISTER_MODULE(PaHeadImpl, PaHeadCpuImplCreator);
|
|
|
|
} // namespace mmdeploy::mmocr
|