1 files changed, 237 insertions, 0 deletions
diff --git a/third_party/libwebrtc/rtc_tools/frame_analyzer/video_color_aligner.cc b/third_party/libwebrtc/rtc_tools/frame_analyzer/video_color_aligner.cc
new file mode 100644
index 0000000000..5983e47f69
--- /dev/null
+++ b/third_party/libwebrtc/rtc_tools/frame_analyzer/video_color_aligner.cc
@@ -0,0 +1,237 @@
+/*
+ *  Copyright (c) 2018 The WebRTC project authors. All Rights Reserved.
+ *
+ *  Use of this source code is governed by a BSD-style license
+ *  that can be found in the LICENSE file in the root of the source
+ *  tree. An additional intellectual property rights grant can be found
+ *  in the file PATENTS.  All contributing project authors may
+ *  be found in the AUTHORS file in the root of the source tree.
+ */
+
+#include "rtc_tools/frame_analyzer/video_color_aligner.h"
+
+#include <stddef.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <vector>
+
+#include "api/array_view.h"
+#include "api/make_ref_counted.h"
+#include "api/video/i420_buffer.h"
+#include "rtc_base/checks.h"
+#include "rtc_tools/frame_analyzer/linear_least_squares.h"
+#include "third_party/libyuv/include/libyuv/planar_functions.h"
+#include "third_party/libyuv/include/libyuv/scale.h"
+
+namespace webrtc {
+namespace test {
+
+namespace {
+
+// Helper function for AdjustColors(). This functions calculates a single output
+// row for y with the given color coefficients. The u/v channels are assumed to
+// be subsampled by a factor of 2, which is the case of I420.
+void CalculateYChannel(rtc::ArrayView<const uint8_t> y_data,
+                       rtc::ArrayView<const uint8_t> u_data,
+                       rtc::ArrayView<const uint8_t> v_data,
+                       const std::array<float, 4>& coeff,
+                       rtc::ArrayView<uint8_t> output) {
+  RTC_CHECK_EQ(y_data.size(), output.size());
+  // Each u/v element represents two y elements. Make sure we have enough to
+  // cover the Y values.
+  RTC_CHECK_GE(u_data.size() * 2, y_data.size());
+  RTC_CHECK_GE(v_data.size() * 2, y_data.size());
+
+  // Do two pixels at a time since u/v are subsampled.
+  for (size_t i = 0; i * 2 < y_data.size() - 1; ++i) {
+    const float uv_contribution =
+        coeff[1] * u_data[i] + coeff[2] * v_data[i] + coeff[3];
+
+    const float val0 = coeff[0] * y_data[i * 2 + 0] + uv_contribution;
+    const float val1 = coeff[0] * y_data[i * 2 + 1] + uv_contribution;
+
+    // Clamp result to a byte.
+    output[i * 2 + 0] = static_cast<uint8_t>(
+        std::round(std::max(0.0f, std::min(val0, 255.0f))));
+    output[i * 2 + 1] = static_cast<uint8_t>(
+        std::round(std::max(0.0f, std::min(val1, 255.0f))));
+  }
+
+  // Handle the last pixel for odd widths.
+  if (y_data.size() % 2 == 1) {
+    const float val = coeff[0] * y_data[y_data.size() - 1] +
+                      coeff[1] * u_data[(y_data.size() - 1) / 2] +
+                      coeff[2] * v_data[(y_data.size() - 1) / 2] + coeff[3];
+    output[y_data.size() - 1] =
+        static_cast<uint8_t>(std::round(std::max(0.0f, std::min(val, 255.0f))));
+  }
+}
+
+// Helper function for AdjustColors(). This functions calculates a single output
+// row for either u or v, with the given color coefficients. Y, U, and V are
+// assumed to be the same size, i.e. no subsampling.
+void CalculateUVChannel(rtc::ArrayView<const uint8_t> y_data,
+                        rtc::ArrayView<const uint8_t> u_data,
+                        rtc::ArrayView<const uint8_t> v_data,
+                        const std::array<float, 4>& coeff,
+                        rtc::ArrayView<uint8_t> output) {
+  RTC_CHECK_EQ(y_data.size(), u_data.size());
+  RTC_CHECK_EQ(y_data.size(), v_data.size());
+  RTC_CHECK_EQ(y_data.size(), output.size());
+
+  for (size_t x = 0; x < y_data.size(); ++x) {
+    const float val = coeff[0] * y_data[x] + coeff[1] * u_data[x] +
+                      coeff[2] * v_data[x] + coeff[3];
+    // Clamp result to a byte.
+    output[x] =
+        static_cast<uint8_t>(std::round(std::max(0.0f, std::min(val, 255.0f))));
+  }
+}
+
+// Convert a frame to four vectors consisting of [y, u, v, 1].
+std::vector<std::vector<uint8_t>> FlattenYuvData(
+    const rtc::scoped_refptr<I420BufferInterface>& frame) {
+  std::vector<std::vector<uint8_t>> result(
+      4, std::vector<uint8_t>(frame->ChromaWidth() * frame->ChromaHeight()));
+
+  // Downscale the Y plane so that all YUV planes are the same size.
+  libyuv::ScalePlane(frame->DataY(), frame->StrideY(), frame->width(),
+                     frame->height(), result[0].data(), frame->ChromaWidth(),
+                     frame->ChromaWidth(), frame->ChromaHeight(),
+                     libyuv::kFilterBox);
+
+  libyuv::CopyPlane(frame->DataU(), frame->StrideU(), result[1].data(),
+                    frame->ChromaWidth(), frame->ChromaWidth(),
+                    frame->ChromaHeight());
+
+  libyuv::CopyPlane(frame->DataV(), frame->StrideV(), result[2].data(),
+                    frame->ChromaWidth(), frame->ChromaWidth(),
+                    frame->ChromaHeight());
+
+  std::fill(result[3].begin(), result[3].end(), 1u);
+
+  return result;
+}
+
+ColorTransformationMatrix VectorToColorMatrix(
+    const std::vector<std::vector<double>>& v) {
+  ColorTransformationMatrix color_transformation;
+  for (int i = 0; i < 3; ++i) {
+    for (int j = 0; j < 4; ++j)
+      color_transformation[i][j] = v[i][j];
+  }
+  return color_transformation;
+}
+
+}  // namespace
+
+ColorTransformationMatrix CalculateColorTransformationMatrix(
+    const rtc::scoped_refptr<I420BufferInterface>& reference_frame,
+    const rtc::scoped_refptr<I420BufferInterface>& test_frame) {
+  IncrementalLinearLeastSquares incremental_lls;
+  incremental_lls.AddObservations(FlattenYuvData(test_frame),
+                                  FlattenYuvData(reference_frame));
+  return VectorToColorMatrix(incremental_lls.GetBestSolution());
+}
+
+ColorTransformationMatrix CalculateColorTransformationMatrix(
+    const rtc::scoped_refptr<Video>& reference_video,
+    const rtc::scoped_refptr<Video>& test_video) {
+  RTC_CHECK_GE(reference_video->number_of_frames(),
+               test_video->number_of_frames());
+
+  IncrementalLinearLeastSquares incremental_lls;
+  for (size_t i = 0; i < test_video->number_of_frames(); ++i) {
+    incremental_lls.AddObservations(
+        FlattenYuvData(test_video->GetFrame(i)),
+        FlattenYuvData(reference_video->GetFrame(i)));
+  }
+
+  return VectorToColorMatrix(incremental_lls.GetBestSolution());
+}
+
+rtc::scoped_refptr<Video> AdjustColors(
+    const ColorTransformationMatrix& color_transformation,
+    const rtc::scoped_refptr<Video>& video) {
+  class ColorAdjustedVideo : public Video {
+   public:
+    ColorAdjustedVideo(const ColorTransformationMatrix& color_transformation,
+                       const rtc::scoped_refptr<Video>& video)
+        : color_transformation_(color_transformation), video_(video) {}
+
+    int width() const override { return video_->width(); }
+    int height() const override { return video_->height(); }
+    size_t number_of_frames() const override {
+      return video_->number_of_frames();
+    }
+
+    rtc::scoped_refptr<I420BufferInterface> GetFrame(
+        size_t index) const override {
+      return AdjustColors(color_transformation_, video_->GetFrame(index));
+    }
+
+   private:
+    const ColorTransformationMatrix color_transformation_;
+    const rtc::scoped_refptr<Video> video_;
+  };
+
+  return rtc::make_ref_counted<ColorAdjustedVideo>(color_transformation, video);
+}
+
+rtc::scoped_refptr<I420BufferInterface> AdjustColors(
+    const ColorTransformationMatrix& color_matrix,
+    const rtc::scoped_refptr<I420BufferInterface>& frame) {
+  // Allocate I420 buffer that will hold the color adjusted frame.
+  rtc::scoped_refptr<I420Buffer> adjusted_frame =
+      I420Buffer::Create(frame->width(), frame->height());
+
+  // Create a downscaled Y plane with the same size as the U/V planes to
+  // simplify converting the U/V planes.
+  std::vector<uint8_t> downscaled_y_plane(frame->ChromaWidth() *
+                                          frame->ChromaHeight());
+  libyuv::ScalePlane(frame->DataY(), frame->StrideY(), frame->width(),
+                     frame->height(), downscaled_y_plane.data(),
+                     frame->ChromaWidth(), frame->ChromaWidth(),
+                     frame->ChromaHeight(), libyuv::kFilterBox);
+
+  // Fill in the adjusted data row by row.
+  for (int y = 0; y < frame->height(); ++y) {
+    const int half_y = y / 2;
+    rtc::ArrayView<const uint8_t> y_row(frame->DataY() + frame->StrideY() * y,
+                                        frame->width());
+    rtc::ArrayView<const uint8_t> u_row(
+        frame->DataU() + frame->StrideU() * half_y, frame->ChromaWidth());
+    rtc::ArrayView<const uint8_t> v_row(
+        frame->DataV() + frame->StrideV() * half_y, frame->ChromaWidth());
+    rtc::ArrayView<uint8_t> output_y_row(
+        adjusted_frame->MutableDataY() + adjusted_frame->StrideY() * y,
+        frame->width());
+
+    CalculateYChannel(y_row, u_row, v_row, color_matrix[0], output_y_row);
+
+    // Chroma channels only exist every second row for I420.
+    if (y % 2 == 0) {
+      rtc::ArrayView<const uint8_t> downscaled_y_row(
+          downscaled_y_plane.data() + frame->ChromaWidth() * half_y,
+          frame->ChromaWidth());
+      rtc::ArrayView<uint8_t> output_u_row(
+          adjusted_frame->MutableDataU() + adjusted_frame->StrideU() * half_y,
+          frame->ChromaWidth());
+      rtc::ArrayView<uint8_t> output_v_row(
+          adjusted_frame->MutableDataV() + adjusted_frame->StrideV() * half_y,
+          frame->ChromaWidth());
+
+      CalculateUVChannel(downscaled_y_row, u_row, v_row, color_matrix[1],
+                         output_u_row);
+      CalculateUVChannel(downscaled_y_row, u_row, v_row, color_matrix[2],
+                         output_v_row);
+    }
+  }
+
+  return adjusted_frame;
+}
+
+}  // namespace test
+}  // namespace webrtc