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

use std::{fmt, io, slice};
15
use std::fs::File;
16
use std::io::prelude::*;
Kyle Siefring's avatar
Kyle Siefring committed
17
use std::sync::Arc;
18
use std::time::Instant;
19
use y4m;
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
31
32
33
34
pub struct CliOptions {
  pub io: EncoderIO,
  pub enc: EncoderConfig,
  pub limit: usize,
  pub verbose: bool,
}

pub fn parse_cli() -> CliOptions {
Luca Barbato's avatar
Luca Barbato committed
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
  let matches = App::new("rav1e")
    .version("0.1.0")
    .about("AV1 video encoder")
    .arg(
      Arg::with_name("INPUT")
        .help("Uncompressed YUV4MPEG2 video input")
        .required(true)
        .index(1)
    ).arg(
      Arg::with_name("OUTPUT")
        .help("Compressed AV1 in IVF video output")
        .short("o")
        .long("output")
        .required(true)
        .takes_value(true)
50
51
52
53
54
    ).arg(
      Arg::with_name("RECONSTRUCTION")
        .short("r")
        .takes_value(true)
    ).arg(
Luca Barbato's avatar
Luca Barbato committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
      Arg::with_name("LIMIT")
        .help("Maximum number of frames to encode")
        .short("l")
        .long("limit")
        .takes_value(true)
        .default_value("0")
    ).arg(
      Arg::with_name("QP")
        .help("Quantizer (0-255)")
        .long("quantizer")
        .takes_value(true)
        .default_value("100")
    ).arg(
      Arg::with_name("SPEED")
        .help("Speed level (0(slow)-10(fast))")
        .short("s")
        .long("speed")
        .takes_value(true)
        .default_value("3")
Josh Holmer's avatar
Josh Holmer committed
74
75
76
77
78
79
80
    ).arg(
      Arg::with_name("MIN_KEYFRAME_INTERVAL")
        .help("Minimum interval between keyframes")
        .short("i")
        .long("min-keyint")
        .takes_value(true)
        .default_value("12")
81
82
    ).arg(
      Arg::with_name("KEYFRAME_INTERVAL")
Josh Holmer's avatar
Josh Holmer committed
83
        .help("Maximum interval between keyframes")
84
85
86
        .short("I")
        .long("keyint")
        .takes_value(true)
87
        .default_value("240")
88
89
90
91
92
    ).arg(
      Arg::with_name("LOW_LATENCY")
        .help("low latency mode. true or false")
        .long("low_latency")
        .takes_value(true)
Frank Bossen's avatar
Frank Bossen committed
93
        .default_value("true")
Luca Barbato's avatar
Luca Barbato committed
94
95
96
97
98
99
100
    ).arg(
      Arg::with_name("TUNE")
        .help("Quality tuning (Will enforce partition sizes >= 8x8)")
        .long("tune")
        .possible_values(&Tune::variants())
        .default_value("psnr")
        .case_insensitive(true)
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
    ).arg(
      Arg::with_name("COLOR_PRIMARIES")
      .help("Color primaries used to describe color parameters.")
      .long("primaries")
      .possible_values(&ColorPrimaries::variants())
      .default_value("unspecified")
      .case_insensitive(true)
    ).arg(
      Arg::with_name("TRANSFER_CHARACTERISTICS")
      .help("Transfer characteristics used to describe color parameters.")
      .long("transfer")
      .possible_values(&TransferCharacteristics::variants())
      .default_value("unspecified")
      .case_insensitive(true)
    ).arg(
      Arg::with_name("MATRIX_COEFFICIENTS")
      .help("Color primaries used to describe color parameters.")
      .long("matrix")
      .possible_values(&MatrixCoefficients::variants())
      .default_value("unspecified")
      .case_insensitive(true)
122
123
124
125
126
    ).arg(
      Arg::with_name("VERBOSE")
        .help("verbose logging, output info for every frame")
        .long("verbose")
        .short("v")
127
128
129
130
    ).arg(
      Arg::with_name("PSNR")
        .help("calculate and display PSNR metrics")
        .long("psnr")
Luca Barbato's avatar
Luca Barbato committed
131
    ).get_matches();
132

Luca Barbato's avatar
Luca Barbato committed
133
134
135
136
137
138
139
140
141
142
143
144
145
  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
146

147
148
  CliOptions {
    io,
149
    enc: parse_config(&matches),
150
151
152
153
154
    limit: matches.value_of("LIMIT").unwrap().parse().unwrap(),
    verbose: matches.is_present("VERBOSE"),
  }
}

155
156
157
fn parse_config(matches: &ArgMatches) -> EncoderConfig {
  let speed = matches.value_of("SPEED").unwrap().parse().unwrap();
  let quantizer = matches.value_of("QP").unwrap().parse().unwrap();
Josh Holmer's avatar
Josh Holmer committed
158
159
  let min_interval = matches.value_of("MIN_KEYFRAME_INTERVAL").unwrap().parse().unwrap();
  let max_interval = matches.value_of("KEYFRAME_INTERVAL").unwrap().parse().unwrap();
160
161
162
163
164
165

  // Validate arguments
  if quantizer == 0 {
    unimplemented!("Lossless encoding not yet implemented");
  } else if quantizer > 255 || speed > 10 {
    panic!("argument out of range");
Josh Holmer's avatar
Josh Holmer committed
166
167
  } else if min_interval > max_interval {
    panic!("Maximum keyframe interval must be greater than or equal to minimum keyframe interval");
168
169
  }

170
171
172
173
  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();

174
  let mut cfg = EncoderConfig::with_speed_preset(speed);
Josh Holmer's avatar
Josh Holmer committed
175
176
  cfg.max_key_frame_interval = min_interval;
  cfg.max_key_frame_interval = max_interval;
177
178
  cfg.low_latency = matches.value_of("LOW_LATENCY").unwrap().parse().unwrap();
  cfg.tune = matches.value_of("TUNE").unwrap().parse().unwrap();
179
180
181
182
183
184
185
186
187
188
189
190
  cfg.color_description = if color_primaries == ColorPrimaries::Unspecified &&
    transfer_characteristics == TransferCharacteristics::Unspecified &&
    matrix_coefficients == MatrixCoefficients::Unspecified {
      // No need to set a color description with all parameters unspecified.
      None
    } else {
      Some(ColorDescription {
        color_primaries,
        transfer_characteristics,
        matrix_coefficients
      })
    };
191
  cfg.quantizer = quantizer;
192
  cfg.show_psnr = matches.is_present("PSNR");
193

194
195
196
  cfg
}

197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
pub fn map_y4m_color_space(color_space: y4m::Colorspace) -> (ChromaSampling, ChromaSamplePosition) {
  match color_space {
    y4m::Colorspace::C420jpeg
    | y4m::Colorspace::C420paldv => (ChromaSampling::Cs420, ChromaSamplePosition::Unknown),
    y4m::Colorspace::C420mpeg2 => (ChromaSampling::Cs420, ChromaSamplePosition::Vertical),
    y4m::Colorspace::C420
    | y4m::Colorspace::C420p10
    | y4m::Colorspace::C420p12 => (ChromaSampling::Cs420, ChromaSamplePosition::Colocated),
    y4m::Colorspace::C422
    | y4m::Colorspace::C422p10
    | y4m::Colorspace::C422p12 => (ChromaSampling::Cs422, ChromaSamplePosition::Colocated),
    y4m::Colorspace::C444
    | y4m::Colorspace::C444p10
    | y4m::Colorspace::C444p12 => (ChromaSampling::Cs444, ChromaSamplePosition::Colocated),
    _ => panic!("Chroma characteristics unknown for the specified color space.")
  }
}

215
216
#[derive(Debug, Clone, Copy)]
pub struct FrameSummary {
217
  // Frame size in bytes
218
219
  pub size: usize,
  pub number: u64,
220
  pub frame_type: FrameType,
221
  // PSNR for Y, U, and V planes
222
  pub psnr: Option<(f64, f64, f64)>,
223
224
225
226
227
228
229
230
}

impl From<Packet> for FrameSummary {
  fn from(packet: Packet) -> Self {
    Self {
      size: packet.data.len(),
      number: packet.number,
      frame_type: packet.frame_type,
231
      psnr: packet.psnr,
232
233
234
235
236
237
238
239
    }
  }
}

impl fmt::Display for FrameSummary {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(
      f,
240
      "Frame {} - {} - {} bytes{}",
241
242
      self.number,
      self.frame_type,
243
244
245
246
      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() }
247
248
    )
  }
249
}
250

251
252
// Encode and write a frame.
// Returns frame information in a `Result`.
Kyle Siefring's avatar
Kyle Siefring committed
253
254
255
pub fn process_frame(
  ctx: &mut Context, output_file: &mut dyn Write,
  y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
256
  mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
257
) -> Result<Vec<FrameSummary>, ()> {
Kyle Siefring's avatar
Kyle Siefring committed
258
259
260
261
262
263
  let width = y4m_dec.get_width();
  let height = y4m_dec.get_height();
  let y4m_bits = y4m_dec.get_bit_depth();
  let y4m_bytes = y4m_dec.get_bytes_per_sample();
  let bit_depth = y4m_dec.get_colorspace().get_bit_depth();

264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
  if ctx.needs_more_frames(ctx.get_frame_count()) {
    match y4m_dec.read_frame() {
      Ok(y4m_frame) => {
        let y4m_y = y4m_frame.get_y_plane();
        let y4m_u = y4m_frame.get_u_plane();
        let y4m_v = y4m_frame.get_v_plane();
        let mut input = ctx.new_frame();
        {
          let input = Arc::get_mut(&mut input).unwrap();
          input.planes[0].copy_from_raw_u8(&y4m_y, width * y4m_bytes, y4m_bytes);
          input.planes[1].copy_from_raw_u8(
            &y4m_u,
            width * y4m_bytes / 2,
            y4m_bytes
          );
          input.planes[2].copy_from_raw_u8(
            &y4m_v,
            width * y4m_bytes / 2,
            y4m_bytes
          );
        }
Kyle Siefring's avatar
Kyle Siefring committed
285

286
287
288
289
        match y4m_bits {
          8 | 10 | 12 => {}
          _ => panic!("unknown input bit depth!")
        }
Kyle Siefring's avatar
Kyle Siefring committed
290

291
292
293
294
295
296
297
        let _ = ctx.send_frame(input);
      }
      _ => {
        let frames_to_be_coded = ctx.get_frame_count();
        ctx.set_frames_to_be_coded(frames_to_be_coded);
        ctx.flush();
      }
fbossen's avatar
fbossen committed
298
    }
299
300
  } else {
    ctx.flush();
fbossen's avatar
fbossen committed
301
302
  };

303
  let mut frame_summaries = Vec::new();
Josh Holmer's avatar
Josh Holmer committed
304
  loop {
fbossen's avatar
fbossen committed
305
306
    let pkt_wrapped = ctx.receive_packet();
    match pkt_wrapped {
307
308
309
310
311
312
      Ok(pkt) => {
        write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
        if let Some(y4m_enc_uw) = y4m_enc.as_mut() {
          if let Some(ref rec) = pkt.rec {
            let pitch_y = if bit_depth > 8 { width * 2 } else { width };
            let pitch_uv = pitch_y / 2;
Kyle Siefring's avatar
Kyle Siefring committed
313

314
315
316
317
318
            let (mut rec_y, mut rec_u, mut rec_v) = (
              vec![128u8; pitch_y * height],
              vec![128u8; pitch_uv * (height / 2)],
              vec![128u8; pitch_uv * (height / 2)]
            );
Kyle Siefring's avatar
Kyle Siefring committed
319

320
321
322
323
            let (stride_y, stride_u, stride_v) = (
              rec.planes[0].cfg.stride,
              rec.planes[1].cfg.stride,
              rec.planes[2].cfg.stride
Kyle Siefring's avatar
Kyle Siefring committed
324
            );
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342

            for (line, line_out) in rec.planes[0]
              .data_origin()
              .chunks(stride_y)
              .zip(rec_y.chunks_mut(pitch_y))
            {
              if bit_depth > 8 {
                unsafe {
                  line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                    line.as_ptr() as (*const u8),
                    pitch_y
                  ));
                }
              } else {
                line_out.copy_from_slice(
                  &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]
                );
              }
343
            }
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
            for (line, line_out) in rec.planes[1]
              .data_origin()
              .chunks(stride_u)
              .zip(rec_u.chunks_mut(pitch_uv))
            {
              if bit_depth > 8 {
                unsafe {
                  line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                    line.as_ptr() as (*const u8),
                    pitch_uv
                  ));
                }
              } else {
                line_out.copy_from_slice(
                  &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
                );
              }
361
            }
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
            for (line, line_out) in rec.planes[2]
              .data_origin()
              .chunks(stride_v)
              .zip(rec_v.chunks_mut(pitch_uv))
            {
              if bit_depth > 8 {
                unsafe {
                  line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                    line.as_ptr() as (*const u8),
                    pitch_uv
                  ));
                }
              } else {
                line_out.copy_from_slice(
                  &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
                );
              }
            }

            let rec_frame = y4m::Frame::new([&rec_y, &rec_u, &rec_v], None);
            y4m_enc_uw.write_frame(&rec_frame).unwrap();
Kyle Siefring's avatar
Kyle Siefring committed
383
384
          }
        }
385
        frame_summaries.push(pkt.into());
386
      },
Josh Holmer's avatar
Josh Holmer committed
387
      _ => { break; }
388
389
    }
  }
390
  Ok(frame_summaries)
391
}
Kyle Siefring's avatar
Kyle Siefring committed
392

393
394
#[derive(Debug, Clone)]
pub struct ProgressInfo {
395
  // Frame rate of the video
396
  frame_rate: y4m::Ratio,
397
  // The length of the whole video, in frames, if known
398
  total_frames: Option<usize>,
399
  // The time the encode was started
400
  time_started: Instant,
401
  // List of frames encoded so far
402
  frame_info: Vec<FrameSummary>,
403
404
405
406
  // 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.
407
  encoded_size: usize,
408
  // Whether to display PSNR statistics during and at end of encode
409
  show_psnr: bool,
410
411
412
}

impl ProgressInfo {
413
  pub fn new(frame_rate: y4m::Ratio, total_frames: Option<usize>, show_psnr: bool) -> Self {
414
415
416
417
418
419
    Self {
      frame_rate,
      total_frames,
      time_started: Instant::now(),
      frame_info: Vec::with_capacity(total_frames.unwrap_or_default()),
      encoded_size: 0,
420
      show_psnr,
421
422
423
424
425
426
427
428
429
430
431
432
433
    }
  }

  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 {
434
435
    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)
436
437
438
439
440
441
  }

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

442
  // Returns the bitrate of the frames so far, in bits/second
443
444
445
446
447
448
  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
  }

449
  // Estimates the final filesize in bytes, if the number of frames is known
450
451
452
453
454
455
  pub fn estimated_size(&self) -> usize {
    self.total_frames
      .map(|frames| self.encoded_size * frames / self.frames_encoded())
      .unwrap_or_default()
  }

456
  // Number of frames of given type which appear in the video
457
458
459
460
461
462
  pub fn get_frame_type_count(&self, frame_type: FrameType) -> usize {
    self.frame_info.iter()
      .filter(|frame| frame.frame_type == frame_type)
      .count()
  }

463
  // Size in bytes of all frames of given frame type
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
  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()
  }

  pub fn print_stats(&self) -> String {
    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));
    format!("\
    Key Frames: {:>6}    avg size: {:>7} B\n\
    Inter:      {:>6}    avg size: {:>7} B\n\
    Intra Only: {:>6}    avg size: {:>7} B\n\
480
481
    Switch:     {:>6}    avg size: {:>7} B\
    {}",
482
      key, key_size / key,
483
      inter, inter_size.checked_div(inter).unwrap_or(0),
484
      ionly, ionly_size / key,
485
486
487
488
489
490
491
492
493
      switch, switch_size / key,
      if self.show_psnr {
        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;
        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() }
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
    )
  }
}

impl fmt::Display for ProgressInfo {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    if let Some(total_frames) = self.total_frames {
      write!(
        f,
        "encoded {}/{} frames, {:.2} fps, {:.2} Kb/s, est. size: {:.2} MB",
        self.frames_encoded(),
        total_frames,
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64,
        self.estimated_size() as f64 / (1024 * 1024) as f64
      )
    } else {
      write!(
        f,
        "encoded {} frames, {:.2} fps, {:.2} Kb/s",
        self.frames_encoded(),
        self.encoding_fps(),
        self.bitrate() as f64 / 1024f64
      )
518
    }
Kyle Siefring's avatar
Kyle Siefring committed
519
  }
520
}