51 lines
1.7 KiB
C++
51 lines
1.7 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "core/utils/device_utils.h"
|
|
#include "opencv_utils.h"
|
|
#include "preprocess/transform/pad.h"
|
|
|
|
using namespace std;
|
|
|
|
namespace mmdeploy {
|
|
namespace cpu {
|
|
|
|
class PadImpl : public ::mmdeploy::PadImpl {
|
|
public:
|
|
PadImpl(const Value& args) : ::mmdeploy::PadImpl(args) {
|
|
static map<string, int> border_map{{"constant", cv::BORDER_CONSTANT},
|
|
{"edge", cv::BORDER_REPLICATE},
|
|
{"reflect", cv::BORDER_REFLECT_101},
|
|
{"symmetric", cv::BORDER_REFLECT}};
|
|
if (border_map.find(arg_.padding_mode) == border_map.end()) {
|
|
ERROR("unsupported padding_mode '{}'", arg_.padding_mode);
|
|
throw std::invalid_argument("unsupported padding_mode");
|
|
}
|
|
border_type_ = border_map[arg_.padding_mode];
|
|
}
|
|
|
|
protected:
|
|
Result<Tensor> PadImage(const Tensor& img, const std::array<int, 4>& padding) override {
|
|
OUTCOME_TRY(auto tensor, MakeAvailableOnDevice(img, device_, stream_));
|
|
cv::Mat dst_mat = Pad(Tensor2CVMat(tensor), padding[1], padding[0], padding[3], padding[2],
|
|
border_type_, arg_.pad_val);
|
|
return CVMat2Tensor(dst_mat);
|
|
}
|
|
|
|
private:
|
|
int border_type_;
|
|
};
|
|
|
|
class PadImplCreator : public Creator<::mmdeploy::PadImpl> {
|
|
public:
|
|
const char* GetName() const override { return "cpu"; }
|
|
int GetVersion() const override { return 1; }
|
|
ReturnType Create(const Value& args) override { return make_unique<PadImpl>(args); }
|
|
};
|
|
|
|
} // namespace cpu
|
|
} // namespace mmdeploy
|
|
|
|
using mmdeploy::PadImpl;
|
|
using mmdeploy::cpu::PadImplCreator;
|
|
REGISTER_MODULE(PadImpl, PadImplCreator);
|