1 /*
2 * Copyright (c) 2018 The WebRTC project authors. All Rights Reserved.
3 *
4 * Use of this source code is governed by a BSD-style license
5 * that can be found in the LICENSE file in the root of the source
6 * tree. An additional intellectual property rights grant can be found
7 * in the file PATENTS. All contributing project authors may
8 * be found in the AUTHORS file in the root of the source tree.
9 */
10
11 #include "rtc_tools/frame_analyzer/video_color_aligner.h"
12
13 #include <stddef.h>
14
15 #include <algorithm>
16 #include <cmath>
17 #include <cstdint>
18 #include <vector>
19
20 #include "api/array_view.h"
21 #include "api/make_ref_counted.h"
22 #include "api/video/i420_buffer.h"
23 #include "rtc_base/checks.h"
24 #include "rtc_tools/frame_analyzer/linear_least_squares.h"
25 #include "third_party/libyuv/include/libyuv/planar_functions.h"
26 #include "third_party/libyuv/include/libyuv/scale.h"
27
28 namespace webrtc {
29 namespace test {
30
31 namespace {
32
33 // Helper function for AdjustColors(). This functions calculates a single output
34 // row for y with the given color coefficients. The u/v channels are assumed to
35 // be subsampled by a factor of 2, which is the case of I420.
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)36 void CalculateYChannel(rtc::ArrayView<const uint8_t> y_data,
37 rtc::ArrayView<const uint8_t> u_data,
38 rtc::ArrayView<const uint8_t> v_data,
39 const std::array<float, 4>& coeff,
40 rtc::ArrayView<uint8_t> output) {
41 RTC_CHECK_EQ(y_data.size(), output.size());
42 // Each u/v element represents two y elements. Make sure we have enough to
43 // cover the Y values.
44 RTC_CHECK_GE(u_data.size() * 2, y_data.size());
45 RTC_CHECK_GE(v_data.size() * 2, y_data.size());
46
47 // Do two pixels at a time since u/v are subsampled.
48 for (size_t i = 0; i * 2 < y_data.size() - 1; ++i) {
49 const float uv_contribution =
50 coeff[1] * u_data[i] + coeff[2] * v_data[i] + coeff[3];
51
52 const float val0 = coeff[0] * y_data[i * 2 + 0] + uv_contribution;
53 const float val1 = coeff[0] * y_data[i * 2 + 1] + uv_contribution;
54
55 // Clamp result to a byte.
56 output[i * 2 + 0] = static_cast<uint8_t>(
57 std::round(std::max(0.0f, std::min(val0, 255.0f))));
58 output[i * 2 + 1] = static_cast<uint8_t>(
59 std::round(std::max(0.0f, std::min(val1, 255.0f))));
60 }
61
62 // Handle the last pixel for odd widths.
63 if (y_data.size() % 2 == 1) {
64 const float val = coeff[0] * y_data[y_data.size() - 1] +
65 coeff[1] * u_data[(y_data.size() - 1) / 2] +
66 coeff[2] * v_data[(y_data.size() - 1) / 2] + coeff[3];
67 output[y_data.size() - 1] =
68 static_cast<uint8_t>(std::round(std::max(0.0f, std::min(val, 255.0f))));
69 }
70 }
71
72 // Helper function for AdjustColors(). This functions calculates a single output
73 // row for either u or v, with the given color coefficients. Y, U, and V are
74 // assumed to be the same size, i.e. no subsampling.
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)75 void CalculateUVChannel(rtc::ArrayView<const uint8_t> y_data,
76 rtc::ArrayView<const uint8_t> u_data,
77 rtc::ArrayView<const uint8_t> v_data,
78 const std::array<float, 4>& coeff,
79 rtc::ArrayView<uint8_t> output) {
80 RTC_CHECK_EQ(y_data.size(), u_data.size());
81 RTC_CHECK_EQ(y_data.size(), v_data.size());
82 RTC_CHECK_EQ(y_data.size(), output.size());
83
84 for (size_t x = 0; x < y_data.size(); ++x) {
85 const float val = coeff[0] * y_data[x] + coeff[1] * u_data[x] +
86 coeff[2] * v_data[x] + coeff[3];
87 // Clamp result to a byte.
88 output[x] =
89 static_cast<uint8_t>(std::round(std::max(0.0f, std::min(val, 255.0f))));
90 }
91 }
92
93 // Convert a frame to four vectors consisting of [y, u, v, 1].
FlattenYuvData(const rtc::scoped_refptr<I420BufferInterface> & frame)94 std::vector<std::vector<uint8_t>> FlattenYuvData(
95 const rtc::scoped_refptr<I420BufferInterface>& frame) {
96 std::vector<std::vector<uint8_t>> result(
97 4, std::vector<uint8_t>(frame->ChromaWidth() * frame->ChromaHeight()));
98
99 // Downscale the Y plane so that all YUV planes are the same size.
100 libyuv::ScalePlane(frame->DataY(), frame->StrideY(), frame->width(),
101 frame->height(), result[0].data(), frame->ChromaWidth(),
102 frame->ChromaWidth(), frame->ChromaHeight(),
103 libyuv::kFilterBox);
104
105 libyuv::CopyPlane(frame->DataU(), frame->StrideU(), result[1].data(),
106 frame->ChromaWidth(), frame->ChromaWidth(),
107 frame->ChromaHeight());
108
109 libyuv::CopyPlane(frame->DataV(), frame->StrideV(), result[2].data(),
110 frame->ChromaWidth(), frame->ChromaWidth(),
111 frame->ChromaHeight());
112
113 std::fill(result[3].begin(), result[3].end(), 1u);
114
115 return result;
116 }
117
VectorToColorMatrix(const std::vector<std::vector<double>> & v)118 ColorTransformationMatrix VectorToColorMatrix(
119 const std::vector<std::vector<double>>& v) {
120 ColorTransformationMatrix color_transformation;
121 for (int i = 0; i < 3; ++i) {
122 for (int j = 0; j < 4; ++j)
123 color_transformation[i][j] = v[i][j];
124 }
125 return color_transformation;
126 }
127
128 } // namespace
129
CalculateColorTransformationMatrix(const rtc::scoped_refptr<I420BufferInterface> & reference_frame,const rtc::scoped_refptr<I420BufferInterface> & test_frame)130 ColorTransformationMatrix CalculateColorTransformationMatrix(
131 const rtc::scoped_refptr<I420BufferInterface>& reference_frame,
132 const rtc::scoped_refptr<I420BufferInterface>& test_frame) {
133 IncrementalLinearLeastSquares incremental_lls;
134 incremental_lls.AddObservations(FlattenYuvData(test_frame),
135 FlattenYuvData(reference_frame));
136 return VectorToColorMatrix(incremental_lls.GetBestSolution());
137 }
138
CalculateColorTransformationMatrix(const rtc::scoped_refptr<Video> & reference_video,const rtc::scoped_refptr<Video> & test_video)139 ColorTransformationMatrix CalculateColorTransformationMatrix(
140 const rtc::scoped_refptr<Video>& reference_video,
141 const rtc::scoped_refptr<Video>& test_video) {
142 RTC_CHECK_GE(reference_video->number_of_frames(),
143 test_video->number_of_frames());
144
145 IncrementalLinearLeastSquares incremental_lls;
146 for (size_t i = 0; i < test_video->number_of_frames(); ++i) {
147 incremental_lls.AddObservations(
148 FlattenYuvData(test_video->GetFrame(i)),
149 FlattenYuvData(reference_video->GetFrame(i)));
150 }
151
152 return VectorToColorMatrix(incremental_lls.GetBestSolution());
153 }
154
AdjustColors(const ColorTransformationMatrix & color_transformation,const rtc::scoped_refptr<Video> & video)155 rtc::scoped_refptr<Video> AdjustColors(
156 const ColorTransformationMatrix& color_transformation,
157 const rtc::scoped_refptr<Video>& video) {
158 class ColorAdjustedVideo : public Video {
159 public:
160 ColorAdjustedVideo(const ColorTransformationMatrix& color_transformation,
161 const rtc::scoped_refptr<Video>& video)
162 : color_transformation_(color_transformation), video_(video) {}
163
164 int width() const override { return video_->width(); }
165 int height() const override { return video_->height(); }
166 size_t number_of_frames() const override {
167 return video_->number_of_frames();
168 }
169
170 rtc::scoped_refptr<I420BufferInterface> GetFrame(
171 size_t index) const override {
172 return AdjustColors(color_transformation_, video_->GetFrame(index));
173 }
174
175 private:
176 const ColorTransformationMatrix color_transformation_;
177 const rtc::scoped_refptr<Video> video_;
178 };
179
180 return rtc::make_ref_counted<ColorAdjustedVideo>(color_transformation, video);
181 }
182
AdjustColors(const ColorTransformationMatrix & color_matrix,const rtc::scoped_refptr<I420BufferInterface> & frame)183 rtc::scoped_refptr<I420BufferInterface> AdjustColors(
184 const ColorTransformationMatrix& color_matrix,
185 const rtc::scoped_refptr<I420BufferInterface>& frame) {
186 // Allocate I420 buffer that will hold the color adjusted frame.
187 rtc::scoped_refptr<I420Buffer> adjusted_frame =
188 I420Buffer::Create(frame->width(), frame->height());
189
190 // Create a downscaled Y plane with the same size as the U/V planes to
191 // simplify converting the U/V planes.
192 std::vector<uint8_t> downscaled_y_plane(frame->ChromaWidth() *
193 frame->ChromaHeight());
194 libyuv::ScalePlane(frame->DataY(), frame->StrideY(), frame->width(),
195 frame->height(), downscaled_y_plane.data(),
196 frame->ChromaWidth(), frame->ChromaWidth(),
197 frame->ChromaHeight(), libyuv::kFilterBox);
198
199 // Fill in the adjusted data row by row.
200 for (int y = 0; y < frame->height(); ++y) {
201 const int half_y = y / 2;
202 rtc::ArrayView<const uint8_t> y_row(frame->DataY() + frame->StrideY() * y,
203 frame->width());
204 rtc::ArrayView<const uint8_t> u_row(
205 frame->DataU() + frame->StrideU() * half_y, frame->ChromaWidth());
206 rtc::ArrayView<const uint8_t> v_row(
207 frame->DataV() + frame->StrideV() * half_y, frame->ChromaWidth());
208 rtc::ArrayView<uint8_t> output_y_row(
209 adjusted_frame->MutableDataY() + adjusted_frame->StrideY() * y,
210 frame->width());
211
212 CalculateYChannel(y_row, u_row, v_row, color_matrix[0], output_y_row);
213
214 // Chroma channels only exist every second row for I420.
215 if (y % 2 == 0) {
216 rtc::ArrayView<const uint8_t> downscaled_y_row(
217 downscaled_y_plane.data() + frame->ChromaWidth() * half_y,
218 frame->ChromaWidth());
219 rtc::ArrayView<uint8_t> output_u_row(
220 adjusted_frame->MutableDataU() + adjusted_frame->StrideU() * half_y,
221 frame->ChromaWidth());
222 rtc::ArrayView<uint8_t> output_v_row(
223 adjusted_frame->MutableDataV() + adjusted_frame->StrideV() * half_y,
224 frame->ChromaWidth());
225
226 CalculateUVChannel(downscaled_y_row, u_row, v_row, color_matrix[1],
227 output_u_row);
228 CalculateUVChannel(downscaled_y_row, u_row, v_row, color_matrix[2],
229 output_v_row);
230 }
231 }
232
233 return adjusted_frame;
234 }
235
236 } // namespace test
237 } // namespace webrtc
238