Internal update.

PiperOrigin-RevId: 561184322
This commit is contained in:
MediaPipe Team 2023-08-29 17:32:04 -07:00 committed by Copybara-Service
parent e18e749e3e
commit 6c2638592e
2 changed files with 109 additions and 10 deletions

View File

@ -248,9 +248,11 @@ cc_library(
":annotation_overlay_calculator_cc_proto", ":annotation_overlay_calculator_cc_proto",
"//mediapipe/framework:calculator_framework", "//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_options_cc_proto", "//mediapipe/framework:calculator_options_cc_proto",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:image_format_cc_proto", "//mediapipe/framework/formats:image_format_cc_proto",
"//mediapipe/framework/formats:image_frame", "//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:image_frame_opencv", "//mediapipe/framework/formats:image_frame_opencv",
"//mediapipe/framework/formats:image_opencv",
"//mediapipe/framework/formats:video_stream_header", "//mediapipe/framework/formats:video_stream_header",
"//mediapipe/framework/port:logging", "//mediapipe/framework/port:logging",
"//mediapipe/framework/port:opencv_core", "//mediapipe/framework/port:opencv_core",

View File

@ -18,9 +18,11 @@
#include "mediapipe/calculators/util/annotation_overlay_calculator.pb.h" #include "mediapipe/calculators/util/annotation_overlay_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h" #include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_options.pb.h" #include "mediapipe/framework/calculator_options.pb.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/image_format.pb.h" #include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/formats/image_frame.h" #include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_frame_opencv.h" #include "mediapipe/framework/formats/image_frame_opencv.h"
#include "mediapipe/framework/formats/image_opencv.h"
#include "mediapipe/framework/formats/video_stream_header.h" #include "mediapipe/framework/formats/video_stream_header.h"
#include "mediapipe/framework/port/logging.h" #include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/opencv_core_inc.h" #include "mediapipe/framework/port/opencv_core_inc.h"
@ -45,6 +47,7 @@ namespace {
constexpr char kVectorTag[] = "VECTOR"; constexpr char kVectorTag[] = "VECTOR";
constexpr char kGpuBufferTag[] = "IMAGE_GPU"; constexpr char kGpuBufferTag[] = "IMAGE_GPU";
constexpr char kImageFrameTag[] = "IMAGE"; constexpr char kImageFrameTag[] = "IMAGE";
constexpr char kImageTag[] = "UIMAGE"; // Universal Image
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES }; enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
@ -57,13 +60,16 @@ size_t RoundUp(size_t n, size_t m) { return ((n + m - 1) / m) * m; } // NOLINT
constexpr uchar kAnnotationBackgroundColor = 2; // Grayscale value. constexpr uchar kAnnotationBackgroundColor = 2; // Grayscale value.
// Future Image type. // Future Image type.
inline bool HasImageTag(mediapipe::CalculatorContext* cc) { return false; } inline bool HasImageTag(mediapipe::CalculatorContext* cc) {
return cc->Inputs().HasTag(kImageTag);
}
} // namespace } // namespace
// A calculator for rendering data on images. // A calculator for rendering data on images.
// //
// Inputs: // Inputs:
// 1. IMAGE or IMAGE_GPU (optional): An ImageFrame (or GpuBuffer), // 1. IMAGE or IMAGE_GPU (optional): An ImageFrame (or GpuBuffer),
// or UIMAGE (an Image).
// containing the input image. // containing the input image.
// If output is CPU, and input isn't provided, the renderer creates a // If output is CPU, and input isn't provided, the renderer creates a
// blank canvas with the width, height and color provided in the options. // blank canvas with the width, height and color provided in the options.
@ -76,6 +82,7 @@ inline bool HasImageTag(mediapipe::CalculatorContext* cc) { return false; }
// //
// Output: // Output:
// 1. IMAGE or IMAGE_GPU: A rendered ImageFrame (or GpuBuffer), // 1. IMAGE or IMAGE_GPU: A rendered ImageFrame (or GpuBuffer),
// or UIMAGE (an Image).
// Note: Output types should match their corresponding input stream type. // Note: Output types should match their corresponding input stream type.
// //
// For CPU input frames, only SRGBA, SRGB and GRAY8 format are supported. The // For CPU input frames, only SRGBA, SRGB and GRAY8 format are supported. The
@ -135,6 +142,9 @@ class AnnotationOverlayCalculator : public CalculatorBase {
absl::Status CreateRenderTargetCpu(CalculatorContext* cc, absl::Status CreateRenderTargetCpu(CalculatorContext* cc,
std::unique_ptr<cv::Mat>& image_mat, std::unique_ptr<cv::Mat>& image_mat,
ImageFormat::Format* target_format); ImageFormat::Format* target_format);
absl::Status CreateRenderTargetCpuImage(CalculatorContext* cc,
std::unique_ptr<cv::Mat>& image_mat,
ImageFormat::Format* target_format);
template <typename Type, const char* Tag> template <typename Type, const char* Tag>
absl::Status CreateRenderTargetGpu(CalculatorContext* cc, absl::Status CreateRenderTargetGpu(CalculatorContext* cc,
std::unique_ptr<cv::Mat>& image_mat); std::unique_ptr<cv::Mat>& image_mat);
@ -176,14 +186,14 @@ absl::Status AnnotationOverlayCalculator::GetContract(CalculatorContract* cc) {
bool use_gpu = false; bool use_gpu = false;
if (cc->Inputs().HasTag(kImageFrameTag) && RET_CHECK(cc->Inputs().HasTag(kImageFrameTag) +
cc->Inputs().HasTag(kGpuBufferTag)) { cc->Inputs().HasTag(kGpuBufferTag) +
return absl::InternalError("Cannot have multiple input images."); cc->Inputs().HasTag(kImageTag) <=
} 1);
if (cc->Inputs().HasTag(kGpuBufferTag) != RET_CHECK(cc->Outputs().HasTag(kImageFrameTag) +
cc->Outputs().HasTag(kGpuBufferTag)) { cc->Outputs().HasTag(kGpuBufferTag) +
return absl::InternalError("GPU output must have GPU input."); cc->Outputs().HasTag(kImageTag) ==
} 1);
// Input image to render onto copy of. Should be same type as output. // Input image to render onto copy of. Should be same type as output.
#if !MEDIAPIPE_DISABLE_GPU #if !MEDIAPIPE_DISABLE_GPU
@ -198,6 +208,14 @@ absl::Status AnnotationOverlayCalculator::GetContract(CalculatorContract* cc) {
RET_CHECK(cc->Outputs().HasTag(kImageFrameTag)); RET_CHECK(cc->Outputs().HasTag(kImageFrameTag));
} }
if (cc->Inputs().HasTag(kImageTag)) {
cc->Inputs().Tag(kImageTag).Set<mediapipe::Image>();
RET_CHECK(cc->Outputs().HasTag(kImageTag));
#if !MEDIAPIPE_DISABLE_GPU
use_gpu = true; // Prepare GPU resources because images can come in on GPU.
#endif
}
// Data streams to render. // Data streams to render.
for (CollectionItemId id = cc->Inputs().BeginId(); id < cc->Inputs().EndId(); for (CollectionItemId id = cc->Inputs().BeginId(); id < cc->Inputs().EndId();
++id) { ++id) {
@ -220,6 +238,9 @@ absl::Status AnnotationOverlayCalculator::GetContract(CalculatorContract* cc) {
if (cc->Outputs().HasTag(kImageFrameTag)) { if (cc->Outputs().HasTag(kImageFrameTag)) {
cc->Outputs().Tag(kImageFrameTag).Set<ImageFrame>(); cc->Outputs().Tag(kImageFrameTag).Set<ImageFrame>();
} }
if (cc->Outputs().HasTag(kImageTag)) {
cc->Outputs().Tag(kImageTag).Set<mediapipe::Image>();
}
if (use_gpu) { if (use_gpu) {
#if !MEDIAPIPE_DISABLE_GPU #if !MEDIAPIPE_DISABLE_GPU
@ -252,9 +273,13 @@ absl::Status AnnotationOverlayCalculator::Open(CalculatorContext* cc) {
renderer_ = absl::make_unique<AnnotationRenderer>(); renderer_ = absl::make_unique<AnnotationRenderer>();
renderer_->SetFlipTextVertically(options_.flip_text_vertically()); renderer_->SetFlipTextVertically(options_.flip_text_vertically());
if (use_gpu_) renderer_->SetScaleFactor(options_.gpu_scale_factor()); if (use_gpu_) renderer_->SetScaleFactor(options_.gpu_scale_factor());
if (renderer_->GetScaleFactor() < 1.0 && HasImageTag(cc))
LOG(WARNING) << "Annotation scale factor only supports GPU backed Image.";
// Set the output header based on the input header (if present). // Set the output header based on the input header (if present).
const char* tag = use_gpu_ ? kGpuBufferTag : kImageFrameTag; const char* tag = HasImageTag(cc) ? kImageTag
: use_gpu_ ? kGpuBufferTag
: kImageFrameTag;
if (image_frame_available_ && !cc->Inputs().Tag(tag).Header().IsEmpty()) { if (image_frame_available_ && !cc->Inputs().Tag(tag).Header().IsEmpty()) {
const auto& input_header = const auto& input_header =
cc->Inputs().Tag(tag).Header().Get<VideoHeader>(); cc->Inputs().Tag(tag).Header().Get<VideoHeader>();
@ -280,6 +305,12 @@ absl::Status AnnotationOverlayCalculator::Process(CalculatorContext* cc) {
cc->Inputs().Tag(kImageFrameTag).IsEmpty()) { cc->Inputs().Tag(kImageFrameTag).IsEmpty()) {
return absl::OkStatus(); return absl::OkStatus();
} }
if (cc->Inputs().HasTag(kImageTag) && cc->Inputs().Tag(kImageTag).IsEmpty()) {
return absl::OkStatus();
}
if (HasImageTag(cc)) {
use_gpu_ = cc->Inputs().Tag(kImageTag).Get<mediapipe::Image>().UsesGpu();
}
// Initialize render target, drawn with OpenCV. // Initialize render target, drawn with OpenCV.
std::unique_ptr<cv::Mat> image_mat; std::unique_ptr<cv::Mat> image_mat;
@ -289,10 +320,17 @@ absl::Status AnnotationOverlayCalculator::Process(CalculatorContext* cc) {
if (!gpu_initialized_) { if (!gpu_initialized_) {
MP_RETURN_IF_ERROR( MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, cc]() -> absl::Status { gpu_helper_.RunInGlContext([this, cc]() -> absl::Status {
if (HasImageTag(cc)) {
return GlSetup<mediapipe::Image, kImageTag>(cc);
}
return GlSetup<mediapipe::GpuBuffer, kGpuBufferTag>(cc); return GlSetup<mediapipe::GpuBuffer, kGpuBufferTag>(cc);
})); }));
gpu_initialized_ = true; gpu_initialized_ = true;
} }
if (HasImageTag(cc)) {
MP_RETURN_IF_ERROR(
(CreateRenderTargetGpu<mediapipe::Image, kImageTag>(cc, image_mat)));
}
if (cc->Inputs().HasTag(kGpuBufferTag)) { if (cc->Inputs().HasTag(kGpuBufferTag)) {
MP_RETURN_IF_ERROR( MP_RETURN_IF_ERROR(
(CreateRenderTargetGpu<mediapipe::GpuBuffer, kGpuBufferTag>( (CreateRenderTargetGpu<mediapipe::GpuBuffer, kGpuBufferTag>(
@ -300,6 +338,10 @@ absl::Status AnnotationOverlayCalculator::Process(CalculatorContext* cc) {
} }
#endif // !MEDIAPIPE_DISABLE_GPU #endif // !MEDIAPIPE_DISABLE_GPU
} else { } else {
if (cc->Outputs().HasTag(kImageTag)) {
MP_RETURN_IF_ERROR(
CreateRenderTargetCpuImage(cc, image_mat, &target_format));
}
if (cc->Outputs().HasTag(kImageFrameTag)) { if (cc->Outputs().HasTag(kImageFrameTag)) {
MP_RETURN_IF_ERROR(CreateRenderTargetCpu(cc, image_mat, &target_format)); MP_RETURN_IF_ERROR(CreateRenderTargetCpu(cc, image_mat, &target_format));
} }
@ -339,6 +381,9 @@ absl::Status AnnotationOverlayCalculator::Process(CalculatorContext* cc) {
uchar* image_mat_ptr = image_mat->data; uchar* image_mat_ptr = image_mat->data;
MP_RETURN_IF_ERROR( MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, cc, image_mat_ptr]() -> absl::Status { gpu_helper_.RunInGlContext([this, cc, image_mat_ptr]() -> absl::Status {
if (HasImageTag(cc)) {
return RenderToGpu<mediapipe::Image, kImageTag>(cc, image_mat_ptr);
}
return RenderToGpu<mediapipe::GpuBuffer, kGpuBufferTag>( return RenderToGpu<mediapipe::GpuBuffer, kGpuBufferTag>(
cc, image_mat_ptr); cc, image_mat_ptr);
})); }));
@ -381,6 +426,10 @@ absl::Status AnnotationOverlayCalculator::RenderToCpu(
ImageFrame::kDefaultAlignmentBoundary); ImageFrame::kDefaultAlignmentBoundary);
#endif // !MEDIAPIPE_DISABLE_GPU #endif // !MEDIAPIPE_DISABLE_GPU
if (HasImageTag(cc)) {
auto out = std::make_unique<mediapipe::Image>(std::move(output_frame));
cc->Outputs().Tag(kImageTag).Add(out.release(), cc->InputTimestamp());
}
if (cc->Outputs().HasTag(kImageFrameTag)) { if (cc->Outputs().HasTag(kImageFrameTag)) {
cc->Outputs() cc->Outputs()
.Tag(kImageFrameTag) .Tag(kImageFrameTag)
@ -487,6 +536,54 @@ absl::Status AnnotationOverlayCalculator::CreateRenderTargetCpu(
return absl::OkStatus(); return absl::OkStatus();
} }
absl::Status AnnotationOverlayCalculator::CreateRenderTargetCpuImage(
CalculatorContext* cc, std::unique_ptr<cv::Mat>& image_mat,
ImageFormat::Format* target_format) {
if (image_frame_available_) {
const auto& input_frame =
cc->Inputs().Tag(kImageTag).Get<mediapipe::Image>();
int target_mat_type;
switch (input_frame.image_format()) {
case ImageFormat::SRGBA:
*target_format = ImageFormat::SRGBA;
target_mat_type = CV_8UC4;
break;
case ImageFormat::SRGB:
*target_format = ImageFormat::SRGB;
target_mat_type = CV_8UC3;
break;
case ImageFormat::GRAY8:
*target_format = ImageFormat::SRGB;
target_mat_type = CV_8UC3;
break;
default:
return absl::UnknownError("Unexpected image frame format.");
break;
}
image_mat = absl::make_unique<cv::Mat>(
input_frame.height(), input_frame.width(), target_mat_type);
auto input_mat = formats::MatView(&input_frame);
if (input_frame.image_format() == ImageFormat::GRAY8) {
cv::Mat rgb_mat;
cv::cvtColor(*input_mat, rgb_mat, cv::COLOR_GRAY2RGB);
rgb_mat.copyTo(*image_mat);
} else {
input_mat->copyTo(*image_mat);
}
} else {
image_mat = absl::make_unique<cv::Mat>(
options_.canvas_height_px(), options_.canvas_width_px(), CV_8UC3,
cv::Scalar(options_.canvas_color().r(), options_.canvas_color().g(),
options_.canvas_color().b()));
*target_format = ImageFormat::SRGB;
}
return absl::OkStatus();
}
template <typename Type, const char* Tag> template <typename Type, const char* Tag>
absl::Status AnnotationOverlayCalculator::CreateRenderTargetGpu( absl::Status AnnotationOverlayCalculator::CreateRenderTargetGpu(
CalculatorContext* cc, std::unique_ptr<cv::Mat>& image_mat) { CalculatorContext* cc, std::unique_ptr<cv::Mat>& image_mat) {