From a4487903009c6513b4d015ff6e6632f2fd3b3b25 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 11 Apr 2023 03:17:01 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 523351901 --- mediapipe/calculators/tensor/BUILD | 25 ++ .../tensor/image_to_tensor_calculator.cc | 17 +- .../image_to_tensor_converter_frame_buffer.cc | 246 ++++++++++++++++++ .../image_to_tensor_converter_frame_buffer.h | 34 +++ mediapipe/framework/BUILD | 9 + mediapipe/framework/formats/image.h | 5 +- mediapipe/framework/port/BUILD | 12 + 7 files changed, 344 insertions(+), 4 deletions(-) create mode 100644 mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.cc create mode 100644 mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.h diff --git a/mediapipe/calculators/tensor/BUILD b/mediapipe/calculators/tensor/BUILD index 9ae88425..a574e11e 100644 --- a/mediapipe/calculators/tensor/BUILD +++ b/mediapipe/calculators/tensor/BUILD @@ -988,6 +988,11 @@ cc_library( }) + select({ "//mediapipe/framework/port:disable_opencv": [], "//conditions:default": [":image_to_tensor_converter_opencv"], + }) + select({ + "//mediapipe/framework/port:enable_halide": [ + ":image_to_tensor_converter_frame_buffer", + ], + "//conditions:default": [], }), alwayslink = 1, ) @@ -1121,6 +1126,26 @@ cc_library( ], ) +cc_library( + name = "image_to_tensor_converter_frame_buffer", + srcs = ["image_to_tensor_converter_frame_buffer.cc"], + hdrs = ["image_to_tensor_converter_frame_buffer.h"], + deps = [ + ":image_to_tensor_converter", + ":image_to_tensor_utils", + "//mediapipe/framework:calculator_context", + "//mediapipe/framework/formats:frame_buffer", + "//mediapipe/framework/formats:image", + "//mediapipe/framework/formats:tensor", + "//mediapipe/framework/port:status", + "//mediapipe/gpu:frame_buffer_view", + "//mediapipe/util/frame_buffer:frame_buffer_util", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:str_format", + ], +) + cc_library( name = "image_to_tensor_converter_gl_buffer", srcs = ["image_to_tensor_converter_gl_buffer.cc"], diff --git a/mediapipe/calculators/tensor/image_to_tensor_calculator.cc b/mediapipe/calculators/tensor/image_to_tensor_calculator.cc index d15d3508..344e12da 100644 --- a/mediapipe/calculators/tensor/image_to_tensor_calculator.cc +++ b/mediapipe/calculators/tensor/image_to_tensor_calculator.cc @@ -34,6 +34,8 @@ #if !MEDIAPIPE_DISABLE_OPENCV #include "mediapipe/calculators/tensor/image_to_tensor_converter_opencv.h" +#elif MEDIAPIPE_ENABLE_HALIDE +#include "mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.h" #endif #if !MEDIAPIPE_DISABLE_GPU @@ -273,10 +275,19 @@ class ImageToTensorCalculator : public Node { CreateOpenCvConverter( cc, GetBorderMode(options_.border_mode()), GetOutputTensorType(/*uses_gpu=*/false, params_))); +// TODO: FrameBuffer-based converter needs to call GetGpuBuffer() +// to get access to a FrameBuffer view. Investigate if GetGpuBuffer() can be +// made available even with MEDIAPIPE_DISABLE_GPU set. +#elif MEDIAPIPE_ENABLE_HALIDE + ASSIGN_OR_RETURN(cpu_converter_, + CreateFrameBufferConverter( + cc, GetBorderMode(options_.border_mode()), + GetOutputTensorType(/*uses_gpu=*/false, params_))); #else - LOG(FATAL) << "Cannot create image to tensor opencv converter since " - "MEDIAPIPE_DISABLE_OPENCV is defined."; -#endif // !MEDIAPIPE_DISABLE_OPENCV + LOG(FATAL) << "Cannot create image to tensor CPU converter since " + "MEDIAPIPE_DISABLE_OPENCV is defined and " + "MEDIAPIPE_ENABLE_HALIDE is not defined."; +#endif // !MEDIAPIPE_DISABLE_HALIDE } } return absl::OkStatus(); diff --git a/mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.cc b/mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.cc new file mode 100644 index 00000000..093f50d7 --- /dev/null +++ b/mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.cc @@ -0,0 +1,246 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.h" + +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_format.h" +#include "mediapipe/calculators/tensor/image_to_tensor_converter.h" +#include "mediapipe/calculators/tensor/image_to_tensor_utils.h" +#include "mediapipe/framework/calculator_context.h" +#include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/framework/formats/image.h" +#include "mediapipe/framework/formats/tensor.h" +#include "mediapipe/framework/port/status_macros.h" +#include "mediapipe/gpu/frame_buffer_view.h" +#include "mediapipe/util/frame_buffer/frame_buffer_util.h" + +namespace mediapipe { + +namespace { + +// Converts from radians (clockwise) to degrees (counter-clockwise) in [0,360). +int RadiansToDegrees(float radians) { + int degrees = static_cast(std::round(-radians * 180 / M_PI)) % 360; + if (degrees < 0) { + degrees += 360; + } + return degrees; +} + +// FrameBuffer-based implementation of ImageToTensorConverter. +class FrameBufferProcessor : public ImageToTensorConverter { + public: + FrameBufferProcessor(Tensor::ElementType tensor_type) + : tensor_type_(tensor_type) {} + + absl::Status Convert(const mediapipe::Image& input, const RotatedRect& roi, + float range_min, float range_max, + int tensor_buffer_offset, + Tensor& output_tensor) override; + + private: + absl::Status ValidateTensorShape(const Tensor::Shape& output_shape); + // Crops, rotates and resizes the input based on the provided + // region-of-interest. + absl::Status CropRotateResize90Degrees( + std::shared_ptr input, const RotatedRect& roi, + std::shared_ptr output); + // Converts the input FrameBuffer to a float Tensor. Output tensor must have + // type kFloat32. + absl::Status ConvertToFloatTensor( + std::shared_ptr input_frame, float range_min, + float range_max, Tensor& output_tensor); + + Tensor::ElementType tensor_type_; + + // Temporary buffers and their respective sizes. + std::unique_ptr cropped_buffer_; + size_t cropped_buffer_size_ = 0; + std::unique_ptr rotated_buffer_; + size_t rotated_buffer_size_ = 0; + std::unique_ptr output_buffer_; + size_t output_buffer_size_ = 0; +}; + +absl::Status FrameBufferProcessor::Convert(const mediapipe::Image& input, + const RotatedRect& roi, + float range_min, float range_max, + int tensor_buffer_offset, + Tensor& output_tensor) { + // TODO: add support for non-zero tensor buffer offset. + RET_CHECK_EQ(tensor_buffer_offset, 0) + << "Non-zero tensor_buffer_offset input is not supported yet."; + + // Range other than [0,255] is not supported for uint8 tensor outputs. + if (tensor_type_ == Tensor::ElementType::kUInt8) { + RET_CHECK(static_cast(range_min) == 0 && + static_cast(range_max) == 255); + } + + auto input_frame = input.GetGpuBuffer().GetReadView(); + const auto& output_shape = output_tensor.shape(); + MP_RETURN_IF_ERROR(ValidateTensorShape(output_shape)); + FrameBuffer::Dimension output_dimension{/*width=*/output_shape.dims[2], + /*height=*/output_shape.dims[1]}; + + // Optimized path for multiples of 90°. + if (RadiansToDegrees(roi.rotation) % 90 == 0) { + if (tensor_type_ == Tensor::ElementType::kUInt8) { + auto view = output_tensor.GetCpuWriteView(); + uint8_t* data = view.buffer(); + auto output_frame = + frame_buffer::CreateFromRgbRawBuffer(data, output_dimension); + return CropRotateResize90Degrees(input_frame, roi, output_frame); + } else { + size_t output_buffer_size = frame_buffer::GetFrameBufferByteSize( + output_dimension, FrameBuffer::Format::kRGB); + if (output_buffer_size > output_buffer_size_) { + output_buffer_ = std::make_unique(output_buffer_size); + output_buffer_size_ = output_buffer_size; + } + auto output_frame = frame_buffer::CreateFromRgbRawBuffer( + output_buffer_.get(), output_dimension); + MP_RETURN_IF_ERROR( + CropRotateResize90Degrees(input_frame, roi, output_frame)); + return ConvertToFloatTensor(output_frame, range_min, range_max, + output_tensor); + } + } else { + // TODO: add support for arbitrary rotations + return absl::UnimplementedError( + "FrameBufferConverter doesn't yet support rotations that are not " + "multiples of 90°."); + } + return absl::OkStatus(); +} + +absl::Status FrameBufferProcessor::ValidateTensorShape( + const Tensor::Shape& shape) { + RET_CHECK_EQ(shape.dims.size(), 4) + << "Wrong output dims size: " << shape.dims.size(); + RET_CHECK_EQ(shape.dims[0], 1) + << "Handling batch dimension not equal to 1 is not implemented in this " + "converter."; + RET_CHECK_EQ(shape.dims[3], 3) << "Wrong output channel: " << shape.dims[3]; + return absl::OkStatus(); +} + +absl::Status FrameBufferProcessor::CropRotateResize90Degrees( + std::shared_ptr input, const RotatedRect& roi, + std::shared_ptr output) { + int rotation_degrees = RadiansToDegrees(roi.rotation); + bool rotation_required = rotation_degrees != 0; + bool conversion_required = input->format() != output->format(); + + // First, crop and resize. + std::shared_ptr cropped = output; + FrameBuffer::Dimension cropped_dims = output->dimension(); + int left, right, top, bottom; + if (rotation_degrees % 180 != 0) { + cropped_dims.Swap(); + left = roi.center_x - roi.height / 2; + right = left + roi.height - 1; + top = roi.center_y - roi.width / 2; + bottom = top + roi.width - 1; + } else { + left = roi.center_x - roi.width / 2; + right = left + roi.width - 1; + top = roi.center_y - roi.height / 2; + bottom = top + roi.height - 1; + } + if (rotation_required || conversion_required) { + // Create temporary FrameBuffer from recycled buffer. + size_t cropped_buffer_size = + frame_buffer::GetFrameBufferByteSize(cropped_dims, input->format()); + if (cropped_buffer_size > cropped_buffer_size_) { + cropped_buffer_ = std::make_unique(cropped_buffer_size); + cropped_buffer_size_ = cropped_buffer_size; + } + ASSIGN_OR_RETURN(cropped, + frame_buffer::CreateFromRawBuffer( + cropped_buffer_.get(), cropped_dims, input->format())); + } + MP_RETURN_IF_ERROR( + frame_buffer::Crop(*input, left, top, right, bottom, cropped.get())); + + // Then rotate if needed. + std::shared_ptr rotated = output; + if (rotation_required) { + if (conversion_required) { + // Create temporary FrameBuffer from recycled buffer. + FrameBuffer::Dimension rotated_dims = output->dimension(); + size_t rotated_buffer_size = + frame_buffer::GetFrameBufferByteSize(rotated_dims, cropped->format()); + if (rotated_buffer_size > rotated_buffer_size_) { + rotated_buffer_ = std::make_unique(rotated_buffer_size); + rotated_buffer_size_ = rotated_buffer_size; + } + ASSIGN_OR_RETURN(auto rotated, frame_buffer::CreateFromRawBuffer( + rotated_buffer_.get(), rotated_dims, + cropped->format())); + } + MP_RETURN_IF_ERROR( + frame_buffer::Rotate(*cropped, rotation_degrees, rotated.get())); + } else { + rotated = cropped; + } + + // Then convert if needed. + if (conversion_required) { + return frame_buffer::Convert(*rotated, output.get()); + } + return absl::OkStatus(); +} + +absl::Status FrameBufferProcessor::ConvertToFloatTensor( + std::shared_ptr input_frame, float range_min, + float range_max, Tensor& output_tensor) { + RET_CHECK(output_tensor.element_type() == Tensor::ElementType::kFloat32); + constexpr float kInputImageRangeMin = 0.0f; + constexpr float kInputImageRangeMax = 255.0f; + ASSIGN_OR_RETURN(auto transform, GetValueRangeTransformation( + kInputImageRangeMin, kInputImageRangeMax, + range_min, range_max)); + return frame_buffer::ToFloatTensor(*input_frame, transform.scale, + transform.offset, output_tensor); +} + +} // namespace + +absl::StatusOr> +CreateFrameBufferConverter(CalculatorContext* cc, BorderMode border_mode, + Tensor::ElementType tensor_type) { + if (tensor_type != Tensor::ElementType::kUInt8 && + tensor_type != Tensor::ElementType::kFloat32) { + return absl::InvalidArgumentError( + absl::StrFormat("Tensor type is currently not supported by " + "FrameBufferProcessor, type: %d.", + tensor_type)); + } + // TODO: add support for BorderMode:kZero. + if (border_mode == BorderMode::kZero) { + return absl::UnimplementedError( + "BorderMode::kZero is not yet supported by FrameBufferProcessor"); + } + return std::make_unique(tensor_type); +} + +} // namespace mediapipe diff --git a/mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.h b/mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.h new file mode 100644 index 00000000..4bffdd72 --- /dev/null +++ b/mediapipe/calculators/tensor/image_to_tensor_converter_frame_buffer.h @@ -0,0 +1,34 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MEDIAPIPE_CALCULATORS_TENSOR_IMAGE_TO_TENSOR_CONVERTER_FRAME_BUFFER_H_ +#define MEDIAPIPE_CALCULATORS_TENSOR_IMAGE_TO_TENSOR_CONVERTER_FRAME_BUFFER_H_ + +#include + +#include "absl/status/statusor.h" +#include "mediapipe/calculators/tensor/image_to_tensor_converter.h" +#include "mediapipe/calculators/tensor/image_to_tensor_utils.h" +#include "mediapipe/framework/calculator_context.h" + +namespace mediapipe { + +// Creates FrameBuffer-based image-to-tensor converter relying on Halide. +absl::StatusOr> +CreateFrameBufferConverter(CalculatorContext* cc, BorderMode border_mode, + Tensor::ElementType tensor_type); + +} // namespace mediapipe + +#endif // MEDIAPIPE_CALCULATORS_TENSOR_IMAGE_TO_TENSOR_CONVERTER_FRAME_BUFFER_H_ diff --git a/mediapipe/framework/BUILD b/mediapipe/framework/BUILD index 19fa6f46..ae788ed5 100644 --- a/mediapipe/framework/BUILD +++ b/mediapipe/framework/BUILD @@ -969,6 +969,15 @@ cc_library( }) + select({ "//conditions:default": [], "//mediapipe/framework/port:disable_opencv": ["MEDIAPIPE_DISABLE_OPENCV=1"], + }) + select({ + "//conditions:default": [], + # TODO: Improve this. This only sets MEDIAPIPE_DISABLE_OPENCV as a "defines" Make + # value, not as a bazel "--define" variable, which has effects in C++ code but not in + # select() statements. + "//mediapipe/framework/port:enable_halide": [ + "MEDIAPIPE_ENABLE_HALIDE=1", + "MEDIAPIPE_DISABLE_OPENCV=1", + ], }) + select({ "//conditions:default": [], "//mediapipe/framework:disable_rtti_and_exceptions": [ diff --git a/mediapipe/framework/formats/image.h b/mediapipe/framework/formats/image.h index 44578a7d..ffb6362f 100644 --- a/mediapipe/framework/formats/image.h +++ b/mediapipe/framework/formats/image.h @@ -111,12 +111,15 @@ class Image { return gpu_buffer_.internal_storage(); } #endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER +#endif // !MEDIAPIPE_DISABLE_GPU + // Get a GPU view. Automatically uploads from CPU if needed. const mediapipe::GpuBuffer GetGpuBuffer() const { +#if !MEDIAPIPE_DISABLE_GPU if (use_gpu_ == false) ConvertToGpu(); +#endif // !MEDIAPIPE_DISABLE_GPU return gpu_buffer_; } -#endif // !MEDIAPIPE_DISABLE_GPU // Returns image properties. int width() const; diff --git a/mediapipe/framework/port/BUILD b/mediapipe/framework/port/BUILD index 1039dc1c..5b0a0fec 100644 --- a/mediapipe/framework/port/BUILD +++ b/mediapipe/framework/port/BUILD @@ -222,6 +222,18 @@ config_setting( }, ) +# Enable Halide when defining MEDIAPIPE_ENABLE_HALIDE to 1 in bazel command. +# This will automatically disable opencv, i.e. will set MEDIAPIPE_DISABLE_OPENCV to 1. +# Note that this only applies to a select few calculators/framework components currently. +# TODO: Improve this. This only sets MEDIAPIPE_DISABLE_OPENCV as a "defines" Make value, +# not as a bazel "--define" variable, which has effects in C++ code but not in select() statements. +config_setting( + name = "enable_halide", + define_values = { + "MEDIAPIPE_ENABLE_HALIDE": "1", + }, +) + cc_library( name = "opencv_core", hdrs = ["opencv_core_inc.h"],