common.rs 23.2 KB
Newer Older
Raphaël Zumer's avatar
Raphaël Zumer committed
1
2
3
4
5
6
7
8
9
// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

10
use crate::muxer::{create_muxer, Muxer};
11
use crate::{ColorPrimaries, MatrixCoefficients, TransferCharacteristics};
12
use clap::{App, AppSettings, Arg, ArgMatches, SubCommand, Shell};
13
use rav1e::partition::BlockSize;
14
use rav1e::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
15

16
use std::fs::File;
17
use std::io::prelude::*;
18
use std::path::PathBuf;
19
use std::time::Instant;
20
use std::{fmt, io};
21
22

pub struct EncoderIO {
Kyle Siefring's avatar
Kyle Siefring committed
23
  pub input: Box<dyn Read>,
24
  pub output: Box<dyn Muxer>,
Kyle Siefring's avatar
Kyle Siefring committed
25
  pub rec: Option<Box<dyn Write>>
26
27
}

28
29
30
31
pub struct CliOptions {
  pub io: EncoderIO,
  pub enc: EncoderConfig,
  pub limit: usize,
Vibhoothi's avatar
Vibhoothi committed
32
  pub skip: usize,
33
  pub verbose: bool,
Luca Barbato's avatar
Luca Barbato committed
34
  pub threads: usize,
35
36
37
}

pub fn parse_cli() -> CliOptions {
38
39
40
  let ver_short = version::short();
  let ver_long = version::full();
  let ver = version::full();
41
  let mut app = App::new("rav1e")
42
43
44
    .version(ver.as_str())
    .long_version(ver_long.as_str())
    .version_short(ver_short.as_str())
Luca Barbato's avatar
Luca Barbato committed
45
    .about("AV1 video encoder")
46
    .setting(AppSettings::DeriveDisplayOrder)
47
    .setting(AppSettings::SubcommandsNegateReqs)
48
49
50
    .arg(Arg::with_name("FULLHELP")
      .help("Prints more detailed help information")
      .long("fullhelp"))
Luca Barbato's avatar
Luca Barbato committed
51
52
53
54
55
56
57
58
    // THREADS
    .arg(
      Arg::with_name("THREADS")
        .help("Set the threadpool size")
        .long("threads")
        .takes_value(true)
        .default_value("0")
    )
59
    // INPUT/OUTPUT
Luca Barbato's avatar
Luca Barbato committed
60
61
62
    .arg(
      Arg::with_name("INPUT")
        .help("Uncompressed YUV4MPEG2 video input")
63
        .required_unless("FULLHELP")
Luca Barbato's avatar
Luca Barbato committed
64
        .index(1)
65
66
    )
    .arg(
Luca Barbato's avatar
Luca Barbato committed
67
68
69
70
      Arg::with_name("OUTPUT")
        .help("Compressed AV1 in IVF video output")
        .short("o")
        .long("output")
71
        .required_unless("FULLHELP")
Luca Barbato's avatar
Luca Barbato committed
72
        .takes_value(true)
73
74
75
76
77
78
79
    )
    .arg(
      Arg::with_name("STATS_FILE")
        .help("Custom location for first-pass stats file")
        .long("stats")
        .takes_value(true)
        .default_value("rav1e_stats.json")
80
81
82
83
84
85
86
    )
    // ENCODING SETTINGS
    .arg(
      Arg::with_name("PASS")
        .help("Specify first-pass or second-pass to run as a two-pass encode; If not provided, will run a one-pass encode")
        .short("p")
        .long("pass")
87
        .takes_value(true)
88
89
90
        .possible_values(&["1", "2"])
    )
    .arg(
Luca Barbato's avatar
Luca Barbato committed
91
92
93
94
95
96
      Arg::with_name("LIMIT")
        .help("Maximum number of frames to encode")
        .short("l")
        .long("limit")
        .takes_value(true)
        .default_value("0")
97
    )
Vibhoothi's avatar
Vibhoothi committed
98
99
100
101
102
103
104
    .arg(
      Arg::with_name("SKIP")
        .help("Skip n number of frames and encode")
        .long("skip")
        .takes_value(true)
        .default_value("0")
    )
105
    .arg(
Luca Barbato's avatar
Luca Barbato committed
106
      Arg::with_name("QP")
107
        .help("Quantizer (0-255), smaller values are higher quality [default: 100]")
Luca Barbato's avatar
Luca Barbato committed
108
109
        .long("quantizer")
        .takes_value(true)
110
111
112
113
114
115
116
    )
    .arg(
      Arg::with_name("BITRATE")
        .help("Bitrate (kbps)")
        .short("b")
        .long("bitrate")
        .takes_value(true)
117
118
    )
    .arg(
Luca Barbato's avatar
Luca Barbato committed
119
      Arg::with_name("SPEED")
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
        .help("Speed level (0 is best quality, 10 is fastest)\n\
        Speeds 10 and 0 are extremes and are generally not recommended")
        .long_help("Speed level (0 is best quality, 10 is fastest)\n\
        Speeds 10 and 0 are extremes and are generally not recommended\n\
        - 10 (fastest):\n\
        Min block size 64x64, TX domain distortion, fast deblock, no scenechange detection\n\
        - 9:\n\
        Min block size 64x64, TX domain distortion, fast deblock\n\
        - 8:\n\
        Min block size 8x8, reduced TX set, TX domain distortion, fast deblock\n\
        - 7:\n\
        Min block size 8x8, reduced TX set, TX domain distortion\n\
        - 6:\n\
        Min block size 8x8, reduced TX set, TX domain distortion\n\
        - 5 (default):\n\
        Min block size 8x8, reduced TX set, TX domain distortion, complex pred modes for keyframes\n\
        - 4:\n\
        Min block size 8x8, TX domain distortion, complex pred modes for keyframes\n\
        - 3:\n\
        Min block size 8x8, TX domain distortion, complex pred modes for keyframes, RDO TX decision\n\
        - 2:\n\
        Min block size 8x8, TX domain distortion, complex pred modes for keyframes, RDO TX decision, include near MVs\n\
        - 1:\n\
        Min block size 8x8, TX domain distortion, complex pred modes, RDO TX decision, include near MVs\n\
        - 0 (slowest):\n\
        Min block size 4x4, TX domain distortion, complex pred modes, RDO TX decision, include near MVs, bottom-up encoding\n")
Luca Barbato's avatar
Luca Barbato committed
146
147
148
        .short("s")
        .long("speed")
        .takes_value(true)
149
        .default_value("5")
150
151
    )
    .arg(
Josh Holmer's avatar
Josh Holmer committed
152
153
154
155
156
157
      Arg::with_name("MIN_KEYFRAME_INTERVAL")
        .help("Minimum interval between keyframes")
        .short("i")
        .long("min-keyint")
        .takes_value(true)
        .default_value("12")
158
159
    )
    .arg(
160
      Arg::with_name("KEYFRAME_INTERVAL")
Josh Holmer's avatar
Josh Holmer committed
161
        .help("Maximum interval between keyframes")
162
163
164
        .short("I")
        .long("keyint")
        .takes_value(true)
165
        .default_value("240")
166
    )
167
168
169
170
171
172
173
    .arg(
      Arg::with_name("RESERVOIR_FRAME_DELAY")
        .help("Number of frames over which rate control should distribute the reservoir [default: max(240, 1.5x keyint)]\n\
         A minimum value of 12 is enforced.")
        .long("reservoir_frame_delay")
        .takes_value(true)
    )
174
    .arg(
175
      Arg::with_name("LOW_LATENCY")
176
177
        .help("Low latency mode; disables frame reordering\n\
            Has a significant speed-to-quality trade-off")
178
        .long("low_latency")
179
180
    )
    .arg(
Luca Barbato's avatar
Luca Barbato committed
181
      Arg::with_name("TUNE")
182
        .help("Quality tuning")
Luca Barbato's avatar
Luca Barbato committed
183
184
        .long("tune")
        .possible_values(&Tune::variants())
185
        .default_value("Psychovisual")
Luca Barbato's avatar
Luca Barbato committed
186
        .case_insensitive(true)
187
    )
188
189
190
191
192
193
194
195
196
197
198
199
200
201
    .arg(
      Arg::with_name("TILE_ROWS_LOG2")
        .help("Log2 of number of tile rows")
        .long("tile-rows-log2")
        .takes_value(true)
        .default_value("0")
    )
    .arg(
      Arg::with_name("TILE_COLS_LOG2")
        .help("Log2 of number of tile columns")
        .long("tile-cols-log2")
        .takes_value(true)
        .default_value("0")
    )
202
203
    // MASTERING
    .arg(
204
      Arg::with_name("PIXEL_RANGE")
205
206
207
208
209
210
211
        .help("Pixel range")
        .long("range")
        .possible_values(&PixelRange::variants())
        .default_value("unspecified")
        .case_insensitive(true)
    )
    .arg(
212
      Arg::with_name("COLOR_PRIMARIES")
213
214
215
216
217
218
219
        .help("Color primaries used to describe color parameters")
        .long("primaries")
        .possible_values(&ColorPrimaries::variants())
        .default_value("unspecified")
        .case_insensitive(true)
    )
    .arg(
220
      Arg::with_name("TRANSFER_CHARACTERISTICS")
221
222
223
224
225
226
227
        .help("Transfer characteristics used to describe color parameters")
        .long("transfer")
        .possible_values(&TransferCharacteristics::variants())
        .default_value("unspecified")
        .case_insensitive(true)
    )
    .arg(
228
      Arg::with_name("MATRIX_COEFFICIENTS")
229
230
231
232
233
234
235
        .help("Matrix coefficients used to describe color parameters")
        .long("matrix")
        .possible_values(&MatrixCoefficients::variants())
        .default_value("unspecified")
        .case_insensitive(true)
    )
    .arg(
236
      Arg::with_name("MASTERING_DISPLAY")
237
238
239
240
241
242
        .help("Mastering display primaries in the form of G(x,y)B(x,y)R(x,y)WP(x,y)L(max,min)")
        .long("mastering_display")
        .default_value("unspecified")
        .case_insensitive(true)
    )
    .arg(
243
      Arg::with_name("CONTENT_LIGHT")
244
245
246
247
248
249
250
        .help("Content light level used to describe content luminosity (cll,fall)")
        .long("content_light")
        .default_value("0,0")
        .case_insensitive(true)
    )
    // DEBUGGING
    .arg(
251
      Arg::with_name("VERBOSE")
252
        .help("Verbose logging; outputs info for every frame")
253
254
        .long("verbose")
        .short("v")
255
256
    )
    .arg(
257
      Arg::with_name("PSNR")
258
        .help("Calculate and display PSNR metrics")
259
        .long("psnr")
260
261
262
263
    )
    .arg(
      Arg::with_name("RECONSTRUCTION")
        .help("Outputs a Y4M file containing the output from the decoder")
264
        .long("reconstruction")
265
266
267
        .short("r")
        .takes_value(true)
    )
268
269
    .arg(
      Arg::with_name("SPEED_TEST")
270
        .help("Run an encode using default encoding settings, manually adjusting only the settings specified; allows benchmarking settings in isolation")
271
272
273
274
        .hidden(true)
        .long("speed-test")
        .takes_value(true)
    )
275
276
277
278
    .arg(
      Arg::with_name("train-rdo")
        .long("train-rdo")
    )
279
    .subcommand(SubCommand::with_name("advanced")
280
                .setting(AppSettings::Hidden)
281
282
283
284
285
286
287
288
289
290
291
                .about("Advanced features")
                .arg(Arg::with_name("SHELL")
                     .help("Output to stdout the completion definition for the shell")
                     .short("c")
                     .long("completion")
                     .takes_value(true)
                     .possible_values(&Shell::variants())
                )
    );

  let matches = app.clone().get_matches();
292

293
294
295
296
297
  if matches.is_present("FULLHELP") {
    app.print_long_help().unwrap();
    std::process::exit(0);
  }

Luca Barbato's avatar
Luca Barbato committed
298
  let threads = matches.value_of("THREADS").map(|v| v.parse().expect("Threads must be an integer")).unwrap();
Luca Barbato's avatar
Luca Barbato committed
299

300
301
302
303
304
305
306
  if let Some(matches) = matches.subcommand_matches("advanced") {
    if let Some(shell) = matches.value_of("SHELL").map(|v| v.parse().unwrap()) {
      app.gen_completions_to("rav1e", shell, &mut std::io::stdout());
      std::process::exit(0);
    }
  }

Luca Barbato's avatar
Luca Barbato committed
307
308
309
310
311
  let io = EncoderIO {
    input: match matches.value_of("INPUT").unwrap() {
      "-" => Box::new(io::stdin()) as Box<dyn Read>,
      f => Box::new(File::open(&f).unwrap()) as Box<dyn Read>
    },
312
    output: create_muxer(matches.value_of("OUTPUT").unwrap()),
Luca Barbato's avatar
Luca Barbato committed
313
314
315
316
    rec: matches
      .value_of("RECONSTRUCTION")
      .map(|f| Box::new(File::create(&f).unwrap()) as Box<dyn Write>)
  };
Kyle Siefring's avatar
Kyle Siefring committed
317

318
319
  CliOptions {
    io,
320
    enc: parse_config(&matches),
321
    limit: matches.value_of("LIMIT").unwrap().parse().unwrap(),
Vibhoothi's avatar
Vibhoothi committed
322
    skip: matches.value_of("SKIP").unwrap().parse().unwrap(),
323
    verbose: matches.is_present("VERBOSE"),
Luca Barbato's avatar
Luca Barbato committed
324
    threads,
325
326
327
  }
}

328
fn parse_config(matches: &ArgMatches<'_>) -> EncoderConfig {
329
330
331
332
333
334
335
336
337
338
339
340
341
  let maybe_quantizer = matches.value_of("QP").map(|qp| qp.parse().unwrap());
  let maybe_bitrate =
    matches.value_of("BITRATE").map(|bitrate| bitrate.parse().unwrap());
  let quantizer = maybe_quantizer.unwrap_or_else(|| {
    if maybe_bitrate.is_some() {
      // If a bitrate is specified, the quantizer is the maximum allowed (e.g.,
      //  the minimum quality allowed), which by default should be
      //  unconstrained.
      255
    } else {
      100
    }
  });
342
  let bitrate: i32 = maybe_bitrate.unwrap_or(0);
343
  let train_rdo = matches.is_present("train-rdo");
344
345
  if quantizer == 0 {
    unimplemented!("Lossless encoding not yet implemented");
346
347
  } else if quantizer > 255 {
    panic!("Quantizer must be between 0-255");
348
349
  }

350
  let mut cfg = if let Some(settings) = matches.value_of("SPEED_TEST") {
351
    eprintln!("Running in speed test mode--ignoring other settings");
352
353
354
355
356
    let mut cfg = EncoderConfig::default();
    settings
      .split_whitespace()
      .for_each(|setting| apply_speed_test_cfg(&mut cfg, setting));
    cfg
357
358
359
360
  } else {
    let speed = matches.value_of("SPEED").unwrap().parse().unwrap();
    let max_interval: u64 = matches.value_of("KEYFRAME_INTERVAL").unwrap().parse().unwrap();
    let mut min_interval: u64 = matches.value_of("MIN_KEYFRAME_INTERVAL").unwrap().parse().unwrap();
361

362
363
364
365
366
    if matches.occurrences_of("MIN_KEYFRAME_INTERVAL") == 0 {
      min_interval = min_interval.min(max_interval);
    }

    // Validate arguments
367
368
369
370
    if max_interval < 1 {
      panic!("Keyframe interval must be greater than 0");
    }

371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
    if speed > 10 {
      panic!("Speed must be between 0-10");
    } else if min_interval > max_interval {
      panic!("Maximum keyframe interval must be greater than or equal to minimum keyframe interval");
    }

    let color_primaries =
      matches.value_of("COLOR_PRIMARIES").unwrap().parse().unwrap_or_default();
    let transfer_characteristics = matches
      .value_of("TRANSFER_CHARACTERISTICS")
      .unwrap()
      .parse()
      .unwrap_or_default();
    let matrix_coefficients = matches
      .value_of("MATRIX_COEFFICIENTS")
      .unwrap()
      .parse()
      .unwrap_or_default();
389

390
391
392
393
394
395
396
397
    let mut cfg = EncoderConfig::with_speed_preset(speed);
    cfg.max_key_frame_interval = min_interval;
    cfg.max_key_frame_interval = max_interval;

    cfg.pixel_range = matches.value_of("PIXEL_RANGE").unwrap().parse().unwrap_or_default();
    cfg.color_description = if color_primaries == ColorPrimaries::Unspecified &&
      transfer_characteristics == TransferCharacteristics::Unspecified &&
      matrix_coefficients == MatrixCoefficients::Unspecified {
398
399
400
401
402
403
404
405
406
      // No need to set a color description with all parameters unspecified.
      None
    } else {
      Some(ColorDescription {
        color_primaries,
        transfer_characteristics,
        matrix_coefficients
      })
    };
407
408
409

    let mastering_display_opt = matches.value_of("MASTERING_DISPLAY").unwrap();
    cfg.mastering_display = if mastering_display_opt == "unspecified" { None } else {
410
      let (g_x, g_y, b_x, b_y, r_x, r_y, wp_x, wp_y, max_lum, min_lum) = scan_fmt!(mastering_display_opt, "G({},{})B({},{})R({},{})WP({},{})L({},{})", f64, f64, f64, f64, f64, f64, f64, f64, f64, f64).expect("Cannot parse the mastering display option");
411
412
413
      Some(MasteringDisplay {
        primaries: [
          Point {
414
415
            x: (r_x * ((1 << 16) as f64)).round() as u16,
            y: (r_y * ((1 << 16) as f64)).round() as u16,
416
417
          },
          Point {
418
419
            x: (g_x * ((1 << 16) as f64)).round() as u16,
            y: (g_y * ((1 << 16) as f64)).round() as u16,
420
421
          },
          Point {
422
423
            x: (b_x * ((1 << 16) as f64)).round() as u16,
            y: (b_y * ((1 << 16) as f64)).round() as u16,
424
425
426
          }
        ],
        white_point: Point {
427
428
          x: (wp_x * ((1 << 16) as f64)).round() as u16,
          y: (wp_y * ((1 << 16) as f64)).round() as u16,
429
        },
430
431
        max_luminance: (max_lum * ((1 << 8) as f64)).round() as u32,
        min_luminance: (min_lum * ((1 << 14) as f64)).round() as u32,
432
433
434
435
      })
    };

    let content_light_opt = matches.value_of("CONTENT_LIGHT").unwrap();
436
437
    let (cll, fall) = scan_fmt!(content_light_opt, "{},{}", u16, u16).expect("Cannot parse the content light option");
    cfg.content_light = if cll == 0 && fall == 0 { None } else {
438
      Some(ContentLight {
439
440
        max_content_light_level: cll,
        max_frame_average_light_level: fall
441
442
443
444
445
      })
    };
    cfg
  };

446
  cfg.quantizer = quantizer;
447
  cfg.bitrate = bitrate.checked_mul(1000).expect("Bitrate too high");
448
  cfg.reservoir_frame_delay = matches.value_of("RESERVOIR_FRAME_DELAY").map(|reservior_frame_delay| reservior_frame_delay.parse().unwrap());
449
  cfg.show_psnr = matches.is_present("PSNR");
450
451
452
453
454
455
  cfg.pass = matches.value_of("PASS").map(|pass| pass.parse().unwrap());
  cfg.stats_file = if cfg.pass.is_some() {
    Some(PathBuf::from(matches.value_of("STATS_FILE").unwrap()))
  } else {
    None
  };
456
  cfg.tune = matches.value_of("TUNE").unwrap().parse().unwrap();
457
458
459
460
461
462
463
464

  cfg.tile_cols_log2 = matches.value_of("TILE_COLS_LOG2").unwrap().parse().unwrap();
  cfg.tile_rows_log2 = matches.value_of("TILE_ROWS_LOG2").unwrap().parse().unwrap();

  if cfg.tile_cols_log2 > 6 || cfg.tile_rows_log2 > 6 {
    panic!("Log2 of tile columns and rows may not be greater than 6");
  }

465
  cfg.low_latency = matches.is_present("LOW_LATENCY");
466
  cfg.train_rdo = train_rdo;
467
468
  cfg
}
469

470
fn apply_speed_test_cfg(cfg: &mut EncoderConfig, setting: &str) {
471
472
  match setting {
    "baseline" => {
473
      cfg.speed_settings = SpeedSettings::default();
474
475
476
477
478
479
480
481
482
483
    },
    "min_block_size_4x4" => {
      cfg.speed_settings.min_block_size = BlockSize::BLOCK_4X4;
    },
    "min_block_size_8x8" => {
      cfg.speed_settings.min_block_size = BlockSize::BLOCK_8X8;
    },
    "min_block_size_32x32" => {
      cfg.speed_settings.min_block_size = BlockSize::BLOCK_32X32;
    },
484
485
486
    "min_block_size_64x64" => {
      cfg.speed_settings.min_block_size = BlockSize::BLOCK_64X64;
    },
487
488
489
490
491
492
493
494
495
496
497
498
    "multiref" => {
      cfg.speed_settings.multiref = true;
    },
    "fast_deblock" => {
      cfg.speed_settings.fast_deblock = true;
    },
    "reduced_tx_set" => {
      cfg.speed_settings.reduced_tx_set = true;
    },
    "tx_domain_distortion" => {
      cfg.speed_settings.tx_domain_distortion = true;
    },
499
500
501
    "tx_domain_rate" => {
      cfg.speed_settings.tx_domain_rate = true;
    },
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
    "encode_bottomup" => {
      cfg.speed_settings.encode_bottomup = true;
    },
    "rdo_tx_decision" => {
      cfg.speed_settings.rdo_tx_decision = true;
    },
    "prediction_modes_keyframes" => {
      cfg.speed_settings.prediction_modes = PredictionModesSetting::ComplexKeyframes;
    },
    "prediction_modes_all" => {
      cfg.speed_settings.prediction_modes = PredictionModesSetting::ComplexAll;
    },
    "include_near_mvs" => {
      cfg.speed_settings.include_near_mvs = true;
    },
    "no_scene_detection" => {
      cfg.speed_settings.no_scene_detection = true;
    },
520
521
522
    "diamond_me" => {
      cfg.speed_settings.diamond_me = true;
    }
523
524
525
    "cdef" => {
      cfg.speed_settings.cdef = true;
    }
526
527
    setting => {
      panic!("Unrecognized speed test setting {}", setting);
528
    }
529
  };
530
531
}

532
533
#[derive(Debug, Clone, Copy)]
pub struct FrameSummary {
534
  // Frame size in bytes
535
536
  pub size: usize,
  pub number: u64,
537
  pub frame_type: FrameType,
538
  // PSNR for Y, U, and V planes
539
  pub psnr: Option<(f64, f64, f64)>,
540
541
}

542
543
impl<T: Pixel> From<Packet<T>> for FrameSummary {
  fn from(packet: Packet<T>) -> Self {
544
545
546
547
    Self {
      size: packet.data.len(),
      number: packet.number,
      frame_type: packet.frame_type,
548
      psnr: packet.psnr,
549
550
551
552
553
554
555
556
    }
  }
}

impl fmt::Display for FrameSummary {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(
      f,
557
      "Frame {} - {} - {} bytes{}",
558
559
      self.number,
      self.frame_type,
560
561
562
563
      self.size,
      if let Some(psnr) = self.psnr {
        format!(" - PSNR: Y: {:.4}  Cb: {:.4}  Cr: {:.4}", psnr.0, psnr.1, psnr.2)
      } else { String::new() }
564
565
    )
  }
566
}
567

568
569
#[derive(Debug, Clone)]
pub struct ProgressInfo {
570
  // Frame rate of the video
571
  frame_rate: Rational,
572
  // The length of the whole video, in frames, if known
573
  total_frames: Option<usize>,
574
  // The time the encode was started
575
  time_started: Instant,
576
  // List of frames encoded so far
577
  frame_info: Vec<FrameSummary>,
578
579
580
581
  // Video size so far in bytes.
  //
  // This value will be updated in the CLI very frequently, so we cache the previous value
  // to reduce the overall complexity.
582
  encoded_size: usize,
583
  // Whether to display PSNR statistics during and at end of encode
584
  show_psnr: bool,
585
586
587
}

impl ProgressInfo {
588
  pub fn new(frame_rate: Rational, total_frames: Option<usize>, show_psnr: bool) -> Self {
589
590
591
592
593
594
    Self {
      frame_rate,
      total_frames,
      time_started: Instant::now(),
      frame_info: Vec::with_capacity(total_frames.unwrap_or_default()),
      encoded_size: 0,
595
      show_psnr,
596
597
598
599
600
601
602
603
604
605
606
607
608
    }
  }

  pub fn add_frame(&mut self, frame: FrameSummary) {
    self.encoded_size += frame.size;
    self.frame_info.push(frame);
  }

  pub fn frames_encoded(&self) -> usize {
    self.frame_info.len()
  }

  pub fn encoding_fps(&self) -> f64 {
609
610
    let duration = Instant::now().duration_since(self.time_started);
    self.frame_info.len() as f64 / (duration.as_secs() as f64 + duration.subsec_millis() as f64 / 1000f64)
611
612
613
614
615
616
  }

  pub fn video_fps(&self) -> f64 {
    self.frame_rate.num as f64 / self.frame_rate.den as f64
  }

617
  // Returns the bitrate of the frames so far, in bits/second
618
619
620
621
622
623
  pub fn bitrate(&self) -> usize {
    let bits = self.encoded_size * 8;
    let seconds = self.frame_info.len() as f64 / self.video_fps();
    (bits as f64 / seconds) as usize
  }

624
  // Estimates the final filesize in bytes, if the number of frames is known
625
626
627
628
629
  pub fn estimated_size(&self) -> usize {
    self.total_frames
      .map(|frames| self.encoded_size * frames / self.frames_encoded())
      .unwrap_or_default()
  }
630

631
632
633
634
635
636
  // Estimates the remaining encoding time in seconds, if the number of frames is known
  pub fn estimated_time(&self) -> f64 {
    self.total_frames
      .map(|frames| (frames - self.frames_encoded()) as f64 / self.encoding_fps())
      .unwrap_or_default()
  }
637

638
  // Number of frames of given type which appear in the video
639
640
641
642
643
644
  pub fn get_frame_type_count(&self, frame_type: FrameType) -> usize {
    self.frame_info.iter()
      .filter(|frame| frame.frame_type == frame_type)
      .count()
  }

645
  // Size in bytes of all frames of given frame type
646
647
648
649
650
651
652
  pub fn get_frame_type_size(&self, frame_type: FrameType) -> usize {
    self.frame_info.iter()
      .filter(|frame| frame.frame_type == frame_type)
      .map(|frame| frame.size)
      .sum()
  }

653
  pub fn print_summary(&self) -> String {
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
    let (key, key_size) = (
      self.get_frame_type_count(FrameType::KEY),
      self.get_frame_type_size(FrameType::KEY)
    );
    let (inter, inter_size) = (
      self.get_frame_type_count(FrameType::INTER),
      self.get_frame_type_size(FrameType::INTER)
    );
    let (ionly, ionly_size) = (
      self.get_frame_type_count(FrameType::INTRA_ONLY),
      self.get_frame_type_size(FrameType::INTRA_ONLY)
    );
    let (switch, switch_size) = (
      self.get_frame_type_count(FrameType::SWITCH),
      self.get_frame_type_size(FrameType::SWITCH)
    );
670
671
672
673
    format!("\
    Key Frames: {:>6}    avg size: {:>7} B\n\
    Inter:      {:>6}    avg size: {:>7} B\n\
    Intra Only: {:>6}    avg size: {:>7} B\n\
674
675
    Switch:     {:>6}    avg size: {:>7} B\
    {}",
676
      key, key_size / key,
677
      inter, inter_size.checked_div(inter).unwrap_or(0),
678
      ionly, ionly_size / key,
679
680
      switch, switch_size / key,
      if self.show_psnr {
681
682
683
684
685
686
687
688
689
        let psnr_y =
          self.frame_info.iter().map(|fi| fi.psnr.unwrap().0).sum::<f64>()
            / self.frame_info.len() as f64;
        let psnr_u =
          self.frame_info.iter().map(|fi| fi.psnr.unwrap().1).sum::<f64>()
            / self.frame_info.len() as f64;
        let psnr_v =
          self.frame_info.iter().map(|fi| fi.psnr.unwrap().2).sum::<f64>()
            / self.frame_info.len() as f64;
690
691
692
693
        format!("\nMean PSNR: Y: {:.4}  Cb: {:.4}  Cr: {:.4}  Avg: {:.4}",
                psnr_y, psnr_u, psnr_v,
                (psnr_y + psnr_u + psnr_v) / 3.0)
      } else { String::new() }
694
695
696
697
698
699
700
701
702
    )
  }
}

impl fmt::Display for ProgressInfo {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    if let Some(total_frames) = self.total_frames {
      write!(
        f,
703
        "encoded {}/{} frames, {:.3} fps, {:.2} Kb/s, est. size: {:.2} MB, est. time: {:.0} s",
704
705
706
707
        self.frames_encoded(),
        total_frames,
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64,
708
709
        self.estimated_size() as f64 / (1024 * 1024) as f64,
        self.estimated_time()
710
711
712
713
      )
    } else {
      write!(
        f,
714
        "encoded {} frames, {:.3} fps, {:.2} Kb/s",
715
716
717
718
        self.frames_encoded(),
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64
      )
719
    }
Kyle Siefring's avatar
Kyle Siefring committed
720
  }
721
}