noise_model.c (17196B)
1 /* 2 * Copyright (c) 2018, 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 /*!\file 13 * \brief This is an sample binary to create noise params from input video. 14 * 15 * To allow for external denoising applications, this sample binary illustrates 16 * how to create a film grain table (film grain params as a function of time) 17 * from an input video and its corresponding denoised source. 18 * 19 * The --output-grain-table file can be passed as input to the encoder (in 20 * aomenc this is done through the "--film-grain-table" parameter). 21 * 22 * As an example, where the input source is an 854x480 yuv420p 8-bit video 23 * named "input.854_480.yuv" you would use steps similar to the following: 24 * 25 * # Run your denoiser (e.g, using hqdn3d filter): 26 * ffmpeg -vcodec rawvideo -video_size 854x480 -i input.854_480.yuv \ 27 * -vf hqdn3d=5:5:5:5 -vcodec rawvideo -an -f rawvideo \ 28 * denoised.854_480.yuv 29 * 30 * # Model the noise between the denoised version and original source: 31 * ./examples/noise_model --fps=25/1 --width=854 --height=480 --i420 \ 32 * --input-denoised=denoised.854_480.yuv --input=original.854_480.yuv \ 33 * --output-grain-table=film_grain.tbl 34 * 35 * # Encode with your favorite settings (including the grain table): 36 * aomenc --limit=100 --cpu-used=4 --input-bit-depth=8 \ 37 * --i420 -w 854 -h 480 --end-usage=q --cq-level=25 --lag-in-frames=25 \ 38 * --auto-alt-ref=2 --bit-depth=8 --film-grain-table=film_grain.tbl \ 39 * -o denoised_with_grain_params.ivf denoised.854_480.yuv 40 */ 41 #include <math.h> 42 #include <stdio.h> 43 #include <stdlib.h> 44 #include <string.h> 45 46 #include "aom/aom_encoder.h" 47 #include "aom_dsp/aom_dsp_common.h" 48 49 #if CONFIG_AV1_DECODER 50 #include "av1/decoder/grain_synthesis.h" 51 #endif 52 53 #include "aom_dsp/grain_table.h" 54 #include "aom_dsp/noise_model.h" 55 #include "aom_dsp/noise_util.h" 56 #include "aom_mem/aom_mem.h" 57 #include "common/args.h" 58 #include "common/tools_common.h" 59 #include "common/video_writer.h" 60 61 static const char *exec_name; 62 63 void usage_exit(void) { 64 fprintf(stderr, 65 "Usage: %s --input=<input> --input-denoised=<denoised> " 66 "--output-grain-table=<outfile> " 67 "See comments in noise_model.c for more information.\n", 68 exec_name); 69 exit(EXIT_FAILURE); 70 } 71 72 static const arg_def_t help = 73 ARG_DEF(NULL, "help", 0, "Show usage options and exit"); 74 static const arg_def_t width_arg = 75 ARG_DEF("w", "width", 1, "Input width (if rawvideo)"); 76 static const arg_def_t height_arg = 77 ARG_DEF("h", "height", 1, "Input height (if rawvideo)"); 78 static const arg_def_t skip_frames_arg = 79 ARG_DEF("s", "skip-frames", 1, "Number of frames to skip (default = 1)"); 80 static const arg_def_t fps_arg = ARG_DEF(NULL, "fps", 1, "Frame rate"); 81 static const arg_def_t input_arg = ARG_DEF("-i", "input", 1, "Input filename"); 82 static const arg_def_t output_grain_table_arg = 83 ARG_DEF("n", "output-grain-table", 1, "Output noise file"); 84 static const arg_def_t input_denoised_arg = 85 ARG_DEF("d", "input-denoised", 1, "Input denoised filename (YUV) only"); 86 static const arg_def_t flat_block_finder_arg = 87 ARG_DEF("b", "flat-block-finder", 1, "Run the flat block finder"); 88 static const arg_def_t block_size_arg = 89 ARG_DEF("b", "block-size", 1, "Block size"); 90 static const arg_def_t bit_depth_arg = 91 ARG_DEF(NULL, "bit-depth", 1, "Bit depth of input"); 92 static const arg_def_t use_i420 = 93 ARG_DEF(NULL, "i420", 0, "Input file (and denoised) is I420 (default)"); 94 static const arg_def_t use_i422 = 95 ARG_DEF(NULL, "i422", 0, "Input file (and denoised) is I422"); 96 static const arg_def_t use_i444 = 97 ARG_DEF(NULL, "i444", 0, "Input file (and denoised) is I444"); 98 static const arg_def_t debug_file_arg = 99 ARG_DEF(NULL, "debug-file", 1, "File to output debug info"); 100 101 typedef struct { 102 int width; 103 int height; 104 struct aom_rational fps; 105 const char *input; 106 const char *input_denoised; 107 const char *output_grain_table; 108 int img_fmt; 109 int block_size; 110 int bit_depth; 111 int run_flat_block_finder; 112 int force_flat_psd; 113 int skip_frames; 114 const char *debug_file; 115 } noise_model_args_t; 116 117 static void parse_args(noise_model_args_t *noise_args, char **argv) { 118 struct arg arg; 119 static const arg_def_t *main_args[] = { &help, 120 &input_arg, 121 &fps_arg, 122 &width_arg, 123 &height_arg, 124 &block_size_arg, 125 &output_grain_table_arg, 126 &input_denoised_arg, 127 &use_i420, 128 &use_i422, 129 &use_i444, 130 &debug_file_arg, 131 NULL }; 132 for (; *argv; argv++) { 133 if (arg_match(&arg, &help, argv)) { 134 fprintf(stdout, "\nOptions:\n"); 135 arg_show_usage(stdout, main_args); 136 exit(0); 137 } else if (arg_match(&arg, &width_arg, argv)) { 138 noise_args->width = atoi(arg.val); 139 } else if (arg_match(&arg, &height_arg, argv)) { 140 noise_args->height = atoi(arg.val); 141 } else if (arg_match(&arg, &input_arg, argv)) { 142 noise_args->input = arg.val; 143 } else if (arg_match(&arg, &input_denoised_arg, argv)) { 144 noise_args->input_denoised = arg.val; 145 } else if (arg_match(&arg, &output_grain_table_arg, argv)) { 146 noise_args->output_grain_table = arg.val; 147 } else if (arg_match(&arg, &block_size_arg, argv)) { 148 noise_args->block_size = atoi(arg.val); 149 } else if (arg_match(&arg, &bit_depth_arg, argv)) { 150 noise_args->bit_depth = atoi(arg.val); 151 } else if (arg_match(&arg, &flat_block_finder_arg, argv)) { 152 noise_args->run_flat_block_finder = atoi(arg.val); 153 } else if (arg_match(&arg, &fps_arg, argv)) { 154 noise_args->fps = arg_parse_rational(&arg); 155 } else if (arg_match(&arg, &use_i420, argv)) { 156 noise_args->img_fmt = AOM_IMG_FMT_I420; 157 } else if (arg_match(&arg, &use_i422, argv)) { 158 noise_args->img_fmt = AOM_IMG_FMT_I422; 159 } else if (arg_match(&arg, &use_i444, argv)) { 160 noise_args->img_fmt = AOM_IMG_FMT_I444; 161 } else if (arg_match(&arg, &skip_frames_arg, argv)) { 162 noise_args->skip_frames = atoi(arg.val); 163 } else if (arg_match(&arg, &debug_file_arg, argv)) { 164 noise_args->debug_file = arg.val; 165 } else { 166 fprintf(stdout, "Unknown arg: %s\n\nUsage:\n", *argv); 167 arg_show_usage(stdout, main_args); 168 exit(0); 169 } 170 } 171 if (noise_args->bit_depth > 8) { 172 noise_args->img_fmt |= AOM_IMG_FMT_HIGHBITDEPTH; 173 } 174 } 175 176 #if CONFIG_AV1_DECODER 177 static void print_variance_y(FILE *debug_file, aom_image_t *raw, 178 aom_image_t *denoised, const uint8_t *flat_blocks, 179 int block_size, aom_film_grain_t *grain) { 180 aom_image_t renoised; 181 grain->apply_grain = 1; 182 grain->random_seed = 7391; 183 grain->bit_depth = raw->bit_depth; 184 aom_img_alloc(&renoised, raw->fmt, raw->w, raw->h, 1); 185 186 if (av1_add_film_grain(grain, denoised, &renoised)) { 187 fprintf(stderr, "Internal failure in av1_add_film_grain().\n"); 188 aom_img_free(&renoised); 189 return; 190 } 191 192 const int num_blocks_w = (raw->w + block_size - 1) / block_size; 193 const int num_blocks_h = (raw->h + block_size - 1) / block_size; 194 fprintf(debug_file, "x = ["); 195 for (int by = 0; by < num_blocks_h; by++) { 196 for (int bx = 0; bx < num_blocks_w; bx++) { 197 double block_mean = 0; 198 double noise_std = 0, noise_mean = 0; 199 double renoise_std = 0, renoise_mean = 0; 200 for (int yi = 0; yi < block_size; ++yi) { 201 const int y = by * block_size + yi; 202 for (int xi = 0; xi < block_size; ++xi) { 203 const int x = bx * block_size + xi; 204 const double noise_v = (raw->planes[0][y * raw->stride[0] + x] - 205 denoised->planes[0][y * raw->stride[0] + x]); 206 noise_mean += noise_v; 207 noise_std += noise_v * noise_v; 208 209 block_mean += raw->planes[0][y * raw->stride[0] + x]; 210 211 const double renoise_v = 212 (renoised.planes[0][y * raw->stride[0] + x] - 213 denoised->planes[0][y * raw->stride[0] + x]); 214 renoise_mean += renoise_v; 215 renoise_std += renoise_v * renoise_v; 216 } 217 } 218 int n = (block_size * block_size); 219 block_mean /= n; 220 noise_mean /= n; 221 renoise_mean /= n; 222 noise_std = sqrt(noise_std / n - noise_mean * noise_mean); 223 renoise_std = sqrt(renoise_std / n - renoise_mean * renoise_mean); 224 fprintf(debug_file, "%d %3.2lf %3.2lf %3.2lf ", 225 flat_blocks[by * num_blocks_w + bx], block_mean, noise_std, 226 renoise_std); 227 } 228 fprintf(debug_file, "\n"); 229 } 230 fprintf(debug_file, "];\n"); 231 232 if (raw->fmt & AOM_IMG_FMT_HIGHBITDEPTH) { 233 fprintf(stderr, 234 "Detailed debug info not supported for high bit" 235 "depth formats\n"); 236 } else { 237 fprintf(debug_file, "figure(2); clf;\n"); 238 fprintf(debug_file, 239 "scatter(x(:, 2:4:end), x(:, 3:4:end), 'r'); hold on;\n"); 240 fprintf(debug_file, "scatter(x(:, 2:4:end), x(:, 4:4:end), 'b');\n"); 241 fprintf(debug_file, 242 "plot(linspace(0, 255, length(noise_strength_0)), " 243 "noise_strength_0, 'b');\n"); 244 fprintf(debug_file, 245 "title('Scatter plot of intensity vs noise strength');\n"); 246 fprintf(debug_file, 247 "legend('Actual', 'Estimated', 'Estimated strength');\n"); 248 fprintf(debug_file, "figure(3); clf;\n"); 249 fprintf(debug_file, "scatter(x(:, 3:4:end), x(:, 4:4:end), 'k');\n"); 250 fprintf(debug_file, "title('Actual vs Estimated');\n"); 251 fprintf(debug_file, "pause(3);\n"); 252 } 253 aom_img_free(&renoised); 254 } 255 #endif 256 257 static void print_debug_info(FILE *debug_file, aom_image_t *raw, 258 aom_image_t *denoised, uint8_t *flat_blocks, 259 int block_size, aom_noise_model_t *noise_model) { 260 (void)raw; 261 (void)denoised; 262 (void)flat_blocks; 263 (void)block_size; 264 fprintf(debug_file, "figure(3); clf;\n"); 265 fprintf(debug_file, "figure(2); clf;\n"); 266 fprintf(debug_file, "figure(1); clf;\n"); 267 for (int c = 0; c < 3; ++c) { 268 fprintf(debug_file, "noise_strength_%d = [\n", c); 269 const aom_equation_system_t *eqns = 270 &noise_model->combined_state[c].strength_solver.eqns; 271 for (int k = 0; k < eqns->n; ++k) { 272 fprintf(debug_file, "%lf ", eqns->x[k]); 273 } 274 fprintf(debug_file, "];\n"); 275 fprintf(debug_file, "plot(noise_strength_%d); hold on;\n", c); 276 } 277 fprintf(debug_file, "legend('Y', 'cb', 'cr');\n"); 278 fprintf(debug_file, "title('Noise strength function');\n"); 279 280 #if CONFIG_AV1_DECODER 281 aom_film_grain_t grain; 282 aom_noise_model_get_grain_parameters(noise_model, &grain); 283 print_variance_y(debug_file, raw, denoised, flat_blocks, block_size, &grain); 284 #endif 285 fflush(debug_file); 286 } 287 288 int main(int argc, char *argv[]) { 289 noise_model_args_t args = { 0, 0, { 25, 1 }, 0, 0, 0, AOM_IMG_FMT_I420, 290 32, 8, 1, 0, 1, NULL }; 291 aom_image_t raw, denoised; 292 FILE *infile = NULL; 293 AvxVideoInfo info; 294 295 memset(&info, 0, sizeof(info)); 296 297 (void)argc; 298 exec_name = argv[0]; 299 parse_args(&args, argv + 1); 300 301 info.frame_width = args.width; 302 info.frame_height = args.height; 303 info.time_base.numerator = args.fps.den; 304 info.time_base.denominator = args.fps.num; 305 306 if (info.frame_width <= 0 || info.frame_height <= 0 || 307 (info.frame_width % 2) != 0 || (info.frame_height % 2) != 0) { 308 die("Invalid frame size: %dx%d", info.frame_width, info.frame_height); 309 } 310 if (!aom_img_alloc(&raw, args.img_fmt, info.frame_width, info.frame_height, 311 1)) { 312 die("Failed to allocate image."); 313 } 314 if (!aom_img_alloc(&denoised, args.img_fmt, info.frame_width, 315 info.frame_height, 1)) { 316 die("Failed to allocate image."); 317 } 318 infile = fopen(args.input, "rb"); 319 if (!infile) { 320 die("Failed to open input file: %s", args.input); 321 } 322 fprintf(stderr, "Bit depth: %d stride:%d\n", args.bit_depth, raw.stride[0]); 323 324 const int high_bd = args.bit_depth > 8; 325 const int block_size = args.block_size; 326 aom_flat_block_finder_t block_finder; 327 aom_flat_block_finder_init(&block_finder, block_size, args.bit_depth, 328 high_bd); 329 330 const int num_blocks_w = (info.frame_width + block_size - 1) / block_size; 331 const int num_blocks_h = (info.frame_height + block_size - 1) / block_size; 332 uint8_t *flat_blocks = (uint8_t *)aom_malloc(num_blocks_w * num_blocks_h); 333 if (!flat_blocks) die("Failed to allocate block data."); 334 // Sets the random seed on the first entry in the output table 335 int16_t random_seed = 7391; 336 aom_noise_model_t noise_model; 337 aom_noise_model_params_t params = { AOM_NOISE_SHAPE_SQUARE, 3, args.bit_depth, 338 high_bd }; 339 aom_noise_model_init(&noise_model, params); 340 341 FILE *denoised_file = 0; 342 if (args.input_denoised) { 343 denoised_file = fopen(args.input_denoised, "rb"); 344 if (!denoised_file) 345 die("Unable to open input_denoised: %s", args.input_denoised); 346 } else { 347 die("--input-denoised file must be specified"); 348 } 349 FILE *debug_file = 0; 350 if (args.debug_file) { 351 debug_file = fopen(args.debug_file, "w"); 352 } 353 aom_film_grain_table_t grain_table = { 0, 0 }; 354 355 int64_t prev_timestamp = 0; 356 int frame_count = 0; 357 while (aom_img_read(&raw, infile)) { 358 if (args.input_denoised) { 359 if (!aom_img_read(&denoised, denoised_file)) { 360 die("Unable to read input denoised file"); 361 } 362 } 363 if (frame_count % args.skip_frames == 0) { 364 int num_flat_blocks = num_blocks_w * num_blocks_h; 365 memset(flat_blocks, 1, num_flat_blocks); 366 if (args.run_flat_block_finder) { 367 memset(flat_blocks, 0, num_flat_blocks); 368 num_flat_blocks = aom_flat_block_finder_run( 369 &block_finder, raw.planes[0], info.frame_width, info.frame_height, 370 info.frame_width, flat_blocks); 371 fprintf(stdout, "Num flat blocks %d\n", num_flat_blocks); 372 } 373 374 const uint8_t *planes[3] = { raw.planes[0], raw.planes[1], 375 raw.planes[2] }; 376 uint8_t *denoised_planes[3] = { denoised.planes[0], denoised.planes[1], 377 denoised.planes[2] }; 378 int strides[3] = { raw.stride[0] >> high_bd, raw.stride[1] >> high_bd, 379 raw.stride[2] >> high_bd }; 380 int chroma_sub[3] = { raw.x_chroma_shift, raw.y_chroma_shift, 0 }; 381 382 fprintf(stdout, "Updating noise model...\n"); 383 aom_noise_status_t status = aom_noise_model_update( 384 &noise_model, (const uint8_t *const *)planes, 385 (const uint8_t *const *)denoised_planes, info.frame_width, 386 info.frame_height, strides, chroma_sub, flat_blocks, block_size); 387 388 int64_t cur_timestamp = 389 frame_count * 10000000ULL * args.fps.den / args.fps.num; 390 if (status == AOM_NOISE_STATUS_DIFFERENT_NOISE_TYPE) { 391 fprintf(stdout, 392 "Noise type is different, updating parameters for time " 393 "[ %" PRId64 ", %" PRId64 ")\n", 394 prev_timestamp, cur_timestamp); 395 aom_film_grain_t grain; 396 aom_noise_model_get_grain_parameters(&noise_model, &grain); 397 grain.random_seed = random_seed; 398 random_seed = 0; 399 aom_film_grain_table_append(&grain_table, prev_timestamp, cur_timestamp, 400 &grain); 401 aom_noise_model_save_latest(&noise_model); 402 prev_timestamp = cur_timestamp; 403 } 404 if (debug_file) { 405 print_debug_info(debug_file, &raw, &denoised, flat_blocks, block_size, 406 &noise_model); 407 } 408 fprintf(stdout, "Done noise model update, status = %d\n", status); 409 } 410 frame_count++; 411 } 412 413 aom_film_grain_t grain; 414 aom_noise_model_get_grain_parameters(&noise_model, &grain); 415 grain.random_seed = random_seed; 416 aom_film_grain_table_append(&grain_table, prev_timestamp, INT64_MAX, &grain); 417 if (args.output_grain_table) { 418 struct aom_internal_error_info error_info; 419 if (AOM_CODEC_OK != aom_film_grain_table_write(&grain_table, 420 args.output_grain_table, 421 &error_info)) { 422 die("Unable to write output film grain table"); 423 } 424 } 425 aom_film_grain_table_free(&grain_table); 426 427 if (infile) fclose(infile); 428 if (denoised_file) fclose(denoised_file); 429 if (debug_file) fclose(debug_file); 430 aom_img_free(&raw); 431 aom_img_free(&denoised); 432 433 return EXIT_SUCCESS; 434 }