rav1e.rs 6.65 KB
Newer Older
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.

Luca Barbato's avatar
Luca Barbato committed
10
11
#![deny(bare_trait_objects)]

12
13
#[macro_use]
extern crate scan_fmt;
Guillaume Martres's avatar
Guillaume Martres committed
14

15
mod common;
16
mod decoder;
17
mod muxer;
Raphaël Zumer's avatar
Raphaël Zumer committed
18
use crate::common::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
19
use rav1e::*;
20

21
22
use std::io;
use std::io::Write;
23
use std::io::Read;
24
use std::path::Path;
25
use std::sync::Arc;
Raphaël Zumer's avatar
Raphaël Zumer committed
26
27
use crate::decoder::Decoder;
use crate::decoder::VideoDetails;
28
use crate::muxer::*;
29
30
use std::fs::File;
use std::io::BufWriter;
31

32
33
34
35
struct Source<D: Decoder> {
 limit: usize,
 count: usize,
 input: D,
36
 #[cfg(all(unix, feature = "signal-hook"))]
Luca Barbato's avatar
Luca Barbato committed
37
 exit_requested: Arc<std::sync::atomic::AtomicBool>,
38
}
39

40
41
42
impl<D: Decoder> Source<D> {
  fn read_frame<T: Pixel>(&mut self, ctx: &mut Context<T>, video_info: VideoDetails) {
    if self.limit != 0 && self.count == self.limit {
43
      ctx.flush();
44
      return;
45
    }
Luca Barbato's avatar
Luca Barbato committed
46

47
    #[cfg(all(unix, feature = "signal-hook"))] {
Luca Barbato's avatar
Luca Barbato committed
48
49
50
51
52
53
      if self.exit_requested.load(std::sync::atomic::Ordering::SeqCst) {
        ctx.flush();
        return;
      }
    }

54
55
56
57
58
59
60
61
62
63
64
65
66
67
    match self.input.read_frame(&video_info) {
      Ok(frame) => {
        match video_info.bit_depth {
          8 | 10 | 12 => {}
          _ => panic!("unknown input bit depth!")
        }
        self.count += 1;
        let _ = ctx.send_frame(Some(Arc::new(frame)));
      }
      _ => {
        ctx.flush();
      }
    };
  }
68
69
70
71
}

// Encode and write a frame.
// Returns frame information in a `Result`.
72
fn process_frame<T: Pixel, D: Decoder>(
73
  ctx: &mut Context<T>, output_file: &mut dyn Muxer,
74
  source: &mut Source<D>,
75
  mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
Luca Barbato's avatar
Luca Barbato committed
76
) -> Option<Vec<FrameSummary>> {
77
  let y4m_details = source.input.get_video_details();
78
79
  let mut frame_summaries = Vec::new();
  let pkt_wrapped = ctx.receive_packet();
80
81
  match pkt_wrapped {
    Ok(pkt) => {
82
      output_file.write_frame(pkt.number as u64, pkt.data.as_ref());
83
84
85
86
87
      if let (Some(ref mut y4m_enc_uw), Some(ref rec)) = (y4m_enc.as_mut(), &pkt.rec) {
        write_y4m_frame(y4m_enc_uw, rec, y4m_details);
      }
      frame_summaries.push(pkt.into());
    }
Luca Barbato's avatar
Luca Barbato committed
88
    Err(EncoderStatus::NeedMoreData) => {
89
      source.read_frame(ctx, y4m_details);
90
    }
Luca Barbato's avatar
Luca Barbato committed
91
92
    Err(EncoderStatus::EnoughData) => {
      unreachable!();
93
    }
Luca Barbato's avatar
Luca Barbato committed
94
95
96
97
    Err(EncoderStatus::LimitReached) => {
      return None;
    }
    Err(EncoderStatus::Failure) => {
98
      panic!("Failed to encode video");
99
    }
100
    Err(EncoderStatus::Encoded) => {}
101
  }
Luca Barbato's avatar
Luca Barbato committed
102
  Some(frame_summaries)
103
}
Guillaume Martres's avatar
Guillaume Martres committed
104

105
fn write_stats_file<T: Pixel>(ctx: &Context<T>, filename: &Path) -> Result<(), io::Error> {
106
107
  let file = File::create(filename)?;
  let writer = BufWriter::new(file);
Luca Barbato's avatar
Luca Barbato committed
108
  serde_json::to_writer(writer, ctx.get_first_pass_data()).expect("Serialization should not fail");
109
110
111
  Ok(())
}

112
113
fn do_encode<T: Pixel, D: Decoder>(
  cfg: Config, verbose: bool, mut progress: ProgressInfo,
linkmauve's avatar
linkmauve committed
114
  output: &mut dyn Muxer,
115
  source: &mut Source<D>,
116
117
118
119
120
  mut y4m_enc: Option<y4m::Encoder<'_, Box<dyn Write>>>
) {
  let mut ctx: Context<T> = cfg.new_context();


Luca Barbato's avatar
Luca Barbato committed
121
  while let Some(frame_info) =
122
    process_frame(&mut ctx, &mut *output, source, y4m_enc.as_mut())
123
124
125
  {
    for frame in frame_info {
      progress.add_frame(frame);
Luca Barbato's avatar
Luca Barbato committed
126
127
      if verbose {
        eprintln!("{} - {}", frame, progress);
128
      } else {
Luca Barbato's avatar
Luca Barbato committed
129
        eprint!("\r{}                    ", progress);
130
131
132
133
134
135
136
137
138
139
      };
    }

    output.flush().unwrap();
  }

  if cfg.enc.pass == Some(1) {
    if let Err(e) =
      write_stats_file(&ctx, cfg.enc.stats_file.as_ref().unwrap())
    {
Luca Barbato's avatar
Luca Barbato committed
140
      eprintln!("\nError: Failed to write stats file! {}\n", e);
141
142
    }
  }
Luca Barbato's avatar
Luca Barbato committed
143
  eprint!("\n{}\n", progress.print_summary());
144
145
}

Guillaume Martres's avatar
Guillaume Martres committed
146
fn main() {
147
  let mut cli = parse_cli();
148
  let mut y4m_dec = y4m::decode(&mut cli.io.input).expect("input is not a y4m file");
149
  let video_info = y4m_dec.get_video_details();
150
  let y4m_enc = match cli.io.rec.as_mut() {
151
    Some(rec) => Some(
152
153
154
      y4m::encode(
        video_info.width,
        video_info.height,
Luca Barbato's avatar
Luca Barbato committed
155
        y4m::Ratio::new(video_info.time_base.den as usize, video_info.time_base.num as usize)
156
      ).with_colorspace(y4m_dec.get_colorspace())
157
158
159
        .write_header(rec)
        .unwrap()
    ),
Michael Bebenita's avatar
Michael Bebenita committed
160
161
    None => None
  };
162

163
164
165
166
167
168
  cli.enc.width = video_info.width;
  cli.enc.height = video_info.height;
  cli.enc.bit_depth = video_info.bit_depth;
  cli.enc.chroma_sampling = video_info.chroma_sampling;
  cli.enc.chroma_sample_position = video_info.chroma_sample_position;
  cli.enc.time_base = video_info.time_base;
Kyle Siefring's avatar
Kyle Siefring committed
169
  let cfg = Config {
Luca Barbato's avatar
Luca Barbato committed
170
171
    enc: cli.enc,
    threads: cli.threads,
Kyle Siefring's avatar
Kyle Siefring committed
172
  };
Luca Barbato's avatar
Luca Barbato committed
173

Luca Barbato's avatar
Luca Barbato committed
174
  eprintln!("{}x{} @ {}/{} fps",
175
176
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
177
178
    video_info.time_base.den,
    video_info.time_base.num
179
  );
180

181
  cli.io.output.write_header(
182
183
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
184
185
    video_info.time_base.den as usize,
    video_info.time_base.num as usize
Michael Bebenita's avatar
Michael Bebenita committed
186
  );
Guillaume Martres's avatar
Guillaume Martres committed
187

188
  let progress = ProgressInfo::new(
Luca Barbato's avatar
Luca Barbato committed
189
    Rational { num: video_info.time_base.den, den: video_info.time_base.num },
190
191
    if cli.limit == 0 { None } else { Some(cli.limit) },
      cfg.enc.show_psnr
192
  );
Luca Barbato's avatar
Luca Barbato committed
193

194
195
  for _ in 0..cli.skip {
    y4m_dec.read_frame().expect("Skipped more frames than in the input");
Vibhoothi's avatar
Vibhoothi committed
196
  }
197

198
  #[cfg(all(unix, feature = "signal-hook"))]
Luca Barbato's avatar
Luca Barbato committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
  let exit_requested = {
    use std::sync::atomic::*;
    let e  = Arc::new(AtomicBool::from(false));

    fn setup_signal(sig: i32, e: Arc<AtomicBool>) {
      unsafe {
        signal_hook::register(sig, move || {
          if e.load(Ordering::SeqCst) {
            std::process::exit(128 + sig);
          }
          e.store(true, Ordering::SeqCst);
          eprintln!("\rExit requested, flushing.\n");
        }).expect("Cannot register the signal hooks");
      }
    }

    setup_signal(signal_hook::SIGTERM, e.clone());
    setup_signal(signal_hook::SIGQUIT, e.clone());
    setup_signal(signal_hook::SIGINT, e.clone());

    e
  };

222
  #[cfg(all(unix, feature = "signal-hook"))]
Luca Barbato's avatar
Luca Barbato committed
223
224
225
226
227
228
  let mut source = Source {
    limit: cli.limit,
    input: y4m_dec,
    count: 0,
    exit_requested
  };
229
  #[cfg(not(all(unix, feature = "signal-hook")))]
230
231
  let mut source = Source { limit: cli.limit, input: y4m_dec, count: 0 };

232
  if video_info.bit_depth == 8 {
233
    do_encode::<u8, y4m::Decoder<'_, Box<dyn Read>>>(
234
      cfg, cli.verbose, progress, &mut *cli.io.output, &mut source, y4m_enc
235
236
    )
  } else {
237
    do_encode::<u16, y4m::Decoder<'_, Box<dyn Read>>>(
238
      cfg, cli.verbose, progress, &mut *cli.io.output, &mut source, y4m_enc
239
    )
240
  }
Guillaume Martres's avatar
Guillaume Martres committed
241
}