1 // Copyright (c) 2017-2021, The rav1e contributors. All rights reserved
2 //
3 // This source code is subject to the terms of the BSD 2 Clause License and
4 // the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
5 // was not distributed with this source code in the LICENSE file, you can
6 // obtain it at www.aomedia.org/license/software. If the Alliance for Open
7 // Media Patent License 1.0 was not distributed with this source code in the
8 // PATENTS file, you can obtain it at www.aomedia.org/license/patent.
9
10 #![deny(bare_trait_objects)]
11 #![allow(clippy::cast_lossless)]
12 #![allow(clippy::cast_ptr_alignment)]
13 #![allow(clippy::cognitive_complexity)]
14 #![allow(clippy::needless_range_loop)]
15 #![allow(clippy::too_many_arguments)]
16 #![allow(clippy::verbose_bit_mask)]
17 #![allow(clippy::unreadable_literal)]
18 #![allow(clippy::many_single_char_names)]
19 #![warn(clippy::expl_impl_clone_on_copy)]
20 #![warn(clippy::linkedlist)]
21 #![warn(clippy::map_flatten)]
22 #![warn(clippy::mem_forget)]
23 #![warn(clippy::mut_mut)]
24 #![warn(clippy::mutex_integer)]
25 #![warn(clippy::needless_borrow)]
26 #![warn(clippy::needless_continue)]
27 #![warn(clippy::path_buf_push_overwrite)]
28 #![warn(clippy::range_plus_one)]
29
30 #[macro_use]
31 extern crate log;
32
33 mod common;
34 mod decoder;
35 mod error;
36 #[cfg(feature = "serialize")]
37 mod kv;
38 mod muxer;
39 mod stats;
40
41 use crate::common::*;
42 use crate::error::*;
43 use crate::stats::*;
44 use rav1e::config::CpuFeatureLevel;
45 use rav1e::prelude::*;
46
47 use crate::decoder::{Decoder, FrameBuilder, VideoDetails};
48 use crate::muxer::*;
49 use std::fs::File;
50 use std::io::{Read, Seek, Write};
51 use std::process::exit;
52 use std::sync::Arc;
53
54 impl<T: Pixel> FrameBuilder<T> for Context<T> {
new_frame(&self) -> Frame<T>55 fn new_frame(&self) -> Frame<T> {
56 Context::new_frame(self)
57 }
58 }
59
60 struct Source<D: Decoder> {
61 limit: usize,
62 count: usize,
63 input: D,
64 #[cfg(all(unix, feature = "signal-hook"))]
65 exit_requested: Arc<std::sync::atomic::AtomicBool>,
66 }
67
68 impl<D: Decoder> Source<D> {
69 cfg_if::cfg_if! {
70 if #[cfg(all(unix, feature = "signal-hook"))] {
71 fn new(limit: usize, input: D) -> Self {
72 use signal_hook::{flag, consts};
73
74 // Make sure double CTRL+C and similar kills
75 let exit_requested = Arc::new(std::sync::atomic::AtomicBool::new(false));
76 for sig in consts::TERM_SIGNALS {
77 // When terminated by a second term signal, exit with exit code 1.
78 // This will do nothing the first time (because term_now is false).
79 flag::register_conditional_shutdown(*sig, 1, Arc::clone(&exit_requested)).unwrap();
80 // But this will "arm" the above for the second time, by setting it to true.
81 // The order of registering these is important, if you put this one first, it will
82 // first arm and then terminate ‒ all in the first round.
83 flag::register(*sig, Arc::clone(&exit_requested)).unwrap();
84 }
85
86 Self { limit, input, count: 0, exit_requested, }
87 }
88 } else {
89 fn new(limit: usize, input: D) -> Self {
90 Self { limit, input, count: 0, }
91 }
92 }
93 }
94
read_frame<T: Pixel>( &mut self, ctx: &mut Context<T>, video_info: VideoDetails, ) -> Result<(), CliError>95 fn read_frame<T: Pixel>(
96 &mut self, ctx: &mut Context<T>, video_info: VideoDetails,
97 ) -> Result<(), CliError> {
98 if self.limit != 0 && self.count == self.limit {
99 ctx.flush();
100 return Ok(());
101 }
102
103 #[cfg(all(unix, feature = "signal-hook"))]
104 {
105 if self.exit_requested.load(std::sync::atomic::Ordering::SeqCst) {
106 ctx.flush();
107 return Ok(());
108 }
109 }
110
111 match self.input.read_frame(ctx, &video_info) {
112 Ok(frame) => {
113 match video_info.bit_depth {
114 8 | 10 | 12 => {}
115 _ => return Err(CliError::new("Unsupported bit depth")),
116 }
117 self.count += 1;
118 let _ = ctx.send_frame(Some(Arc::new(frame)));
119 }
120 _ => {
121 ctx.flush();
122 }
123 };
124 Ok(())
125 }
126 }
127
128 // Encode and write a frame.
129 // Returns frame information in a `Result`.
process_frame<T: Pixel, D: Decoder>( ctx: &mut Context<T>, output_file: &mut dyn Muxer, source: &mut Source<D>, pass1file: Option<&mut File>, pass2file: Option<&mut File>, mut y4m_enc: Option<&mut y4m::Encoder<Box<dyn Write + Send>>>, metrics_cli: MetricsEnabled, ) -> Result<Option<Vec<FrameSummary>>, CliError>130 fn process_frame<T: Pixel, D: Decoder>(
131 ctx: &mut Context<T>, output_file: &mut dyn Muxer, source: &mut Source<D>,
132 pass1file: Option<&mut File>, pass2file: Option<&mut File>,
133 mut y4m_enc: Option<&mut y4m::Encoder<Box<dyn Write + Send>>>,
134 metrics_cli: MetricsEnabled,
135 ) -> Result<Option<Vec<FrameSummary>>, CliError> {
136 let y4m_details = source.input.get_video_details();
137 let mut frame_summaries = Vec::new();
138 let mut pass1file = pass1file;
139 let mut pass2file = pass2file;
140
141 // Submit first pass data to pass 2.
142 if let Some(passfile) = pass2file.as_mut() {
143 while ctx.rc_second_pass_data_required() > 0 {
144 let mut buflen = [0u8; 8];
145 passfile
146 .read_exact(&mut buflen)
147 .map_err(|e| e.context("Unable to read the two-pass data file."))?;
148 let mut data = vec![0u8; u64::from_be_bytes(buflen) as usize];
149
150 passfile
151 .read_exact(&mut data)
152 .map_err(|e| e.context("Unable to read the two-pass data file."))?;
153
154 ctx
155 .rc_send_pass_data(&data)
156 .map_err(|e| e.context("Corrupted first pass data"))?;
157 }
158 }
159
160 let pkt_wrapped = ctx.receive_packet();
161 let (ret, emit_pass_data) = match pkt_wrapped {
162 Ok(pkt) => {
163 output_file.write_frame(
164 pkt.input_frameno as u64,
165 pkt.data.as_ref(),
166 pkt.frame_type,
167 );
168 if let (Some(ref mut y4m_enc_uw), Some(ref rec)) =
169 (y4m_enc.as_mut(), &pkt.rec)
170 {
171 write_y4m_frame(y4m_enc_uw, rec, y4m_details);
172 }
173 frame_summaries.push(build_frame_summary(
174 pkt,
175 y4m_details.bit_depth,
176 y4m_details.chroma_sampling,
177 metrics_cli,
178 ));
179 (Ok(Some(frame_summaries)), true)
180 }
181 Err(EncoderStatus::NeedMoreData) => {
182 source.read_frame(ctx, y4m_details)?;
183 (Ok(Some(frame_summaries)), false)
184 }
185 Err(EncoderStatus::EnoughData) => {
186 unreachable!();
187 }
188 Err(EncoderStatus::LimitReached) => (Ok(None), true),
189 Err(e @ EncoderStatus::Failure) => {
190 (Err(e.context("Failed to encode video")), false)
191 }
192 Err(e @ EncoderStatus::NotReady) => {
193 (Err(e.context("Mismanaged handling of two-pass stats data")), false)
194 }
195 Err(EncoderStatus::Encoded) => (Ok(Some(frame_summaries)), true),
196 };
197
198 if ret.is_err() {
199 return ret;
200 }
201
202 // Save first pass data from pass 1.
203 if let Some(passfile) = pass1file.as_mut() {
204 if emit_pass_data {
205 match ctx.rc_receive_pass_data() {
206 Some(RcData::Frame(outbuf)) => {
207 let len = outbuf.len() as u64;
208 passfile.write_all(&len.to_be_bytes()).map_err(|e| {
209 e.context("Unable to write to two-pass data file.")
210 })?;
211
212 passfile.write_all(&outbuf).map_err(|e| {
213 e.context("Unable to write to two-pass data file.")
214 })?;
215 }
216 Some(RcData::Summary(outbuf)) => {
217 // The last packet of rate control data we get is the summary data.
218 // Let's put it at the start of the file.
219 passfile.seek(std::io::SeekFrom::Start(0)).map_err(|e| {
220 e.context("Unable to seek in the two-pass data file.")
221 })?;
222 let len = outbuf.len() as u64;
223
224 passfile.write_all(&len.to_be_bytes()).map_err(|e| {
225 e.context("Unable to write to two-pass data file.")
226 })?;
227
228 passfile.write_all(&outbuf).map_err(|e| {
229 e.context("Unable to write to two-pass data file.")
230 })?;
231 }
232 None => {}
233 }
234 }
235 }
236
237 ret
238 }
239
do_encode<T: Pixel, D: Decoder>( cfg: Config, verbose: Verbose, mut progress: ProgressInfo, output: &mut dyn Muxer, mut source: Source<D>, mut pass1file: Option<File>, mut pass2file: Option<File>, mut y4m_enc: Option<y4m::Encoder<Box<dyn Write + Send>>>, metrics_enabled: MetricsEnabled, ) -> Result<(), CliError>240 fn do_encode<T: Pixel, D: Decoder>(
241 cfg: Config, verbose: Verbose, mut progress: ProgressInfo,
242 output: &mut dyn Muxer, mut source: Source<D>, mut pass1file: Option<File>,
243 mut pass2file: Option<File>,
244 mut y4m_enc: Option<y4m::Encoder<Box<dyn Write + Send>>>,
245 metrics_enabled: MetricsEnabled,
246 ) -> Result<(), CliError> {
247 let mut ctx: Context<T> =
248 cfg.new_context().map_err(|e| e.context("Invalid encoder settings"))?;
249
250 // Let's write down a placeholder.
251 if let Some(passfile) = pass1file.as_mut() {
252 let len = ctx.rc_summary_size();
253 let buf = vec![0u8; len];
254
255 passfile
256 .write_all(&(len as u64).to_be_bytes())
257 .map_err(|e| e.context("Unable to write to two-pass data file."))?;
258
259 passfile
260 .write_all(&buf)
261 .map_err(|e| e.context("Unable to write to two-pass data file."))?;
262 }
263
264 while let Some(frame_info) = process_frame(
265 &mut ctx,
266 &mut *output,
267 &mut source,
268 pass1file.as_mut(),
269 pass2file.as_mut(),
270 y4m_enc.as_mut(),
271 metrics_enabled,
272 )? {
273 if verbose != Verbose::Quiet {
274 for frame in frame_info {
275 progress.add_frame(frame.clone());
276 if verbose == Verbose::Verbose {
277 info!("{} - {}", frame, progress);
278 } else {
279 // Print a one-line progress indicator that overrides itself with every update
280 eprint!("\r{} ", progress);
281 };
282 }
283
284 output.flush().unwrap();
285 }
286 }
287 if verbose != Verbose::Quiet {
288 if verbose == Verbose::Verbose {
289 // Clear out the temporary progress indicator
290 eprint!("\r");
291 }
292 progress.print_summary(verbose == Verbose::Verbose);
293 }
294 Ok(())
295 }
296
main()297 fn main() {
298 #[cfg(feature = "tracing")]
299 use rust_hawktracer::*;
300 init_logger();
301
302 #[cfg(feature = "tracing")]
303 let instance = HawktracerInstance::new();
304 #[cfg(feature = "tracing")]
305 let _listener = instance.create_listener(HawktracerListenerType::ToFile {
306 file_path: "trace.bin".into(),
307 buffer_size: 4096,
308 });
309
310 run().unwrap_or_else(|e| {
311 error::print_error(&e);
312 exit(1);
313 });
314 }
315
init_logger()316 fn init_logger() {
317 use std::str::FromStr;
318 fn level_colored(l: log::Level) -> console::StyledObject<&'static str> {
319 use console::style;
320 use log::Level;
321 match l {
322 Level::Trace => style("??").dim(),
323 Level::Debug => style("? ").dim(),
324 Level::Info => style("> ").green(),
325 Level::Warn => style("! ").yellow(),
326 Level::Error => style("!!").red(),
327 }
328 }
329
330 // this can be changed to flatten
331 let level = std::env::var("RAV1E_LOG")
332 .ok()
333 .map(|l| log::LevelFilter::from_str(&l).ok())
334 .unwrap_or(Some(log::LevelFilter::Info))
335 .unwrap();
336
337 fern::Dispatch::new()
338 .format(move |out, message, record| {
339 out.finish(format_args!(
340 "{level} {message}",
341 level = level_colored(record.level()),
342 message = message,
343 ));
344 })
345 // set the default log level. to filter out verbose log messages from dependencies, set
346 // this to Warn and overwrite the log level for your crate.
347 .level(log::LevelFilter::Warn)
348 // change log levels for individual modules. Note: This looks for the record's target
349 // field which defaults to the module path but can be overwritten with the `target`
350 // parameter:
351 // `info!(target="special_target", "This log message is about special_target");`
352 .level_for("rav1e", level)
353 // output to stdout
354 .chain(std::io::stderr())
355 .apply()
356 .unwrap();
357 }
358
359 cfg_if::cfg_if! {
360 if #[cfg(any(target_os = "windows", target_arch = "wasm32"))] {
361 fn print_rusage() {
362 eprintln!("Windows benchmarking is not supported currently.");
363 }
364 } else {
365 fn print_rusage() {
366 let (utime, stime, maxrss) = unsafe {
367 let mut usage = std::mem::zeroed();
368 let _ = libc::getrusage(libc::RUSAGE_SELF, &mut usage);
369 (usage.ru_utime, usage.ru_stime, usage.ru_maxrss)
370 };
371 eprintln!(
372 "user time: {} s",
373 utime.tv_sec as f64 + utime.tv_usec as f64 / 1_000_000f64
374 );
375 eprintln!(
376 "system time: {} s",
377 stime.tv_sec as f64 + stime.tv_usec as f64 / 1_000_000f64
378 );
379 eprintln!("maximum rss: {} KB", maxrss);
380 }
381 }
382 }
383
run() -> Result<(), error::CliError>384 fn run() -> Result<(), error::CliError> {
385 let mut cli = parse_cli()?;
386 // Maximum frame size by specification + maximum y4m header
387 let limit = y4m::Limits {
388 // Use saturating operations to gracefully handle 32-bit architectures
389 bytes: 64usize
390 .saturating_mul(64)
391 .saturating_mul(4096)
392 .saturating_mul(2304)
393 .saturating_add(1024),
394 };
395 let mut y4m_dec = match y4m::Decoder::new_with_limits(cli.io.input, limit) {
396 Err(e) => {
397 return Err(CliError::new(match e {
398 y4m::Error::ParseError(_) => {
399 "Could not input video. Is it a y4m file?"
400 }
401 y4m::Error::UnknownColorspace => {
402 "Unknown colorspace or unsupported bit depth."
403 }
404 y4m::Error::OutOfMemory => "The video's frame size exceeds the limit.",
405 _ => unreachable!(),
406 }))
407 }
408 Ok(d) => d,
409 };
410 let video_info = y4m_dec.get_video_details();
411 let y4m_enc = cli.io.rec.map(|rec| {
412 y4m::encode(
413 video_info.width,
414 video_info.height,
415 y4m::Ratio::new(
416 video_info.time_base.den as usize,
417 video_info.time_base.num as usize,
418 ),
419 )
420 .with_colorspace(y4m_dec.get_colorspace())
421 .with_pixel_aspect(y4m::Ratio {
422 num: video_info.sample_aspect_ratio.num as usize,
423 den: video_info.sample_aspect_ratio.den as usize,
424 })
425 .write_header(rec)
426 .unwrap()
427 });
428
429 cli.enc.width = video_info.width;
430 cli.enc.height = video_info.height;
431 cli.enc.sample_aspect_ratio = video_info.sample_aspect_ratio;
432 cli.enc.bit_depth = video_info.bit_depth;
433 cli.enc.chroma_sampling = video_info.chroma_sampling;
434 cli.enc.chroma_sample_position = video_info.chroma_sample_position;
435
436 // If no pixel range is specified via CLI, assume limited,
437 // as it is the default for the Y4M format.
438 if !cli.color_range_specified {
439 cli.enc.pixel_range = PixelRange::Limited;
440 }
441
442 if !cli.override_time_base {
443 cli.enc.time_base = video_info.time_base;
444 }
445
446 let mut rc = RateControlConfig::new();
447
448 let pass2file = match cli.pass2file_name {
449 Some(f) => {
450 let mut f = File::open(f).map_err(|e| {
451 e.context("Unable to open file for reading two-pass data")
452 })?;
453 let mut buflen = [0u8; 8];
454 f.read_exact(&mut buflen)
455 .map_err(|e| e.context("Summary data too short"))?;
456 let len = i64::from_be_bytes(buflen);
457 let mut buf = vec![0u8; len as usize];
458
459 f.read_exact(&mut buf)
460 .map_err(|e| e.context("Summary data too short"))?;
461
462 rc = RateControlConfig::from_summary_slice(&buf)
463 .map_err(|e| e.context("Invalid summary"))?;
464
465 Some(f)
466 }
467 None => None,
468 };
469
470 let pass1file = match cli.pass1file_name {
471 Some(f) => {
472 let f = File::create(f).map_err(|e| {
473 e.context("Unable to open file for writing two-pass data")
474 })?;
475 rc = rc.with_emit_data(true);
476 Some(f)
477 }
478 None => None,
479 };
480
481 let cfg = Config::new()
482 .with_encoder_config(cli.enc)
483 .with_threads(cli.threads)
484 .with_rate_control(rc);
485
486 #[cfg(feature = "serialize")]
487 {
488 if let Some(save_config) = cli.save_config {
489 let mut out = File::create(save_config)
490 .map_err(|e| e.context("Cannot create configuration file"))?;
491 let s = toml::to_string(&cli.enc).unwrap();
492 out
493 .write_all(s.as_bytes())
494 .map_err(|e| e.context("Cannot write the configuration file"))?
495 }
496 }
497
498 cli.io.output.write_header(
499 video_info.width,
500 video_info.height,
501 cli.enc.time_base.den as usize,
502 cli.enc.time_base.num as usize,
503 );
504
505 let tiling =
506 cfg.tiling_info().map_err(|e| e.context("Invalid configuration"))?;
507 if cli.verbose != Verbose::Quiet {
508 info!("CPU Feature Level: {}", CpuFeatureLevel::default());
509
510 info!(
511 "Using y4m decoder: {}x{}p @ {}/{} fps, {}, {}-bit",
512 video_info.width,
513 video_info.height,
514 video_info.time_base.den,
515 video_info.time_base.num,
516 video_info.chroma_sampling,
517 video_info.bit_depth
518 );
519 info!("Encoding settings: {}", cli.enc);
520
521 if tiling.tile_count() == 1 {
522 info!("Using 1 tile");
523 } else {
524 info!(
525 "Using {} tiles ({}x{})",
526 tiling.tile_count(),
527 tiling.cols,
528 tiling.rows
529 );
530 }
531 }
532
533 let progress = ProgressInfo::new(
534 Rational { num: video_info.time_base.den, den: video_info.time_base.num },
535 if cli.limit == 0 { None } else { Some(cli.limit) },
536 cli.metrics_enabled,
537 );
538
539 for _ in 0..cli.skip {
540 match y4m_dec.read_frame() {
541 Ok(f) => f,
542 Err(_) => {
543 return Err(CliError::new("Skipped more frames than in the input"))
544 }
545 };
546 }
547
548 let source = Source::new(cli.limit, y4m_dec);
549
550 if video_info.bit_depth == 8 {
551 do_encode::<u8, y4m::Decoder<Box<dyn Read + Send>>>(
552 cfg,
553 cli.verbose,
554 progress,
555 &mut *cli.io.output,
556 source,
557 pass1file,
558 pass2file,
559 y4m_enc,
560 cli.metrics_enabled,
561 )?
562 } else {
563 do_encode::<u16, y4m::Decoder<Box<dyn Read + Send>>>(
564 cfg,
565 cli.verbose,
566 progress,
567 &mut *cli.io.output,
568 source,
569 pass1file,
570 pass2file,
571 y4m_enc,
572 cli.metrics_enabled,
573 )?
574 }
575 if cli.benchmark {
576 print_rusage();
577 }
578
579 Ok(())
580 }
581