common.rs 22.7 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::{ColorPrimaries, MatrixCoefficients, TransferCharacteristics};
11
use clap::{App, AppSettings, Arg, ArgMatches, SubCommand, Shell};
12
use rav1e::partition::BlockSize;
13
use rav1e::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
14

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

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

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

pub fn parse_cli() -> CliOptions {
37
  let mut app = App::new("rav1e")
38
    .version(env!("CARGO_PKG_VERSION"))
Luca Barbato's avatar
Luca Barbato committed
39
    .about("AV1 video encoder")
40
    .setting(AppSettings::DeriveDisplayOrder)
41
    .setting(AppSettings::SubcommandsNegateReqs)
42 43 44
    .arg(Arg::with_name("FULLHELP")
      .help("Prints more detailed help information")
      .long("fullhelp"))
Luca Barbato's avatar
Luca Barbato committed
45 46 47 48 49 50 51 52
    // THREADS
    .arg(
      Arg::with_name("THREADS")
        .help("Set the threadpool size")
        .long("threads")
        .takes_value(true)
        .default_value("0")
    )
53
    // INPUT/OUTPUT
Luca Barbato's avatar
Luca Barbato committed
54 55 56
    .arg(
      Arg::with_name("INPUT")
        .help("Uncompressed YUV4MPEG2 video input")
57
        .required_unless("FULLHELP")
Luca Barbato's avatar
Luca Barbato committed
58
        .index(1)
59 60
    )
    .arg(
Luca Barbato's avatar
Luca Barbato committed
61 62 63 64
      Arg::with_name("OUTPUT")
        .help("Compressed AV1 in IVF video output")
        .short("o")
        .long("output")
65
        .required_unless("FULLHELP")
Luca Barbato's avatar
Luca Barbato committed
66
        .takes_value(true)
67 68 69 70 71 72 73
    )
    .arg(
      Arg::with_name("STATS_FILE")
        .help("Custom location for first-pass stats file")
        .long("stats")
        .takes_value(true)
        .default_value("rav1e_stats.json")
74 75 76 77 78 79 80
    )
    // 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")
81
        .takes_value(true)
82 83 84
        .possible_values(&["1", "2"])
    )
    .arg(
Luca Barbato's avatar
Luca Barbato committed
85 86 87 88 89 90
      Arg::with_name("LIMIT")
        .help("Maximum number of frames to encode")
        .short("l")
        .long("limit")
        .takes_value(true)
        .default_value("0")
91
    )
Vibhoothi's avatar
Vibhoothi committed
92 93 94 95 96 97 98
    .arg(
      Arg::with_name("SKIP")
        .help("Skip n number of frames and encode")
        .long("skip")
        .takes_value(true)
        .default_value("0")
    )
99
    .arg(
Luca Barbato's avatar
Luca Barbato committed
100
      Arg::with_name("QP")
101
        .help("Quantizer (0-255), smaller values are higher quality [default: 100]")
Luca Barbato's avatar
Luca Barbato committed
102 103
        .long("quantizer")
        .takes_value(true)
104 105 106 107 108 109 110
    )
    .arg(
      Arg::with_name("BITRATE")
        .help("Bitrate (kbps)")
        .short("b")
        .long("bitrate")
        .takes_value(true)
111 112
    )
    .arg(
Luca Barbato's avatar
Luca Barbato committed
113
      Arg::with_name("SPEED")
114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139
        .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
140 141 142
        .short("s")
        .long("speed")
        .takes_value(true)
143
        .default_value("5")
144 145
    )
    .arg(
Josh Holmer's avatar
Josh Holmer committed
146 147 148 149 150 151
      Arg::with_name("MIN_KEYFRAME_INTERVAL")
        .help("Minimum interval between keyframes")
        .short("i")
        .long("min-keyint")
        .takes_value(true)
        .default_value("12")
152 153
    )
    .arg(
154
      Arg::with_name("KEYFRAME_INTERVAL")
Josh Holmer's avatar
Josh Holmer committed
155
        .help("Maximum interval between keyframes")
156 157 158
        .short("I")
        .long("keyint")
        .takes_value(true)
159
        .default_value("240")
160 161
    )
    .arg(
162
      Arg::with_name("LOW_LATENCY")
163 164
        .help("Low latency mode; disables frame reordering\n\
            Has a significant speed-to-quality trade-off")
165
        .long("low_latency")
166 167
    )
    .arg(
Luca Barbato's avatar
Luca Barbato committed
168
      Arg::with_name("TUNE")
169
        .help("Quality tuning")
Luca Barbato's avatar
Luca Barbato committed
170 171
        .long("tune")
        .possible_values(&Tune::variants())
172
        .default_value("Psychovisual")
Luca Barbato's avatar
Luca Barbato committed
173
        .case_insensitive(true)
174
    )
175 176 177 178 179 180 181 182 183 184 185 186 187 188
    .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")
    )
189 190
    // MASTERING
    .arg(
191
      Arg::with_name("PIXEL_RANGE")
192 193 194 195 196 197 198
        .help("Pixel range")
        .long("range")
        .possible_values(&PixelRange::variants())
        .default_value("unspecified")
        .case_insensitive(true)
    )
    .arg(
199
      Arg::with_name("COLOR_PRIMARIES")
200 201 202 203 204 205 206
        .help("Color primaries used to describe color parameters")
        .long("primaries")
        .possible_values(&ColorPrimaries::variants())
        .default_value("unspecified")
        .case_insensitive(true)
    )
    .arg(
207
      Arg::with_name("TRANSFER_CHARACTERISTICS")
208 209 210 211 212 213 214
        .help("Transfer characteristics used to describe color parameters")
        .long("transfer")
        .possible_values(&TransferCharacteristics::variants())
        .default_value("unspecified")
        .case_insensitive(true)
    )
    .arg(
215
      Arg::with_name("MATRIX_COEFFICIENTS")
216 217 218 219 220 221 222
        .help("Matrix coefficients used to describe color parameters")
        .long("matrix")
        .possible_values(&MatrixCoefficients::variants())
        .default_value("unspecified")
        .case_insensitive(true)
    )
    .arg(
223
      Arg::with_name("MASTERING_DISPLAY")
224 225 226 227 228 229
        .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(
230
      Arg::with_name("CONTENT_LIGHT")
231 232 233 234 235 236 237
        .help("Content light level used to describe content luminosity (cll,fall)")
        .long("content_light")
        .default_value("0,0")
        .case_insensitive(true)
    )
    // DEBUGGING
    .arg(
238
      Arg::with_name("VERBOSE")
239
        .help("Verbose logging; outputs info for every frame")
240 241
        .long("verbose")
        .short("v")
242 243
    )
    .arg(
244
      Arg::with_name("PSNR")
245
        .help("Calculate and display PSNR metrics")
246
        .long("psnr")
247 248 249 250 251 252 253
    )
    .arg(
      Arg::with_name("RECONSTRUCTION")
        .help("Outputs a Y4M file containing the output from the decoder")
        .short("r")
        .takes_value(true)
    )
254 255
    .arg(
      Arg::with_name("SPEED_TEST")
256
        .help("Run an encode using default encoding settings, manually adjusting only the settings specified; allows benchmarking settings in isolation")
257 258 259 260
        .hidden(true)
        .long("speed-test")
        .takes_value(true)
    )
261 262 263 264
    .arg(
      Arg::with_name("train-rdo")
        .long("train-rdo")
    )
265
    .subcommand(SubCommand::with_name("advanced")
266
                .setting(AppSettings::Hidden)
267 268 269 270 271 272 273 274 275 276 277
                .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();
278

279 280 281 282 283
  if matches.is_present("FULLHELP") {
    app.print_long_help().unwrap();
    std::process::exit(0);
  }

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

286 287 288 289 290 291 292
  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
293 294 295 296 297 298 299 300 301 302 303 304 305
  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>
    },
    output: match matches.value_of("OUTPUT").unwrap() {
      "-" => Box::new(io::stdout()) as Box<dyn Write>,
      f => Box::new(File::create(&f).unwrap()) as Box<dyn Write>
    },
    rec: matches
      .value_of("RECONSTRUCTION")
      .map(|f| Box::new(File::create(&f).unwrap()) as Box<dyn Write>)
  };
Kyle Siefring's avatar
Kyle Siefring committed
306

307 308
  CliOptions {
    io,
309
    enc: parse_config(&matches),
310
    limit: matches.value_of("LIMIT").unwrap().parse().unwrap(),
Vibhoothi's avatar
Vibhoothi committed
311
    skip: matches.value_of("SKIP").unwrap().parse().unwrap(),
312
    verbose: matches.is_present("VERBOSE"),
Luca Barbato's avatar
Luca Barbato committed
313
    threads,
314 315 316
  }
}

317
fn parse_config(matches: &ArgMatches<'_>) -> EncoderConfig {
318 319 320 321 322 323 324 325 326 327 328 329 330 331
  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
    }
  });
  let bitrate = maybe_bitrate.unwrap_or(0);
332
  let train_rdo = matches.is_present("train-rdo");
333 334
  if quantizer == 0 {
    unimplemented!("Lossless encoding not yet implemented");
335 336
  } else if quantizer > 255 {
    panic!("Quantizer must be between 0-255");
337 338
  }

339
  let mut cfg = if let Some(settings) = matches.value_of("SPEED_TEST") {
340
    eprintln!("Running in speed test mode--ignoring other settings");
341 342 343 344 345
    let mut cfg = EncoderConfig::default();
    settings
      .split_whitespace()
      .for_each(|setting| apply_speed_test_cfg(&mut cfg, setting));
    cfg
346 347 348 349
  } 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();
350

351 352 353 354 355
    if matches.occurrences_of("MIN_KEYFRAME_INTERVAL") == 0 {
      min_interval = min_interval.min(max_interval);
    }

    // Validate arguments
356 357 358 359
    if max_interval < 1 {
      panic!("Keyframe interval must be greater than 0");
    }

360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377
    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();
378

379 380 381 382 383 384 385 386
    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 {
387 388 389 390 391 392 393 394 395
      // No need to set a color description with all parameters unspecified.
      None
    } else {
      Some(ColorDescription {
        color_primaries,
        transfer_characteristics,
        matrix_coefficients
      })
    };
396 397 398

    let mastering_display_opt = matches.value_of("MASTERING_DISPLAY").unwrap();
    cfg.mastering_display = if mastering_display_opt == "unspecified" { None } else {
399
      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");
400 401 402
      Some(MasteringDisplay {
        primaries: [
          Point {
403 404
            x: (r_x * ((1 << 16) as f64)).round() as u16,
            y: (r_y * ((1 << 16) as f64)).round() as u16,
405 406
          },
          Point {
407 408
            x: (g_x * ((1 << 16) as f64)).round() as u16,
            y: (g_y * ((1 << 16) as f64)).round() as u16,
409 410
          },
          Point {
411 412
            x: (b_x * ((1 << 16) as f64)).round() as u16,
            y: (b_y * ((1 << 16) as f64)).round() as u16,
413 414 415
          }
        ],
        white_point: Point {
416 417
          x: (wp_x * ((1 << 16) as f64)).round() as u16,
          y: (wp_y * ((1 << 16) as f64)).round() as u16,
418
        },
419 420
        max_luminance: (max_lum * ((1 << 8) as f64)).round() as u32,
        min_luminance: (min_lum * ((1 << 14) as f64)).round() as u32,
421 422 423 424
      })
    };

    let content_light_opt = matches.value_of("CONTENT_LIGHT").unwrap();
425 426
    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 {
427
      Some(ContentLight {
428 429
        max_content_light_level: cll,
        max_frame_average_light_level: fall
430 431 432 433 434
      })
    };
    cfg
  };

435
  cfg.quantizer = quantizer;
436
  cfg.bitrate = bitrate;
437
  cfg.show_psnr = matches.is_present("PSNR");
438 439 440 441 442 443
  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
  };
444
  cfg.tune = matches.value_of("TUNE").unwrap().parse().unwrap();
445 446 447 448 449 450 451 452

  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");
  }

453
  cfg.low_latency = matches.is_present("LOW_LATENCY");
454
  cfg.train_rdo = train_rdo;
455 456
  cfg
}
457

458
fn apply_speed_test_cfg(cfg: &mut EncoderConfig, setting: &str) {
459 460
  match setting {
    "baseline" => {
461
      cfg.speed_settings = SpeedSettings::default();
462 463 464 465 466 467 468 469 470 471
    },
    "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;
    },
472 473 474
    "min_block_size_64x64" => {
      cfg.speed_settings.min_block_size = BlockSize::BLOCK_64X64;
    },
475 476 477 478 479 480 481 482 483 484 485 486
    "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;
    },
487 488 489
    "tx_domain_rate" => {
      cfg.speed_settings.tx_domain_rate = true;
    },
490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507
    "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;
    },
508 509 510
    "diamond_me" => {
      cfg.speed_settings.diamond_me = true;
    }
511 512 513
    "cdef" => {
      cfg.speed_settings.cdef = true;
    }
514 515
    setting => {
      panic!("Unrecognized speed test setting {}", setting);
516
    }
517
  };
518 519
}

520 521
#[derive(Debug, Clone, Copy)]
pub struct FrameSummary {
522
  // Frame size in bytes
523 524
  pub size: usize,
  pub number: u64,
525
  pub frame_type: FrameType,
526
  // PSNR for Y, U, and V planes
527
  pub psnr: Option<(f64, f64, f64)>,
528 529
}

530 531
impl<T: Pixel> From<Packet<T>> for FrameSummary {
  fn from(packet: Packet<T>) -> Self {
532 533 534 535
    Self {
      size: packet.data.len(),
      number: packet.number,
      frame_type: packet.frame_type,
536
      psnr: packet.psnr,
537 538 539 540 541 542 543 544
    }
  }
}

impl fmt::Display for FrameSummary {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(
      f,
545
      "Frame {} - {} - {} bytes{}",
546 547
      self.number,
      self.frame_type,
548 549 550 551
      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() }
552 553
    )
  }
554
}
555

556 557
#[derive(Debug, Clone)]
pub struct ProgressInfo {
558
  // Frame rate of the video
559
  frame_rate: Rational,
560
  // The length of the whole video, in frames, if known
561
  total_frames: Option<usize>,
562
  // The time the encode was started
563
  time_started: Instant,
564
  // List of frames encoded so far
565
  frame_info: Vec<FrameSummary>,
566 567 568 569
  // 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.
570
  encoded_size: usize,
571
  // Whether to display PSNR statistics during and at end of encode
572
  show_psnr: bool,
573 574 575
}

impl ProgressInfo {
576
  pub fn new(frame_rate: Rational, total_frames: Option<usize>, show_psnr: bool) -> Self {
577 578 579 580 581 582
    Self {
      frame_rate,
      total_frames,
      time_started: Instant::now(),
      frame_info: Vec::with_capacity(total_frames.unwrap_or_default()),
      encoded_size: 0,
583
      show_psnr,
584 585 586 587 588 589 590 591 592 593 594 595 596
    }
  }

  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 {
597 598
    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)
599 600 601 602 603 604
  }

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

605
  // Returns the bitrate of the frames so far, in bits/second
606 607 608 609 610 611
  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
  }

612
  // Estimates the final filesize in bytes, if the number of frames is known
613 614 615 616 617
  pub fn estimated_size(&self) -> usize {
    self.total_frames
      .map(|frames| self.encoded_size * frames / self.frames_encoded())
      .unwrap_or_default()
  }
618

619 620 621 622 623 624
  // 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()
  }
625

626
  // Number of frames of given type which appear in the video
627 628 629 630 631 632
  pub fn get_frame_type_count(&self, frame_type: FrameType) -> usize {
    self.frame_info.iter()
      .filter(|frame| frame.frame_type == frame_type)
      .count()
  }

633
  // Size in bytes of all frames of given frame type
634 635 636 637 638 639 640
  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()
  }

641
  pub fn print_summary(&self) -> String {
642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657
    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)
    );
658 659 660 661
    format!("\
    Key Frames: {:>6}    avg size: {:>7} B\n\
    Inter:      {:>6}    avg size: {:>7} B\n\
    Intra Only: {:>6}    avg size: {:>7} B\n\
662 663
    Switch:     {:>6}    avg size: {:>7} B\
    {}",
664
      key, key_size / key,
665
      inter, inter_size.checked_div(inter).unwrap_or(0),
666
      ionly, ionly_size / key,
667 668
      switch, switch_size / key,
      if self.show_psnr {
669 670 671 672 673 674 675 676 677
        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;
678 679 680 681
        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() }
682 683 684 685 686 687 688 689 690
    )
  }
}

impl fmt::Display for ProgressInfo {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    if let Some(total_frames) = self.total_frames {
      write!(
        f,
691
        "encoded {}/{} frames, {:.3} fps, {:.2} Kb/s, est. size: {:.2} MB, est. time: {:.0} s",
692 693 694 695
        self.frames_encoded(),
        total_frames,
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64,
696 697
        self.estimated_size() as f64 / (1024 * 1024) as f64,
        self.estimated_time()
698 699 700 701
      )
    } else {
      write!(
        f,
702
        "encoded {} frames, {:.3} fps, {:.2} Kb/s",
703 704 705 706
        self.frames_encoded(),
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64
      )
707
    }
Kyle Siefring's avatar
Kyle Siefring committed
708
  }
709
}