Internal update.

PiperOrigin-RevId: 561184322

mediapipe/calculators/util/BUILD

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

mediapipe/calculators/util/annotation_overlay_calculator.cc

@@ -18,9 +18,11 @@
 #include "mediapipe/calculators/util/annotation_overlay_calculator.pb.h"
 #include "mediapipe/framework/calculator_framework.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_frame.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/port/logging.h"
 #include "mediapipe/framework/port/opencv_core_inc.h"
@@ -45,6 +47,7 @@ namespace {
 constexpr char kVectorTag[] = "VECTOR";
 constexpr char kGpuBufferTag[] = "IMAGE_GPU";
 constexpr char kImageFrameTag[] = "IMAGE";
+constexpr char kImageTag[] = "UIMAGE";  // Universal Image
 
 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.
 
 // Future Image type.
-inline bool HasImageTag(mediapipe::CalculatorContext* cc) { return false; }
+inline bool HasImageTag(mediapipe::CalculatorContext* cc) {
+  return cc->Inputs().HasTag(kImageTag);
+}
 }  // namespace
 
 // A calculator for rendering data on images.
 //
 // Inputs:
 //  1. IMAGE or IMAGE_GPU (optional): An ImageFrame (or GpuBuffer),
+//     or UIMAGE (an Image).
 //     containing the input image.
 //     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.
@@ -76,6 +82,7 @@ inline bool HasImageTag(mediapipe::CalculatorContext* cc) { return false; }
 //
 // Output:
 //  1. IMAGE or IMAGE_GPU: A rendered ImageFrame (or GpuBuffer),
+//     or UIMAGE (an Image).
 //  Note: Output types should match their corresponding input stream type.
 //
 // 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,
                                      std::unique_ptr<cv::Mat>& image_mat,
                                      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>
   absl::Status CreateRenderTargetGpu(CalculatorContext* cc,
                                      std::unique_ptr<cv::Mat>& image_mat);
@@ -176,14 +186,14 @@ absl::Status AnnotationOverlayCalculator::GetContract(CalculatorContract* cc) {
 
   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.
 #if !MEDIAPIPE_DISABLE_GPU
@@ -198,6 +208,14 @@ absl::Status AnnotationOverlayCalculator::GetContract(CalculatorContract* cc) {
     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.
   for (CollectionItemId id = cc->Inputs().BeginId(); id < cc->Inputs().EndId();
        ++id) {
@@ -220,6 +238,9 @@ absl::Status AnnotationOverlayCalculator::GetContract(CalculatorContract* cc) {
   if (cc->Outputs().HasTag(kImageFrameTag)) {
     cc->Outputs().Tag(kImageFrameTag).Set<ImageFrame>();
   }
+  if (cc->Outputs().HasTag(kImageTag)) {
+    cc->Outputs().Tag(kImageTag).Set<mediapipe::Image>();
+  }
 
   if (use_gpu) {
 #if !MEDIAPIPE_DISABLE_GPU
@@ -252,9 +273,13 @@ absl::Status AnnotationOverlayCalculator::Open(CalculatorContext* cc) {
   renderer_ = absl::make_unique<AnnotationRenderer>();
   renderer_->SetFlipTextVertically(options_.flip_text_vertically());
   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).
-  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()) {
     const auto& input_header =
         cc->Inputs().Tag(tag).Header().Get<VideoHeader>();
@@ -280,6 +305,12 @@ absl::Status AnnotationOverlayCalculator::Process(CalculatorContext* cc) {
       cc->Inputs().Tag(kImageFrameTag).IsEmpty()) {
     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.
   std::unique_ptr<cv::Mat> image_mat;
@@ -289,10 +320,17 @@ absl::Status AnnotationOverlayCalculator::Process(CalculatorContext* cc) {
     if (!gpu_initialized_) {
       MP_RETURN_IF_ERROR(
           gpu_helper_.RunInGlContext([this, cc]() -> absl::Status {
+            if (HasImageTag(cc)) {
+              return GlSetup<mediapipe::Image, kImageTag>(cc);
+            }
             return GlSetup<mediapipe::GpuBuffer, kGpuBufferTag>(cc);
           }));
       gpu_initialized_ = true;
     }
+    if (HasImageTag(cc)) {
+      MP_RETURN_IF_ERROR(
+          (CreateRenderTargetGpu<mediapipe::Image, kImageTag>(cc, image_mat)));
+    }
     if (cc->Inputs().HasTag(kGpuBufferTag)) {
       MP_RETURN_IF_ERROR(
           (CreateRenderTargetGpu<mediapipe::GpuBuffer, kGpuBufferTag>(
@@ -300,6 +338,10 @@ absl::Status AnnotationOverlayCalculator::Process(CalculatorContext* cc) {
     }
 #endif  // !MEDIAPIPE_DISABLE_GPU
   } else {
+    if (cc->Outputs().HasTag(kImageTag)) {
+      MP_RETURN_IF_ERROR(
+          CreateRenderTargetCpuImage(cc, image_mat, &target_format));
+    }
     if (cc->Outputs().HasTag(kImageFrameTag)) {
       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;
     MP_RETURN_IF_ERROR(
         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>(
               cc, image_mat_ptr);
         }));
@@ -381,6 +426,10 @@ absl::Status AnnotationOverlayCalculator::RenderToCpu(
                               ImageFrame::kDefaultAlignmentBoundary);
 #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)) {
     cc->Outputs()
         .Tag(kImageFrameTag)
@@ -487,6 +536,54 @@ absl::Status AnnotationOverlayCalculator::CreateRenderTargetCpu(
   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>
 absl::Status AnnotationOverlayCalculator::CreateRenderTargetGpu(
     CalculatorContext* cc, std::unique_ptr<cv::Mat>& image_mat) {