1 /*
2 * Copyright (c) 2021, Alliance for Open Media. All rights reserved
3 *
4 * This source code is subject to the terms of the BSD 2 Clause License and
5 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
6 * was not distributed with this source code in the LICENSE file, you can
7 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
8 * Media Patent License 1.0 was not distributed with this source code in the
9 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
10 */
11
12 #include <math.h>
13
14 #include "av1/encoder/tune_butteraugli.h"
15
16 #include "aom_dsp/butteraugli.h"
17 #include "av1/encoder/encodeframe.h"
18 #include "av1/encoder/encoder_utils.h"
19 #include "av1/encoder/extend.h"
20 #include "av1/encoder/var_based_part.h"
21
22 static const int resize_factor = 2;
23
set_mb_butteraugli_rdmult_scaling(AV1_COMP * cpi,const YV12_BUFFER_CONFIG * source,const YV12_BUFFER_CONFIG * recon,const double K)24 static void set_mb_butteraugli_rdmult_scaling(AV1_COMP *cpi,
25 const YV12_BUFFER_CONFIG *source,
26 const YV12_BUFFER_CONFIG *recon,
27 const double K) {
28 AV1_COMMON *const cm = &cpi->common;
29 SequenceHeader *const seq_params = cm->seq_params;
30 const CommonModeInfoParams *const mi_params = &cm->mi_params;
31 const aom_color_range_t color_range =
32 seq_params->color_range != 0 ? AOM_CR_FULL_RANGE : AOM_CR_STUDIO_RANGE;
33 const int bit_depth = cpi->td.mb.e_mbd.bd;
34 const int width = source->y_crop_width;
35 const int height = source->y_crop_height;
36 const int ss_x = source->subsampling_x;
37 const int ss_y = source->subsampling_y;
38
39 float *diffmap;
40 CHECK_MEM_ERROR(cm, diffmap, aom_malloc(width * height * sizeof(*diffmap)));
41 if (!aom_calc_butteraugli(source, recon, bit_depth,
42 seq_params->matrix_coefficients, color_range,
43 diffmap)) {
44 aom_internal_error(cm->error, AOM_CODEC_ERROR,
45 "Failed to calculate Butteraugli distances.");
46 }
47
48 const int num_mi_w = mi_size_wide[butteraugli_rdo_bsize] / resize_factor;
49 const int num_mi_h = mi_size_high[butteraugli_rdo_bsize] / resize_factor;
50 const int num_cols =
51 (mi_params->mi_cols / resize_factor + num_mi_w - 1) / num_mi_w;
52 const int num_rows =
53 (mi_params->mi_rows / resize_factor + num_mi_h - 1) / num_mi_h;
54 const int block_w = num_mi_w << 2;
55 const int block_h = num_mi_h << 2;
56 double log_sum = 0.0;
57 double blk_count = 0.0;
58
59 // Loop through each block.
60 for (int row = 0; row < num_rows; ++row) {
61 for (int col = 0; col < num_cols; ++col) {
62 const int index = row * num_cols + col;
63 const int y_start = row * block_h;
64 const int x_start = col * block_w;
65 float dbutteraugli = 0.0f;
66 float dmse = 0.0f;
67 float px_count = 0.0f;
68
69 // Loop through each pixel.
70 for (int y = y_start; y < y_start + block_h && y < height; y++) {
71 for (int x = x_start; x < x_start + block_w && x < width; x++) {
72 dbutteraugli += powf(diffmap[y * width + x], 12.0f);
73 float px_diff = source->y_buffer[y * source->y_stride + x] -
74 recon->y_buffer[y * recon->y_stride + x];
75 dmse += px_diff * px_diff;
76 px_count += 1.0f;
77 }
78 }
79 const int y_end = AOMMIN((y_start >> ss_y) + (block_h >> ss_y),
80 (height + ss_y) >> ss_y);
81 for (int y = y_start >> ss_y; y < y_end; y++) {
82 const int x_end = AOMMIN((x_start >> ss_x) + (block_w >> ss_x),
83 (width + ss_x) >> ss_x);
84 for (int x = x_start >> ss_x; x < x_end; x++) {
85 const int src_px_index = y * source->uv_stride + x;
86 const int recon_px_index = y * recon->uv_stride + x;
87 const float px_diff_u = (float)(source->u_buffer[src_px_index] -
88 recon->u_buffer[recon_px_index]);
89 const float px_diff_v = (float)(source->v_buffer[src_px_index] -
90 recon->v_buffer[recon_px_index]);
91 dmse += px_diff_u * px_diff_u + px_diff_v * px_diff_v;
92 px_count += 2.0f;
93 }
94 }
95
96 dbutteraugli = powf(dbutteraugli, 1.0f / 12.0f);
97 dmse = dmse / px_count;
98 const float eps = 0.01f;
99 double weight;
100 if (dbutteraugli < eps || dmse < eps) {
101 weight = -1.0;
102 } else {
103 blk_count += 1.0;
104 weight = dmse / dbutteraugli;
105 weight = AOMMIN(weight, 5.0);
106 weight += K;
107 log_sum += log(weight);
108 }
109 cpi->butteraugli_info.rdmult_scaling_factors[index] = weight;
110 }
111 }
112 // Geometric average of the weights.
113 log_sum = exp(log_sum / blk_count);
114
115 for (int row = 0; row < num_rows; ++row) {
116 for (int col = 0; col < num_cols; ++col) {
117 const int index = row * num_cols + col;
118 double *weight = &cpi->butteraugli_info.rdmult_scaling_factors[index];
119 if (*weight <= 0.0) {
120 *weight = 1.0;
121 } else {
122 *weight /= log_sum;
123 }
124 *weight = AOMMIN(*weight, 2.5);
125 *weight = AOMMAX(*weight, 0.4);
126 }
127 }
128
129 aom_free(diffmap);
130 }
131
av1_set_butteraugli_rdmult(const AV1_COMP * cpi,MACROBLOCK * x,BLOCK_SIZE bsize,int mi_row,int mi_col,int * rdmult)132 void av1_set_butteraugli_rdmult(const AV1_COMP *cpi, MACROBLOCK *x,
133 BLOCK_SIZE bsize, int mi_row, int mi_col,
134 int *rdmult) {
135 assert(cpi->oxcf.tune_cfg.tuning == AOM_TUNE_BUTTERAUGLI);
136 if (!cpi->butteraugli_info.recon_set) {
137 return;
138 }
139 const AV1_COMMON *const cm = &cpi->common;
140
141 const int num_mi_w = mi_size_wide[butteraugli_rdo_bsize];
142 const int num_mi_h = mi_size_high[butteraugli_rdo_bsize];
143 const int num_cols = (cm->mi_params.mi_cols + num_mi_w - 1) / num_mi_w;
144 const int num_rows = (cm->mi_params.mi_rows + num_mi_h - 1) / num_mi_h;
145 const int num_bcols = (mi_size_wide[bsize] + num_mi_w - 1) / num_mi_w;
146 const int num_brows = (mi_size_high[bsize] + num_mi_h - 1) / num_mi_h;
147 double num_of_mi = 0.0;
148 double geom_mean_of_scale = 0.0;
149
150 for (int row = mi_row / num_mi_w;
151 row < num_rows && row < mi_row / num_mi_w + num_brows; ++row) {
152 for (int col = mi_col / num_mi_h;
153 col < num_cols && col < mi_col / num_mi_h + num_bcols; ++col) {
154 const int index = row * num_cols + col;
155 geom_mean_of_scale +=
156 log(cpi->butteraugli_info.rdmult_scaling_factors[index]);
157 num_of_mi += 1.0;
158 }
159 }
160 geom_mean_of_scale = exp(geom_mean_of_scale / num_of_mi);
161
162 *rdmult = (int)((double)(*rdmult) * geom_mean_of_scale + 0.5);
163 *rdmult = AOMMAX(*rdmult, 0);
164 av1_set_error_per_bit(&x->errorperbit, *rdmult);
165 }
166
copy_plane(const uint8_t * src,int src_stride,uint8_t * dst,int dst_stride,int w,int h)167 static void copy_plane(const uint8_t *src, int src_stride, uint8_t *dst,
168 int dst_stride, int w, int h) {
169 for (int row = 0; row < h; row++) {
170 memcpy(dst, src, w);
171 src += src_stride;
172 dst += dst_stride;
173 }
174 }
175
copy_img(const YV12_BUFFER_CONFIG * src,YV12_BUFFER_CONFIG * dst,int width,int height)176 static void copy_img(const YV12_BUFFER_CONFIG *src, YV12_BUFFER_CONFIG *dst,
177 int width, int height) {
178 copy_plane(src->y_buffer, src->y_stride, dst->y_buffer, dst->y_stride, width,
179 height);
180 const int width_uv = (width + src->subsampling_x) >> src->subsampling_x;
181 const int height_uv = (height + src->subsampling_y) >> src->subsampling_y;
182 copy_plane(src->u_buffer, src->uv_stride, dst->u_buffer, dst->uv_stride,
183 width_uv, height_uv);
184 copy_plane(src->v_buffer, src->uv_stride, dst->v_buffer, dst->uv_stride,
185 width_uv, height_uv);
186 }
187
zero_plane(uint8_t * dst,int dst_stride,int h)188 static void zero_plane(uint8_t *dst, int dst_stride, int h) {
189 for (int row = 0; row < h; row++) {
190 memset(dst, 0, dst_stride);
191 dst += dst_stride;
192 }
193 }
194
zero_img(YV12_BUFFER_CONFIG * dst)195 static void zero_img(YV12_BUFFER_CONFIG *dst) {
196 zero_plane(dst->y_buffer, dst->y_stride, dst->y_height);
197 zero_plane(dst->u_buffer, dst->uv_stride, dst->uv_height);
198 zero_plane(dst->v_buffer, dst->uv_stride, dst->uv_height);
199 }
200
av1_setup_butteraugli_source(AV1_COMP * cpi)201 void av1_setup_butteraugli_source(AV1_COMP *cpi) {
202 YV12_BUFFER_CONFIG *const dst = &cpi->butteraugli_info.source;
203 AV1_COMMON *const cm = &cpi->common;
204 const int width = cpi->source->y_crop_width;
205 const int height = cpi->source->y_crop_height;
206 const int bit_depth = cpi->td.mb.e_mbd.bd;
207 const int ss_x = cpi->source->subsampling_x;
208 const int ss_y = cpi->source->subsampling_y;
209 if (dst->buffer_alloc_sz == 0) {
210 aom_alloc_frame_buffer(
211 dst, width, height, ss_x, ss_y, cm->seq_params->use_highbitdepth,
212 cpi->oxcf.border_in_pixels, cm->features.byte_alignment);
213 }
214 av1_copy_and_extend_frame(cpi->source, dst);
215
216 YV12_BUFFER_CONFIG *const resized_dst = &cpi->butteraugli_info.resized_source;
217 if (resized_dst->buffer_alloc_sz == 0) {
218 aom_alloc_frame_buffer(
219 resized_dst, width / resize_factor, height / resize_factor, ss_x, ss_y,
220 cm->seq_params->use_highbitdepth, cpi->oxcf.border_in_pixels,
221 cm->features.byte_alignment);
222 }
223 av1_resize_and_extend_frame_nonnormative(cpi->source, resized_dst, bit_depth,
224 av1_num_planes(cm));
225
226 zero_img(cpi->source);
227 copy_img(resized_dst, cpi->source, width / resize_factor,
228 height / resize_factor);
229 }
230
av1_setup_butteraugli_rdmult_and_restore_source(AV1_COMP * cpi,double K)231 void av1_setup_butteraugli_rdmult_and_restore_source(AV1_COMP *cpi, double K) {
232 av1_copy_and_extend_frame(&cpi->butteraugli_info.source, cpi->source);
233 AV1_COMMON *const cm = &cpi->common;
234 const int width = cpi->source->y_crop_width;
235 const int height = cpi->source->y_crop_height;
236 const int ss_x = cpi->source->subsampling_x;
237 const int ss_y = cpi->source->subsampling_y;
238
239 YV12_BUFFER_CONFIG resized_recon;
240 memset(&resized_recon, 0, sizeof(resized_recon));
241 aom_alloc_frame_buffer(
242 &resized_recon, width / resize_factor, height / resize_factor, ss_x, ss_y,
243 cm->seq_params->use_highbitdepth, cpi->oxcf.border_in_pixels,
244 cm->features.byte_alignment);
245 copy_img(&cpi->common.cur_frame->buf, &resized_recon, width / resize_factor,
246 height / resize_factor);
247
248 set_mb_butteraugli_rdmult_scaling(cpi, &cpi->butteraugli_info.resized_source,
249 &resized_recon, K);
250 cpi->butteraugli_info.recon_set = true;
251 aom_free_frame_buffer(&resized_recon);
252 }
253
av1_setup_butteraugli_rdmult(AV1_COMP * cpi)254 void av1_setup_butteraugli_rdmult(AV1_COMP *cpi) {
255 AV1_COMMON *const cm = &cpi->common;
256 const AV1EncoderConfig *const oxcf = &cpi->oxcf;
257 const QuantizationCfg *const q_cfg = &oxcf->q_cfg;
258 const int q_index = 96;
259
260 // Setup necessary params for encoding, including frame source, etc.
261 if (cm->current_frame.frame_type == KEY_FRAME) copy_frame_prob_info(cpi);
262 av1_set_frame_size(cpi, cm->superres_upscaled_width,
263 cm->superres_upscaled_height);
264
265 cpi->source =
266 av1_scale_if_required(cm, cpi->unscaled_source, &cpi->scaled_source,
267 cm->features.interp_filter, 0, false, false);
268 if (cpi->unscaled_last_source != NULL) {
269 cpi->last_source = av1_scale_if_required(
270 cm, cpi->unscaled_last_source, &cpi->scaled_last_source,
271 cm->features.interp_filter, 0, false, false);
272 }
273
274 av1_setup_butteraugli_source(cpi);
275 av1_setup_frame(cpi);
276
277 if (cm->seg.enabled) {
278 if (!cm->seg.update_data && cm->prev_frame) {
279 segfeatures_copy(&cm->seg, &cm->prev_frame->seg);
280 cm->seg.enabled = cm->prev_frame->seg.enabled;
281 } else {
282 av1_calculate_segdata(&cm->seg);
283 }
284 } else {
285 memset(&cm->seg, 0, sizeof(cm->seg));
286 }
287 segfeatures_copy(&cm->cur_frame->seg, &cm->seg);
288 cm->cur_frame->seg.enabled = cm->seg.enabled;
289
290 const PARTITION_SEARCH_TYPE partition_search_type =
291 cpi->sf.part_sf.partition_search_type;
292 const BLOCK_SIZE fixed_partition_size = cpi->sf.part_sf.fixed_partition_size;
293 // Enable a quicker pass by uncommenting the following lines:
294 // cpi->sf.part_sf.partition_search_type = FIXED_PARTITION;
295 // cpi->sf.part_sf.fixed_partition_size = BLOCK_32X32;
296
297 av1_set_quantizer(cm, q_cfg->qm_minlevel, q_cfg->qm_maxlevel, q_index,
298 q_cfg->enable_chroma_deltaq);
299 av1_set_speed_features_qindex_dependent(cpi, oxcf->speed);
300 if (q_cfg->deltaq_mode != NO_DELTA_Q || q_cfg->enable_chroma_deltaq)
301 av1_init_quantizer(&cpi->enc_quant_dequant_params, &cm->quant_params,
302 cm->seq_params->bit_depth);
303
304 av1_set_variance_partition_thresholds(cpi, q_index, 0);
305 av1_encode_frame(cpi);
306
307 av1_setup_butteraugli_rdmult_and_restore_source(cpi, 0.3);
308 cpi->sf.part_sf.partition_search_type = partition_search_type;
309 cpi->sf.part_sf.fixed_partition_size = fixed_partition_size;
310 }
311