Skip to content

Commit

Permalink
Support interpolate flags in image_to_tensor_converter_opencv.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 607863874
  • Loading branch information
MediaPipe Team authored and copybara-github committed Feb 17, 2024
1 parent 1906a11 commit 1d3105a
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 6 deletions.
12 changes: 7 additions & 5 deletions mediapipe/calculators/tensor/image_to_tensor_converter_opencv.cc
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,9 @@ namespace {

class OpenCvProcessor : public ImageToTensorConverter {
public:
OpenCvProcessor(BorderMode border_mode, Tensor::ElementType tensor_type)
: tensor_type_(tensor_type) {
OpenCvProcessor(BorderMode border_mode, Tensor::ElementType tensor_type,
cv::InterpolationFlags flags)
: tensor_type_(tensor_type), flags_(flags) {
switch (border_mode) {
case BorderMode::kReplicate:
border_mode_ = cv::BORDER_REPLICATE;
Expand Down Expand Up @@ -148,7 +149,7 @@ class OpenCvProcessor : public ImageToTensorConverter {
cv::Mat transformed;
cv::warpPerspective(*src, transformed, projection_matrix,
cv::Size(dst_width, dst_height),
/*flags=*/cv::INTER_LINEAR,
/*flags=*/flags_,
/*borderMode=*/border_mode_);

if (transformed.channels() > output_channels) {
Expand Down Expand Up @@ -181,6 +182,7 @@ class OpenCvProcessor : public ImageToTensorConverter {

enum cv::BorderTypes border_mode_;
Tensor::ElementType tensor_type_;
cv::InterpolationFlags flags_;
int mat_type_;
int mat_gray_type_;
};
Expand All @@ -189,15 +191,15 @@ class OpenCvProcessor : public ImageToTensorConverter {

absl::StatusOr<std::unique_ptr<ImageToTensorConverter>> CreateOpenCvConverter(
CalculatorContext* cc, BorderMode border_mode,
Tensor::ElementType tensor_type) {
Tensor::ElementType tensor_type, cv::InterpolationFlags flags) {
if (tensor_type != Tensor::ElementType::kInt8 &&
tensor_type != Tensor::ElementType::kFloat32 &&
tensor_type != Tensor::ElementType::kUInt8) {
return absl::InvalidArgumentError(absl::StrCat(
"Tensor type is currently not supported by OpenCvProcessor, type: ",
tensor_type));
}
return absl::make_unique<OpenCvProcessor>(border_mode, tensor_type);
return std::make_unique<OpenCvProcessor>(border_mode, tensor_type, flags);
}

} // namespace mediapipe
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,16 @@

#include "mediapipe/calculators/tensor/image_to_tensor_converter.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
#include "mediapipe/framework/port/statusor.h"

namespace mediapipe {

// Creates OpenCV image-to-tensor converter.
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>> CreateOpenCvConverter(
CalculatorContext* cc, BorderMode border_mode,
Tensor::ElementType tensor_type);
Tensor::ElementType tensor_type,
cv::InterpolationFlags flags = cv::INTER_LINEAR);

} // namespace mediapipe

Expand Down

0 comments on commit 1d3105a

Please sign in to comment.