|
| 1 | +// Copyright (c) OpenMMLab. All rights reserved. |
| 2 | + |
| 3 | +#include "mmdeploy/core/utils/formatter.h" |
| 4 | +#include "mmdeploy/operation/vision.h" |
| 5 | +#include "ppl/cv/cuda/warpaffine.h" |
| 6 | + |
| 7 | +namespace mmdeploy::operation::cuda { |
| 8 | + |
| 9 | +class WarpAffineImpl : public WarpAffine { |
| 10 | + public: |
| 11 | + explicit WarpAffineImpl(ppl::cv::InterpolationType interp) : interp_(interp) {} |
| 12 | + |
| 13 | + Result<void> apply(const Tensor& src, Tensor& dst, const float affine_matrix[6], int dst_h, |
| 14 | + int dst_w) override { |
| 15 | + assert(src.device() == device()); |
| 16 | + |
| 17 | + TensorDesc desc{device(), src.data_type(), {1, dst_h, dst_w, src.shape(3)}, src.name()}; |
| 18 | + Tensor dst_tensor(desc); |
| 19 | + |
| 20 | + const auto m = affine_matrix; |
| 21 | + auto inv = Invert(affine_matrix); |
| 22 | + |
| 23 | + auto cuda_stream = GetNative<cudaStream_t>(stream()); |
| 24 | + if (src.data_type() == DataType::kINT8) { |
| 25 | + OUTCOME_TRY(Dispatch<uint8_t>(src, dst_tensor, inv.data(), cuda_stream)); |
| 26 | + } else if (src.data_type() == DataType::kFLOAT) { |
| 27 | + OUTCOME_TRY(Dispatch<float>(src, dst_tensor, inv.data(), cuda_stream)); |
| 28 | + } else { |
| 29 | + MMDEPLOY_ERROR("unsupported data type {}", src.data_type()); |
| 30 | + return Status(eNotSupported); |
| 31 | + } |
| 32 | + |
| 33 | + dst = std::move(dst_tensor); |
| 34 | + return success(); |
| 35 | + } |
| 36 | + |
| 37 | + private: |
| 38 | + // ppl.cv uses inverted transform |
| 39 | + // https://github.com/opencv/opencv/blob/bc6544c0bcfa9ca5db5e0d0551edf5c8e7da3852/modules/imgproc/src/imgwarp.cpp#L3478 |
| 40 | + static std::array<float, 6> Invert(const float affine_matrix[6]) { |
| 41 | + const auto* M = affine_matrix; |
| 42 | + std::array<float, 6> inv{}; |
| 43 | + auto iM = inv.data(); |
| 44 | + |
| 45 | + auto D = M[0] * M[3 + 1] - M[1] * M[3]; |
| 46 | + D = D != 0.f ? 1.f / D : 0.f; |
| 47 | + auto A11 = M[3 + 1] * D, A22 = M[0] * D, A12 = -M[1] * D, A21 = -M[3] * D; |
| 48 | + auto b1 = -A11 * M[2] - A12 * M[3 + 2]; |
| 49 | + auto b2 = -A21 * M[2] - A22 * M[3 + 2]; |
| 50 | + |
| 51 | + iM[0] = A11; |
| 52 | + iM[1] = A12; |
| 53 | + iM[2] = b1; |
| 54 | + iM[3] = A21; |
| 55 | + iM[3 + 1] = A22; |
| 56 | + iM[3 + 2] = b2; |
| 57 | + |
| 58 | + return inv; |
| 59 | + } |
| 60 | + |
| 61 | + template <typename T> |
| 62 | + auto Select(int channels) -> decltype(&ppl::cv::cuda::WarpAffine<T, 1>) { |
| 63 | + switch (channels) { |
| 64 | + case 1: |
| 65 | + return &ppl::cv::cuda::WarpAffine<T, 1>; |
| 66 | + case 3: |
| 67 | + return &ppl::cv::cuda::WarpAffine<T, 3>; |
| 68 | + case 4: |
| 69 | + return &ppl::cv::cuda::WarpAffine<T, 4>; |
| 70 | + default: |
| 71 | + MMDEPLOY_ERROR("unsupported channels {}", channels); |
| 72 | + return nullptr; |
| 73 | + } |
| 74 | + } |
| 75 | + |
| 76 | + template <class T> |
| 77 | + Result<void> Dispatch(const Tensor& src, Tensor& dst, const float affine_matrix[6], |
| 78 | + cudaStream_t stream) { |
| 79 | + int h = (int)src.shape(1); |
| 80 | + int w = (int)src.shape(2); |
| 81 | + int c = (int)src.shape(3); |
| 82 | + int dst_h = (int)dst.shape(1); |
| 83 | + int dst_w = (int)dst.shape(2); |
| 84 | + |
| 85 | + auto input = src.data<T>(); |
| 86 | + auto output = dst.data<T>(); |
| 87 | + |
| 88 | + ppl::common::RetCode ret = 0; |
| 89 | + |
| 90 | + if (auto warp_affine = Select<T>(c); warp_affine) { |
| 91 | + ret = warp_affine(stream, h, w, w * c, input, dst_h, dst_w, dst_w * c, output, affine_matrix, |
| 92 | + interp_, ppl::cv::BORDER_CONSTANT, 0); |
| 93 | + } else { |
| 94 | + return Status(eNotSupported); |
| 95 | + } |
| 96 | + |
| 97 | + return ret == 0 ? success() : Result<void>(Status(eFail)); |
| 98 | + } |
| 99 | + |
| 100 | + ppl::cv::InterpolationType interp_; |
| 101 | +}; |
| 102 | + |
| 103 | +static auto Create(const string_view& interp) { |
| 104 | + ppl::cv::InterpolationType type{}; |
| 105 | + if (interp == "bilinear") { |
| 106 | + type = ppl::cv::InterpolationType::INTERPOLATION_LINEAR; |
| 107 | + } else if (interp == "nearest") { |
| 108 | + type = ppl::cv::InterpolationType::INTERPOLATION_NEAREST_POINT; |
| 109 | + } else { |
| 110 | + MMDEPLOY_ERROR("unsupported interpolation method: {}", interp); |
| 111 | + throw_exception(eNotSupported); |
| 112 | + } |
| 113 | + return std::make_unique<WarpAffineImpl>(type); |
| 114 | +} |
| 115 | + |
| 116 | +MMDEPLOY_REGISTER_FACTORY_FUNC(WarpAffine, (cuda, 0), Create); |
| 117 | + |
| 118 | +} // namespace mmdeploy::operation::cuda |
0 commit comments