Internal update.
PiperOrigin-RevId: 561184322
This commit is contained in:
parent
e18e749e3e
commit
6c2638592e
|
@ -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",
|
||||||
|
|
|
@ -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) {
|
||||||
|
|
Loading…
Reference in New Issue
Block a user