Project import generated by Copybara.

GitOrigin-RevId: 53a42bf7ad836321123cb7b6c80b0f2e13fbf83e

mediapipe/calculators/image/BUILD

@@ -86,6 +86,14 @@ mediapipe_cc_proto_library(
     deps = [":opencv_image_encoder_calculator_proto"],
 )
 
+mediapipe_cc_proto_library(
+    name = "opencv_encoded_image_to_image_frame_calculator_cc_proto",
+    srcs = ["opencv_encoded_image_to_image_frame_calculator.proto"],
+    cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
+    visibility = ["//visibility:public"],
+    deps = [":opencv_encoded_image_to_image_frame_calculator_proto"],
+)
+
 mediapipe_cc_proto_library(
     name = "mask_overlay_calculator_cc_proto",
     srcs = ["mask_overlay_calculator.proto"],
@@ -172,6 +180,7 @@ cc_library(
     srcs = ["opencv_encoded_image_to_image_frame_calculator.cc"],
     visibility = ["//visibility:public"],
     deps = [
+        ":opencv_encoded_image_to_image_frame_calculator_cc_proto",
         "//mediapipe/framework:calculator_framework",
         "//mediapipe/framework/formats:image_frame_opencv",
         "//mediapipe/framework/port:opencv_imgcodecs",
@@ -557,6 +566,27 @@ proto_library(
     deps = ["//mediapipe/framework:calculator_proto"],
 )
 
+proto_library(
+    name = "opencv_encoded_image_to_image_frame_calculator_proto",
+    srcs = ["opencv_encoded_image_to_image_frame_calculator.proto"],
+    visibility = ["//visibility:public"],
+    deps = ["//mediapipe/framework:calculator_proto"],
+)
+
+proto_library(
+    name = "feature_detector_calculator_proto",
+    srcs = ["feature_detector_calculator.proto"],
+    deps = ["//mediapipe/framework:calculator_proto"],
+)
+
+mediapipe_cc_proto_library(
+    name = "feature_detector_calculator_cc_proto",
+    srcs = ["feature_detector_calculator.proto"],
+    cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
+    visibility = ["//visibility:public"],
+    deps = [":feature_detector_calculator_proto"],
+)
+
 cc_library(
     name = "mask_overlay_calculator",
     srcs = ["mask_overlay_calculator.cc"],
@@ -572,3 +602,30 @@ cc_library(
     ],
     alwayslink = 1,
 )
+
+cc_library(
+    name = "feature_detector_calculator",
+    srcs = ["feature_detector_calculator.cc"],
+    visibility = ["//mediapipe:__subpackages__"],
+    deps = [
+        ":feature_detector_calculator_cc_proto",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/formats:image_frame",
+        "//mediapipe/framework/formats:image_frame_opencv",
+        "//mediapipe/framework/formats:landmark_cc_proto",
+        "//mediapipe/framework/formats:video_stream_header",
+        "//mediapipe/framework/port:integral_types",
+        "//mediapipe/framework/port:logging",
+        "//mediapipe/framework/port:opencv_core",
+        "//mediapipe/framework/port:opencv_features2d",
+        "//mediapipe/framework/port:opencv_imgproc",
+        "//mediapipe/framework/port:ret_check",
+        "//mediapipe/framework/port:status",
+        "//mediapipe/framework/port:threadpool",
+        "//mediapipe/framework/tool:options_util",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/synchronization",
+        "@org_tensorflow//tensorflow/lite:framework",
+    ],
+    alwayslink = 1,
+)

mediapipe/calculators/image/feature_detector_calculator.cc

@@ -0,0 +1,210 @@
+// Copyright 2020 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 <memory>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "absl/synchronization/blocking_counter.h"
+#include "mediapipe/calculators/image/feature_detector_calculator.pb.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/image_frame.h"
+#include "mediapipe/framework/formats/image_frame_opencv.h"
+#include "mediapipe/framework/formats/landmark.pb.h"
+#include "mediapipe/framework/formats/video_stream_header.h"
+#include "mediapipe/framework/port/integral_types.h"
+#include "mediapipe/framework/port/logging.h"
+#include "mediapipe/framework/port/opencv_core_inc.h"
+#include "mediapipe/framework/port/opencv_features2d_inc.h"
+#include "mediapipe/framework/port/opencv_imgproc_inc.h"
+#include "mediapipe/framework/port/ret_check.h"
+#include "mediapipe/framework/port/status.h"
+#include "mediapipe/framework/port/threadpool.h"
+#include "mediapipe/framework/tool/options_util.h"
+#include "tensorflow/lite/interpreter.h"
+
+namespace mediapipe {
+
+const char kOptionsTag[] = "OPTIONS";
+const int kPatchSize = 32;
+const int kNumThreads = 16;
+
+// A calculator to apply local feature detection.
+// Input stream:
+//   IMAGE: Input image frame of type ImageFrame from video stream.
+// Output streams:
+//   FEATURES: The detected keypoints from input image as vector<cv::KeyPoint>.
+//   PATCHES:  Optional output the extracted patches as vector<cv::Mat>
+class FeatureDetectorCalculator : public CalculatorBase {
+ public:
+  ~FeatureDetectorCalculator() override = default;
+
+  static ::mediapipe::Status GetContract(CalculatorContract* cc);
+
+  ::mediapipe::Status Open(CalculatorContext* cc) override;
+  ::mediapipe::Status Process(CalculatorContext* cc) override;
+
+ private:
+  FeatureDetectorCalculatorOptions options_;
+  cv::Ptr<cv::Feature2D> feature_detector_;
+  std::unique_ptr<::mediapipe::ThreadPool> pool_;
+
+  // Create image pyramid based on input image.
+  void ComputeImagePyramid(const cv::Mat& input_image,
+                           std::vector<cv::Mat>* image_pyramid);
+
+  // Extract the patch for single feature with image pyramid.
+  cv::Mat ExtractPatch(const cv::KeyPoint& feature,
+                       const std::vector<cv::Mat>& image_pyramid);
+};
+
+REGISTER_CALCULATOR(FeatureDetectorCalculator);
+
+::mediapipe::Status FeatureDetectorCalculator::GetContract(
+    CalculatorContract* cc) {
+  if (cc->Inputs().HasTag("IMAGE")) {
+    cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
+  }
+  if (cc->Outputs().HasTag("FEATURES")) {
+    cc->Outputs().Tag("FEATURES").Set<std::vector<cv::KeyPoint>>();
+  }
+  if (cc->Outputs().HasTag("LANDMARKS")) {
+    cc->Outputs().Tag("LANDMARKS").Set<NormalizedLandmarkList>();
+  }
+  if (cc->Outputs().HasTag("PATCHES")) {
+    cc->Outputs().Tag("PATCHES").Set<std::vector<TfLiteTensor>>();
+  }
+  return ::mediapipe::OkStatus();
+}
+
+::mediapipe::Status FeatureDetectorCalculator::Open(CalculatorContext* cc) {
+  options_ =
+      tool::RetrieveOptions(cc->Options(), cc->InputSidePackets(), kOptionsTag)
+          .GetExtension(FeatureDetectorCalculatorOptions::ext);
+  feature_detector_ = cv::ORB::create(
+      options_.max_features(), options_.scale_factor(),
+      options_.pyramid_level(), kPatchSize - 1, 0, 2, cv::ORB::FAST_SCORE);
+  pool_ = absl::make_unique<::mediapipe::ThreadPool>("ThreadPool", kNumThreads);
+  pool_->StartWorkers();
+  return ::mediapipe::OkStatus();
+}
+
+::mediapipe::Status FeatureDetectorCalculator::Process(CalculatorContext* cc) {
+  const Timestamp& timestamp = cc->InputTimestamp();
+  if (timestamp == Timestamp::PreStream()) {
+    // Indicator packet.
+    return ::mediapipe::OkStatus();
+  }
+  InputStream* input_frame = &(cc->Inputs().Tag("IMAGE"));
+  cv::Mat input_view = formats::MatView(&input_frame->Get<ImageFrame>());
+  cv::Mat grayscale_view;
+  cv::cvtColor(input_view, grayscale_view, cv::COLOR_RGB2GRAY);
+
+  std::vector<cv::KeyPoint> keypoints;
+  feature_detector_->detect(grayscale_view, keypoints);
+  if (keypoints.size() > options_.max_features()) {
+    keypoints.resize(options_.max_features());
+  }
+
+  if (cc->Outputs().HasTag("FEATURES")) {
+    auto features_ptr = absl::make_unique<std::vector<cv::KeyPoint>>(keypoints);
+    cc->Outputs().Tag("FEATURES").Add(features_ptr.release(), timestamp);
+  }
+
+  if (cc->Outputs().HasTag("LANDMARKS")) {
+    auto landmarks_ptr = absl::make_unique<NormalizedLandmarkList>();
+    for (int j = 0; j < keypoints.size(); ++j) {
+      auto feature_landmark = landmarks_ptr->add_landmark();
+      feature_landmark->set_x(keypoints[j].pt.x / grayscale_view.cols);
+      feature_landmark->set_y(keypoints[j].pt.y / grayscale_view.rows);
+    }
+    cc->Outputs().Tag("LANDMARKS").Add(landmarks_ptr.release(), timestamp);
+  }
+
+  if (cc->Outputs().HasTag("PATCHES")) {
+    std::vector<cv::Mat> image_pyramid;
+    ComputeImagePyramid(grayscale_view, &image_pyramid);
+    std::vector<cv::Mat> patch_mat;
+    patch_mat.resize(keypoints.size());
+    absl::BlockingCounter counter(keypoints.size());
+    for (int i = 0; i < keypoints.size(); i++) {
+      pool_->Schedule(
+          [this, &image_pyramid, &keypoints, &patch_mat, i, &counter] {
+            patch_mat[i] = ExtractPatch(keypoints[i], image_pyramid);
+            counter.DecrementCount();
+          });
+    }
+    counter.Wait();
+    const int batch_size = options_.max_features();
+    auto patches = absl::make_unique<std::vector<TfLiteTensor>>();
+    TfLiteTensor tensor;
+    tensor.type = kTfLiteFloat32;
+    tensor.dims = TfLiteIntArrayCreate(4);
+    tensor.dims->data[0] = batch_size;
+    tensor.dims->data[1] = kPatchSize;
+    tensor.dims->data[2] = kPatchSize;
+    tensor.dims->data[3] = 1;
+    int num_bytes = batch_size * kPatchSize * kPatchSize * sizeof(float);
+    tensor.data.data = malloc(num_bytes);
+    tensor.bytes = num_bytes;
+    tensor.allocation_type = kTfLiteArenaRw;
+    float* tensor_buffer = tensor.data.f;
+    for (int i = 0; i < keypoints.size(); i++) {
+      for (int j = 0; j < patch_mat[i].rows; ++j) {
+        for (int k = 0; k < patch_mat[i].cols; ++k) {
+          *tensor_buffer++ = patch_mat[i].at<uchar>(j, k) / 128.0f - 1.0f;
+        }
+      }
+    }
+    for (int i = keypoints.size() * kPatchSize * kPatchSize; i < num_bytes / 4;
+         i++) {
+      *tensor_buffer++ = 0;
+    }
+
+    patches->emplace_back(tensor);
+    cc->Outputs().Tag("PATCHES").Add(patches.release(), timestamp);
+  }
+
+  return ::mediapipe::OkStatus();
+}
+
+void FeatureDetectorCalculator::ComputeImagePyramid(
+    const cv::Mat& input_image, std::vector<cv::Mat>* image_pyramid) {
+  cv::Mat tmp_image = input_image;
+  cv::Mat src_image = input_image;
+  for (int i = 0; i < options_.pyramid_level(); ++i) {
+    image_pyramid->push_back(src_image);
+    cv::resize(src_image, tmp_image, cv::Size(), 1.0f / options_.scale_factor(),
+               1.0f / options_.scale_factor());
+    src_image = tmp_image;
+  }
+}
+
+cv::Mat FeatureDetectorCalculator::ExtractPatch(
+    const cv::KeyPoint& feature, const std::vector<cv::Mat>& image_pyramid) {
+  cv::Mat img = image_pyramid[feature.octave];
+  float scale_factor = 1 / pow(options_.scale_factor(), feature.octave);
+  cv::Point2f center =
+      cv::Point2f(feature.pt.x * scale_factor, feature.pt.y * scale_factor);
+  cv::Mat rot = cv::getRotationMatrix2D(center, feature.angle, 1.0);
+  rot.at<double>(0, 2) += kPatchSize / 2 - center.x;
+  rot.at<double>(1, 2) += kPatchSize / 2 - center.y;
+  cv::Mat cropped_img;
+  // perform the affine transformation
+  cv::warpAffine(img, cropped_img, rot, cv::Size(kPatchSize, kPatchSize),
+                 cv::INTER_LINEAR);
+  return cropped_img;
+}
+
+}  // namespace mediapipe

mediapipe/calculators/image/feature_detector_calculator.proto

@@ -0,0 +1,24 @@
+// Options for FeatureDetectorCalculator
+syntax = "proto2";
+
+package mediapipe;
+
+import "mediapipe/framework/calculator.proto";
+
+message FeatureDetectorCalculatorOptions {
+  extend CalculatorOptions {
+    optional FeatureDetectorCalculatorOptions ext = 278741680;
+  }
+
+  // Set to true if output patches, otherwise only output cv::KeyPoint
+  optional bool output_patch = 1;
+
+  // The max number of detected features.
+  optional int32 max_features = 2 [default = 200];
+
+  // The number of pyramid levels.
+  optional int32 pyramid_level = 3 [default = 4];
+
+  // Pyramid decimation ratio.
+  optional float scale_factor = 4 [default = 1.2];
+}

mediapipe/calculators/image/image_cropping_calculator.cc

@@ -219,8 +219,10 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
   const auto& input_img = cc->Inputs().Tag(kImageTag).Get<ImageFrame>();
   cv::Mat input_mat = formats::MatView(&input_img);
 
-  auto [target_width, target_height, rect_center_x, rect_center_y, rotation] =
-      GetCropSpecs(cc, input_img.Width(), input_img.Height());
+  RectSpec specs = GetCropSpecs(cc, input_img.Width(), input_img.Height());
+  int target_width = specs.width, target_height = specs.height,
+      rect_center_x = specs.center_x, rect_center_y = specs.center_y;
+  float rotation = specs.rotation;
 
   // Get border mode and value for OpenCV.
   int border_mode;
@@ -403,8 +405,10 @@ void ImageCroppingCalculator::GetOutputDimensions(CalculatorContext* cc,
                                                   int src_width, int src_height,
                                                   int* dst_width,
                                                   int* dst_height) {
-  auto [crop_width, crop_height, x_center, y_center, rotation] =
-      GetCropSpecs(cc, src_width, src_height);
+  RectSpec specs = GetCropSpecs(cc, src_width, src_height);
+  int crop_width = specs.width, crop_height = specs.height,
+      x_center = specs.center_x, y_center = specs.center_y;
+  float rotation = specs.rotation;
 
   const float half_width = crop_width / 2.0f;
   const float half_height = crop_height / 2.0f;

mediapipe/calculators/image/opencv_encoded_image_to_image_frame_calculator.cc

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "mediapipe/calculators/image/opencv_encoded_image_to_image_frame_calculator.pb.h"
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/formats/image_frame_opencv.h"
 #include "mediapipe/framework/port/opencv_imgcodecs_inc.h"
@@ -34,7 +35,11 @@ namespace mediapipe {
 class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
  public:
   static ::mediapipe::Status GetContract(CalculatorContract* cc);
+  ::mediapipe::Status Open(CalculatorContext* cc) override;
   ::mediapipe::Status Process(CalculatorContext* cc) override;
+
+ private:
+  mediapipe::OpenCvEncodedImageToImageFrameCalculatorOptions options_;
 };
 
 ::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::GetContract(
@@ -44,13 +49,29 @@ class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
   return ::mediapipe::OkStatus();
 }
 
+::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::Open(
+    CalculatorContext* cc) {
+  options_ =
+      cc->Options<mediapipe::OpenCvEncodedImageToImageFrameCalculatorOptions>();
+  return ::mediapipe::OkStatus();
+}
+
 ::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::Process(
     CalculatorContext* cc) {
   const std::string& contents = cc->Inputs().Index(0).Get<std::string>();
   const std::vector<char> contents_vector(contents.begin(), contents.end());
-  cv::Mat decoded_mat =
-      cv::imdecode(contents_vector, -1 /* return the loaded image as-is */);
-
+  cv::Mat decoded_mat;
+  if (options_.apply_orientation_from_exif_data()) {
+    // We want to respect the orientation from the EXIF data, which
+    // IMREAD_UNCHANGED ignores, but otherwise we want to be as permissive as
+    // possible with our reading flags. Therefore, we use IMREAD_ANYCOLOR and
+    // IMREAD_ANYDEPTH.
+    decoded_mat = cv::imdecode(contents_vector,
+                               cv::IMREAD_ANYCOLOR | cv::IMREAD_ANYDEPTH);
+  } else {
+    // Return the loaded image as-is
+    decoded_mat = cv::imdecode(contents_vector, cv::IMREAD_UNCHANGED);
+  }
   ImageFormat::Format image_format = ImageFormat::UNKNOWN;
   cv::Mat output_mat;
   switch (decoded_mat.channels()) {
@@ -70,7 +91,8 @@ class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
              << "Unsupported number of channels: " << decoded_mat.channels();
   }
   std::unique_ptr<ImageFrame> output_frame = absl::make_unique<ImageFrame>(
-      image_format, decoded_mat.size().width, decoded_mat.size().height);
+      image_format, decoded_mat.size().width, decoded_mat.size().height,
+      ImageFrame::kGlDefaultAlignmentBoundary);
   output_mat.copyTo(formats::MatView(output_frame.get()));
   cc->Outputs().Index(0).Add(output_frame.release(), cc->InputTimestamp());
   return ::mediapipe::OkStatus();

mediapipe/calculators/image/opencv_encoded_image_to_image_frame_calculator.proto

@@ -0,0 +1,30 @@
+// Copyright 2020 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.
+
+syntax = "proto2";
+
+package mediapipe;
+
+import "mediapipe/framework/calculator.proto";
+
+message OpenCvEncodedImageToImageFrameCalculatorOptions {
+  extend CalculatorOptions {
+    optional OpenCvEncodedImageToImageFrameCalculatorOptions ext = 303447308;
+  }
+
+  // If set, we will attempt to automatically apply the orientation specified by
+  // the image's EXIF data when loading the image. Otherwise, the image data
+  // will be loaded as-is.
+  optional bool apply_orientation_from_exif_data = 1 [default = false];
+}