tune_butteraugli.c (12687B)
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 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 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 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 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 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 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 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, false, 0); 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, false, 0); 222 } 223 if (!av1_resize_and_extend_frame_nonnormative( 224 cpi->source, resized_dst, bit_depth, av1_num_planes(cm))) { 225 aom_internal_error(cm->error, AOM_CODEC_MEM_ERROR, 226 "Error allocating buffers during resize"); 227 } 228 229 zero_img(cpi->source); 230 copy_img(resized_dst, cpi->source, width / resize_factor, 231 height / resize_factor); 232 } 233 234 void av1_setup_butteraugli_rdmult_and_restore_source(AV1_COMP *cpi, double K) { 235 av1_copy_and_extend_frame(&cpi->butteraugli_info.source, cpi->source); 236 AV1_COMMON *const cm = &cpi->common; 237 const int width = cpi->source->y_crop_width; 238 const int height = cpi->source->y_crop_height; 239 const int ss_x = cpi->source->subsampling_x; 240 const int ss_y = cpi->source->subsampling_y; 241 242 YV12_BUFFER_CONFIG resized_recon; 243 memset(&resized_recon, 0, sizeof(resized_recon)); 244 aom_alloc_frame_buffer( 245 &resized_recon, width / resize_factor, height / resize_factor, ss_x, ss_y, 246 cm->seq_params->use_highbitdepth, cpi->oxcf.border_in_pixels, 247 cm->features.byte_alignment, false, 0); 248 copy_img(&cpi->common.cur_frame->buf, &resized_recon, width / resize_factor, 249 height / resize_factor); 250 251 set_mb_butteraugli_rdmult_scaling(cpi, &cpi->butteraugli_info.resized_source, 252 &resized_recon, K); 253 cpi->butteraugli_info.recon_set = true; 254 aom_free_frame_buffer(&resized_recon); 255 } 256 257 void av1_setup_butteraugli_rdmult(AV1_COMP *cpi) { 258 AV1_COMMON *const cm = &cpi->common; 259 const AV1EncoderConfig *const oxcf = &cpi->oxcf; 260 const QuantizationCfg *const q_cfg = &oxcf->q_cfg; 261 const int q_index = 96; 262 263 // Setup necessary params for encoding, including frame source, etc. 264 if (cm->current_frame.frame_type == KEY_FRAME) copy_frame_prob_info(cpi); 265 av1_set_frame_size(cpi, cm->superres_upscaled_width, 266 cm->superres_upscaled_height); 267 268 cpi->source = av1_realloc_and_scale_if_required( 269 cm, cpi->unscaled_source, &cpi->scaled_source, cm->features.interp_filter, 270 0, false, false, cpi->oxcf.border_in_pixels, cpi->alloc_pyramid); 271 if (cpi->unscaled_last_source != NULL) { 272 cpi->last_source = av1_realloc_and_scale_if_required( 273 cm, cpi->unscaled_last_source, &cpi->scaled_last_source, 274 cm->features.interp_filter, 0, false, false, cpi->oxcf.border_in_pixels, 275 cpi->alloc_pyramid); 276 } 277 278 av1_setup_butteraugli_source(cpi); 279 av1_setup_frame(cpi); 280 281 if (cm->seg.enabled) { 282 if (!cm->seg.update_data && cm->prev_frame) { 283 segfeatures_copy(&cm->seg, &cm->prev_frame->seg); 284 cm->seg.enabled = cm->prev_frame->seg.enabled; 285 } else { 286 av1_calculate_segdata(&cm->seg); 287 } 288 } else { 289 memset(&cm->seg, 0, sizeof(cm->seg)); 290 } 291 segfeatures_copy(&cm->cur_frame->seg, &cm->seg); 292 cm->cur_frame->seg.enabled = cm->seg.enabled; 293 294 const PARTITION_SEARCH_TYPE partition_search_type = 295 cpi->sf.part_sf.partition_search_type; 296 const BLOCK_SIZE fixed_partition_size = cpi->sf.part_sf.fixed_partition_size; 297 // Enable a quicker pass by uncommenting the following lines: 298 // cpi->sf.part_sf.partition_search_type = FIXED_PARTITION; 299 // cpi->sf.part_sf.fixed_partition_size = BLOCK_32X32; 300 301 av1_set_quantizer(cm, q_cfg->qm_minlevel, q_cfg->qm_maxlevel, q_index, 302 q_cfg->enable_chroma_deltaq, q_cfg->enable_hdr_deltaq, 303 oxcf->mode == ALLINTRA, oxcf->tune_cfg.tuning); 304 av1_set_speed_features_qindex_dependent(cpi, oxcf->speed); 305 av1_init_quantizer(&cpi->enc_quant_dequant_params, &cm->quant_params, 306 cm->seq_params->bit_depth, cpi->oxcf.algo_cfg.sharpness); 307 308 av1_set_variance_partition_thresholds(cpi, q_index, 0); 309 av1_encode_frame(cpi); 310 311 av1_setup_butteraugli_rdmult_and_restore_source(cpi, 0.3); 312 cpi->sf.part_sf.partition_search_type = partition_search_type; 313 cpi->sf.part_sf.fixed_partition_size = fixed_partition_size; 314 }