Project import generated by Copybara.

PiperOrigin-RevId: 263889205

mediapipe/calculators/video/BUILD

@@ -113,6 +113,22 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "tvl1_optical_flow_calculator",
+    srcs = ["tvl1_optical_flow_calculator.cc"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/formats:image_frame",
+        "//mediapipe/framework/formats:image_frame_opencv",
+        "//mediapipe/framework/formats/motion:optical_flow_field",
+        "//mediapipe/framework/port:opencv_video",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/synchronization",
+    ],
+    alwayslink = 1,
+)
+
 cc_test(
     name = "opencv_video_decoder_calculator_test",
     srcs = ["opencv_video_decoder_calculator_test.cc"],
@@ -155,3 +171,22 @@ cc_test(
         "//mediapipe/framework/port:parse_text_proto",
     ],
 )
+
+cc_test(
+    name = "tvl1_optical_flow_calculator_test",
+    srcs = ["tvl1_optical_flow_calculator_test.cc"],
+    data = ["//mediapipe/calculators/image/testdata:test_images"],
+    deps = [
+        ":tvl1_optical_flow_calculator",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:calculator_runner",
+        "//mediapipe/framework/deps:file_path",
+        "//mediapipe/framework/formats:image_frame_opencv",
+        "//mediapipe/framework/formats/motion:optical_flow_field",
+        "//mediapipe/framework/port:file_helpers",
+        "//mediapipe/framework/port:gtest_main",
+        "//mediapipe/framework/port:opencv_imgcodecs",
+        "//mediapipe/framework/port:opencv_imgproc",
+        "//mediapipe/framework/port:parse_text_proto",
+    ],
+)

mediapipe/calculators/video/opencv_video_encoder_calculator.cc

@@ -139,9 +139,9 @@ class OpenCvVideoEncoderCalculator : public CalculatorBase {
              << " in OpenCvVideoEncoderCalculator::Process()";
     }
     if (format == ImageFormat::SRGB) {
-      cv::cvtColor(tmp_frame, frame, cv::COLOR_BGR2RGB);
+      cv::cvtColor(tmp_frame, frame, cv::COLOR_RGB2BGR);
     } else if (format == ImageFormat::SRGBA) {
-      cv::cvtColor(tmp_frame, frame, cv::COLOR_BGRA2RGBA);
+      cv::cvtColor(tmp_frame, frame, cv::COLOR_RGBA2BGR);
     } else {
       return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
              << "Unsupported image format: " << format;

mediapipe/calculators/video/tvl1_optical_flow_calculator.cc

@@ -0,0 +1,191 @@
+// Copyright 2019 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 "absl/base/macros.h"
+#include "absl/synchronization/mutex.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/motion/optical_flow_field.h"
+#include "mediapipe/framework/port/opencv_video_inc.h"
+
+namespace mediapipe {
+namespace {
+
+// Checks that img1 and img2 have the same dimensions.
+bool ImageSizesMatch(const ImageFrame& img1, const ImageFrame& img2) {
+  return (img1.Width() == img2.Width()) && (img1.Height() == img2.Height());
+}
+
+// Converts an RGB image to grayscale.
+cv::Mat ConvertToGrayscale(const cv::Mat& image) {
+  if (image.channels() == 1) {
+    return image;
+  }
+  cv::Mat gray;
+  cv::cvtColor(image, gray, cv::COLOR_RGB2GRAY);
+  return gray;
+}
+
+}  // namespace
+
+// Calls OpenCV's DenseOpticalFlow to compute the optical flow between a pair of
+// image frames. The calculator can output forward flow fields (optical flow
+// from the first frame to the second frame), backward flow fields (optical flow
+// from the second frame to the first frame), or both, depending on the tag of
+// the specified output streams. Note that the timestamp of the output optical
+// flow is always tied to the input timestamp. Be aware of the different
+// meanings of the timestamp between the forward and the backward optical flows
+// if the calculator outputs both.
+//
+// If the "max_in_flight" field is set to any value greater than 1, it will
+// enable the calculator to process multiple inputs in parallel. The output
+// packets will be automatically ordered by timestamp before they are passed
+// along to downstream calculators.
+//
+// Inputs:
+//   FIRST_FRAME: An ImageFrame in either SRGB or GRAY8 format.
+//   SECOND_FRAME: An ImageFrame in either SRGB or GRAY8 format.
+// Outputs:
+//   FORWARD_FLOW: The OpticalFlowField from the first frame to the second
+//                 frame, output at the input timestamp.
+//   BACKWARD_FLOW: The OpticalFlowField from the second frame to the first
+//                  frame, output at the input timestamp.
+// Example config:
+//   node {
+//     calculator: "Tvl1OpticalFlowCalculator"
+//     input_stream: "FIRST_FRAME:first_frames"
+//     input_stream: "SECOND_FRAME:second_frames"
+//     output_stream: "FORWARD_FLOW:forward_flow"
+//     output_stream: "BACKWARD_FLOW:backward_flow"
+//     max_in_flight: 10
+//   }
+//   num_threads: 10
+class Tvl1OpticalFlowCalculator : public CalculatorBase {
+ public:
+  static ::mediapipe::Status GetContract(CalculatorContract* cc);
+  ::mediapipe::Status Open(CalculatorContext* cc) override;
+  ::mediapipe::Status Process(CalculatorContext* cc) override;
+
+ private:
+  ::mediapipe::Status CalculateOpticalFlow(const ImageFrame& current_frame,
+                                           const ImageFrame& next_frame,
+                                           OpticalFlowField* flow);
+  bool forward_requested_ = false;
+  bool backward_requested_ = false;
+  // Stores the idle DenseOpticalFlow objects.
+  // cv::DenseOpticalFlow is not thread-safe. Invoking multiple
+  // DenseOpticalFlow::calc() in parallel may lead to memory corruption or
+  // memory leak.
+  std::list<cv::Ptr<cv::DenseOpticalFlow>> tvl1_computers_ GUARDED_BY(mutex_);
+  absl::Mutex mutex_;
+};
+
+::mediapipe::Status Tvl1OpticalFlowCalculator::GetContract(
+    CalculatorContract* cc) {
+  if (!cc->Inputs().HasTag("FIRST_FRAME") ||
+      !cc->Inputs().HasTag("SECOND_FRAME")) {
+    return ::mediapipe::InvalidArgumentError(
+        "Missing required input streams. Both FIRST_FRAME and SECOND_FRAME "
+        "must be specified.");
+  }
+  cc->Inputs().Tag("FIRST_FRAME").Set<ImageFrame>();
+  cc->Inputs().Tag("SECOND_FRAME").Set<ImageFrame>();
+  if (cc->Outputs().HasTag("FORWARD_FLOW")) {
+    cc->Outputs().Tag("FORWARD_FLOW").Set<OpticalFlowField>();
+  }
+  if (cc->Outputs().HasTag("BACKWARD_FLOW")) {
+    cc->Outputs().Tag("BACKWARD_FLOW").Set<OpticalFlowField>();
+  }
+  return ::mediapipe::OkStatus();
+}
+
+::mediapipe::Status Tvl1OpticalFlowCalculator::Open(CalculatorContext* cc) {
+  {
+    absl::MutexLock lock(&mutex_);
+    tvl1_computers_.emplace_back(cv::createOptFlow_DualTVL1());
+  }
+  if (cc->Outputs().HasTag("FORWARD_FLOW")) {
+    forward_requested_ = true;
+  }
+  if (cc->Outputs().HasTag("BACKWARD_FLOW")) {
+    backward_requested_ = true;
+  }
+
+  return ::mediapipe::OkStatus();
+}
+
+::mediapipe::Status Tvl1OpticalFlowCalculator::Process(CalculatorContext* cc) {
+  const ImageFrame& first_frame =
+      cc->Inputs().Tag("FIRST_FRAME").Value().Get<ImageFrame>();
+  const ImageFrame& second_frame =
+      cc->Inputs().Tag("SECOND_FRAME").Value().Get<ImageFrame>();
+  if (forward_requested_) {
+    auto forward_optical_flow_field = absl::make_unique<OpticalFlowField>();
+    RETURN_IF_ERROR(CalculateOpticalFlow(first_frame, second_frame,
+                                         forward_optical_flow_field.get()));
+    cc->Outputs()
+        .Tag("FORWARD_FLOW")
+        .Add(forward_optical_flow_field.release(), cc->InputTimestamp());
+  }
+  if (backward_requested_) {
+    auto backward_optical_flow_field = absl::make_unique<OpticalFlowField>();
+    RETURN_IF_ERROR(CalculateOpticalFlow(second_frame, first_frame,
+                                         backward_optical_flow_field.get()));
+    cc->Outputs()
+        .Tag("BACKWARD_FLOW")
+        .Add(backward_optical_flow_field.release(), cc->InputTimestamp());
+  }
+  return ::mediapipe::OkStatus();
+}
+
+::mediapipe::Status Tvl1OpticalFlowCalculator::CalculateOpticalFlow(
+    const ImageFrame& current_frame, const ImageFrame& next_frame,
+    OpticalFlowField* flow) {
+  CHECK(flow);
+  if (!ImageSizesMatch(current_frame, next_frame)) {
+    return tool::StatusInvalid("Images are different sizes.");
+  }
+  const cv::Mat& first = ConvertToGrayscale(formats::MatView(&current_frame));
+  const cv::Mat& second = ConvertToGrayscale(formats::MatView(&next_frame));
+
+  // Tries getting an idle DenseOpticalFlow object from the cache. If not,
+  // creates a new DenseOpticalFlow.
+  cv::Ptr<cv::DenseOpticalFlow> tvl1_computer;
+  {
+    absl::MutexLock lock(&mutex_);
+    if (!tvl1_computers_.empty()) {
+      std::swap(tvl1_computer, tvl1_computers_.front());
+      tvl1_computers_.pop_front();
+    }
+  }
+  if (tvl1_computer.empty()) {
+    tvl1_computer = cv::createOptFlow_DualTVL1();
+  }
+
+  flow->Allocate(first.cols, first.rows);
+  cv::Mat cv_flow(flow->mutable_flow_data());
+  tvl1_computer->calc(first, second, cv_flow);
+  CHECK_EQ(flow->mutable_flow_data().data, cv_flow.data);
+  // Inserts the idle DenseOpticalFlow object back to the cache for reuse.
+  {
+    absl::MutexLock lock(&mutex_);
+    tvl1_computers_.push_back(tvl1_computer);
+  }
+  return ::mediapipe::OkStatus();
+}
+
+REGISTER_CALCULATOR(Tvl1OpticalFlowCalculator);
+
+}  // namespace mediapipe

mediapipe/calculators/video/tvl1_optical_flow_calculator_test.cc

@@ -0,0 +1,127 @@
+// Copyright 2019 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/framework/calculator_framework.h"
+#include "mediapipe/framework/calculator_runner.h"
+#include "mediapipe/framework/deps/file_path.h"
+#include "mediapipe/framework/formats/image_frame_opencv.h"
+#include "mediapipe/framework/formats/motion/optical_flow_field.h"
+#include "mediapipe/framework/port/file_helpers.h"
+#include "mediapipe/framework/port/gmock.h"
+#include "mediapipe/framework/port/gtest.h"
+#include "mediapipe/framework/port/opencv_imgcodecs_inc.h"
+#include "mediapipe/framework/port/opencv_imgproc_inc.h"
+#include "mediapipe/framework/port/parse_text_proto.h"
+#include "mediapipe/framework/port/status_matchers.h"
+
+namespace mediapipe {
+
+namespace {
+void AddInputPackets(int num_packets, CalculatorGraph* graph) {
+  int width = 127;
+  int height = 227;
+  Packet packet1 = MakePacket<ImageFrame>(ImageFormat::SRGB, width, height);
+  Packet packet2 = MakePacket<ImageFrame>(ImageFormat::SRGB, width, height);
+  cv::Mat mat1 = formats::MatView(&(packet1.Get<ImageFrame>()));
+  cv::Mat mat2 = formats::MatView(&(packet2.Get<ImageFrame>()));
+  for (int r = 0; r < mat1.rows; ++r) {
+    for (int c = 0; c < mat1.cols; ++c) {
+      cv::Vec3b& color1 = mat1.at<cv::Vec3b>(r, c);
+      color1[0] = r + 3;
+      color1[1] = r + 3;
+      color1[2] = 0;
+      cv::Vec3b& color2 = mat2.at<cv::Vec3b>(r, c);
+      color2[0] = r;
+      color2[1] = r;
+      color2[2] = 0;
+    }
+  }
+
+  for (int i = 0; i < num_packets; ++i) {
+    MEDIAPIPE_ASSERT_OK(graph->AddPacketToInputStream(
+        "first_frames", packet1.At(Timestamp(i))));
+    MEDIAPIPE_ASSERT_OK(graph->AddPacketToInputStream(
+        "second_frames", packet2.At(Timestamp(i))));
+  }
+  MEDIAPIPE_ASSERT_OK(graph->CloseAllInputStreams());
+}
+
+void RunTest(int num_input_packets, int max_in_flight) {
+  CalculatorGraphConfig config = ParseTextProtoOrDie<CalculatorGraphConfig>(
+      absl::Substitute(R"(
+    input_stream: "first_frames"
+    input_stream: "second_frames"
+    node {
+      calculator: "Tvl1OpticalFlowCalculator"
+      input_stream: "FIRST_FRAME:first_frames"
+      input_stream: "SECOND_FRAME:second_frames"
+      output_stream: "FORWARD_FLOW:forward_flow"
+      output_stream: "BACKWARD_FLOW:backward_flow"
+      max_in_flight: $0
+    }
+    num_threads: $0
+  )",
+                       max_in_flight));
+  CalculatorGraph graph;
+  MEDIAPIPE_ASSERT_OK(graph.Initialize(config));
+  StatusOrPoller status_or_poller1 =
+      graph.AddOutputStreamPoller("forward_flow");
+  ASSERT_TRUE(status_or_poller1.ok());
+  OutputStreamPoller poller1 = std::move(status_or_poller1.ValueOrDie());
+  StatusOrPoller status_or_poller2 =
+      graph.AddOutputStreamPoller("backward_flow");
+  ASSERT_TRUE(status_or_poller2.ok());
+  OutputStreamPoller poller2 = std::move(status_or_poller2.ValueOrDie());
+
+  MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
+  AddInputPackets(num_input_packets, &graph);
+  Packet packet;
+  std::vector<Packet> forward_optical_flow_packets;
+  while (poller1.Next(&packet)) {
+    forward_optical_flow_packets.emplace_back(packet);
+  }
+  std::vector<Packet> backward_optical_flow_packets;
+  while (poller2.Next(&packet)) {
+    backward_optical_flow_packets.emplace_back(packet);
+  }
+  MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
+  EXPECT_EQ(num_input_packets, forward_optical_flow_packets.size());
+
+  int count = 0;
+  for (const Packet& packet : forward_optical_flow_packets) {
+    cv::Scalar average = cv::mean(packet.Get<OpticalFlowField>().flow_data());
+    EXPECT_NEAR(average[0], 0.0, 0.5) << "Actual mean_dx = " << average[0];
+    EXPECT_NEAR(average[1], 3.0, 0.5) << "Actual mean_dy = " << average[1];
+    EXPECT_EQ(count++, packet.Timestamp().Value());
+  }
+  EXPECT_EQ(num_input_packets, backward_optical_flow_packets.size());
+  count = 0;
+  for (const Packet& packet : backward_optical_flow_packets) {
+    cv::Scalar average = cv::mean(packet.Get<OpticalFlowField>().flow_data());
+    EXPECT_NEAR(average[0], 0.0, 0.5) << "Actual mean_dx = " << average[0];
+    EXPECT_NEAR(average[1], -3.0, 0.5) << "Actual mean_dy = " << average[1];
+    EXPECT_EQ(count++, packet.Timestamp().Value());
+  }
+}
+
+TEST(Tvl1OpticalFlowCalculatorTest, TestSequentialExecution) {
+  RunTest(/*num_input_packets=*/2, /*max_in_flight=*/1);
+}
+
+TEST(Tvl1OpticalFlowCalculatorTest, TestParallelExecution) {
+  RunTest(/*num_input_packets=*/20, /*max_in_flight=*/10);
+}
+
+}  // namespace
+}  // namespace mediapipe