rav1e.rs 9.38 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
mod common;
13
mod decoder;
14
mod muxer;
Raphaël Zumer's avatar
Raphaël Zumer committed
15
use crate::common::*;
Luca Barbato's avatar
Luca Barbato committed
16
use rav1e::prelude::*;
17

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

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

38
39
40
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 {
41
      ctx.flush();
42
      return;
43
    }
Luca Barbato's avatar
Luca Barbato committed
44

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

52
53
54
55
56
57
58
59
60
61
62
63
64
65
    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();
      }
    };
  }
66
67
68
69
}

// Encode and write a frame.
// Returns frame information in a `Result`.
70
fn process_frame<T: Pixel, D: Decoder>(
71
  ctx: &mut Context<T>, output_file: &mut dyn Muxer,
72
  source: &mut Source<D>,
73
74
75
76
  pass1file: Option<&mut File>,
  pass2file: Option<&mut File>,
  buffer: &mut [u8],
  buf_pos: &mut usize,
77
  mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
Luca Barbato's avatar
Luca Barbato committed
78
) -> Option<Vec<FrameSummary>> {
79
  let y4m_details = source.input.get_video_details();
80
  let mut frame_summaries = Vec::new();
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
  let mut pass1file = pass1file;
  let mut pass2file = pass2file;
  // Submit first pass data to pass 2.
  if let Some(passfile) = pass2file.as_mut() {
    loop {
      let mut bytes = ctx.twopass_bytes_needed();
      if bytes == 0 {
        break;
      }
      // Read in some more bytes, if necessary.
      bytes = bytes.min(buffer.len() - *buf_pos);
      if bytes > 0 {
        passfile.read_exact(&mut buffer[*buf_pos..*buf_pos + bytes])
         .expect("Could not read frame data from two-pass data file!");
      }
      // And pass them off.
      let consumed = ctx.twopass_in(&buffer[*buf_pos..*buf_pos + bytes])
       .expect("Error submitting pass data in second pass.");
      // If the encoder consumed the whole buffer, reset it.
      if consumed >= bytes {
        *buf_pos = 0;
      }
      else {
        *buf_pos += consumed;
      }
    }
  }
  // Extract first pass data from pass 1.
  // We call this before encoding any frames to ensure we are in 2-pass mode
  //  and to get the placeholder header data.
  if let Some(passfile) = pass1file.as_mut() {
    if let Some(outbuf) = ctx.twopass_out() {
      passfile.write_all(outbuf)
       .expect("Unable to write to two-pass data file.");
    }
  }

118
  let pkt_wrapped = ctx.receive_packet();
119
120
  match pkt_wrapped {
    Ok(pkt) => {
121
      output_file.write_frame(pkt.input_frameno as u64, pkt.data.as_ref());
122
123
124
125
126
      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
127
    Err(EncoderStatus::NeedMoreData) => {
128
      source.read_frame(ctx, y4m_details);
129
    }
Luca Barbato's avatar
Luca Barbato committed
130
131
    Err(EncoderStatus::EnoughData) => {
      unreachable!();
132
    }
Luca Barbato's avatar
Luca Barbato committed
133
    Err(EncoderStatus::LimitReached) => {
134
135
136
137
138
139
140
141
142
143
144
      if let Some(passfile) = pass1file.as_mut() {
        if let Some(outbuf) = ctx.twopass_out() {
          // The last block of data we get is the summary data that needs to go
          //  at the start of the pass file.
          // Seek to the start so we can write it there.
          passfile.seek(std::io::SeekFrom::Start(0))
           .expect("Unable to seek in two-pass data file.");
          passfile.write_all(outbuf)
           .expect("Unable to write to two-pass data file.");
        }
      }
Luca Barbato's avatar
Luca Barbato committed
145
146
147
      return None;
    }
    Err(EncoderStatus::Failure) => {
148
      panic!("Failed to encode video");
149
    }
150
151
152
    Err(EncoderStatus::NotReady) => {
      panic!("Mis-managed handling of two-pass stats data");
    }
153
    Err(EncoderStatus::Encoded) => {}
154
  }
Luca Barbato's avatar
Luca Barbato committed
155
  Some(frame_summaries)
156
}
Guillaume Martres's avatar
Guillaume Martres committed
157

158
fn write_stats_file<T: Pixel>(ctx: &Context<T>, filename: &Path) -> Result<(), io::Error> {
159
160
  let file = File::create(filename)?;
  let writer = BufWriter::new(file);
Luca Barbato's avatar
Luca Barbato committed
161
  serde_json::to_writer(writer, ctx.get_first_pass_data()).expect("Serialization should not fail");
162
163
164
  Ok(())
}

165
166
fn do_encode<T: Pixel, D: Decoder>(
  cfg: Config, verbose: bool, mut progress: ProgressInfo,
linkmauve's avatar
linkmauve committed
167
  output: &mut dyn Muxer,
168
  source: &mut Source<D>,
169
170
  pass1file_name: Option<&String>,
  pass2file_name: Option<&String>,
171
172
173
174
  mut y4m_enc: Option<y4m::Encoder<'_, Box<dyn Write>>>
) {
  let mut ctx: Context<T> = cfg.new_context();

175
176
177
178
179
180
181
182
183
184
185
  let mut pass2file = pass2file_name.map(|f| {
    File::open(f)
     .unwrap_or_else(|_| panic!("Unable to open \"{}\" for reading two-pass data.", f))
  });
  let mut pass1file = pass1file_name.map(|f| {
    File::create(f)
     .unwrap_or_else(|_| panic!("Unable to open \"{}\" for writing two-pass data.", f))
  });

  let mut buffer: [u8; 80] = [0; 80];
  let mut buf_pos = 0;
186

Luca Barbato's avatar
Luca Barbato committed
187
  while let Some(frame_info) =
188
189
    process_frame(&mut ctx, &mut *output, source, pass1file.as_mut(),
     pass2file.as_mut(), &mut buffer, &mut buf_pos, y4m_enc.as_mut())
190
191
192
  {
    for frame in frame_info {
      progress.add_frame(frame);
Luca Barbato's avatar
Luca Barbato committed
193
194
      if verbose {
        eprintln!("{} - {}", frame, progress);
195
      } else {
Luca Barbato's avatar
Luca Barbato committed
196
        eprint!("\r{}                    ", progress);
197
198
199
200
201
202
203
204
205
206
      };
    }

    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
207
      eprintln!("\nError: Failed to write stats file! {}\n", e);
208
209
    }
  }
Luca Barbato's avatar
Luca Barbato committed
210
  eprint!("\n{}\n", progress.print_summary());
211
212
}

Guillaume Martres's avatar
Guillaume Martres committed
213
fn main() {
214
  let mut cli = parse_cli();
215
  let mut y4m_dec = y4m::decode(&mut cli.io.input).expect("input is not a y4m file");
216
  let video_info = y4m_dec.get_video_details();
217
  let y4m_enc = match cli.io.rec.as_mut() {
218
    Some(rec) => Some(
219
220
221
      y4m::encode(
        video_info.width,
        video_info.height,
Luca Barbato's avatar
Luca Barbato committed
222
        y4m::Ratio::new(video_info.time_base.den as usize, video_info.time_base.num as usize)
223
      ).with_colorspace(y4m_dec.get_colorspace())
224
225
226
        .write_header(rec)
        .unwrap()
    ),
Michael Bebenita's avatar
Michael Bebenita committed
227
228
    None => None
  };
229

230
231
232
233
234
  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;
235
236
237

  // If no pixel range is specified via CLI, assume limited,
  // as it is the default for the Y4M format.
238
  if !cli.color_range_specified {
239
240
241
    cli.enc.pixel_range = PixelRange::Limited;
  }

242
  cli.enc.time_base = video_info.time_base;
Kyle Siefring's avatar
Kyle Siefring committed
243
  let cfg = Config {
Luca Barbato's avatar
Luca Barbato committed
244
245
    enc: cli.enc,
    threads: cli.threads,
Kyle Siefring's avatar
Kyle Siefring committed
246
  };
Luca Barbato's avatar
Luca Barbato committed
247

Luca Barbato's avatar
Luca Barbato committed
248
  eprintln!("{}x{} @ {}/{} fps",
249
250
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
251
252
    video_info.time_base.den,
    video_info.time_base.num
253
  );
254

255
  cli.io.output.write_header(
256
257
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
258
259
    video_info.time_base.den as usize,
    video_info.time_base.num as usize
Michael Bebenita's avatar
Michael Bebenita committed
260
  );
Guillaume Martres's avatar
Guillaume Martres committed
261

262
  let progress = ProgressInfo::new(
Luca Barbato's avatar
Luca Barbato committed
263
    Rational { num: video_info.time_base.den, den: video_info.time_base.num },
264
265
    if cli.limit == 0 { None } else { Some(cli.limit) },
      cfg.enc.show_psnr
266
  );
Luca Barbato's avatar
Luca Barbato committed
267

268
269
  for _ in 0..cli.skip {
    y4m_dec.read_frame().expect("Skipped more frames than in the input");
Vibhoothi's avatar
Vibhoothi committed
270
  }
271

272
  #[cfg(all(unix, feature = "signal-hook"))]
Luca Barbato's avatar
Luca Barbato committed
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
  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
  };

296
  #[cfg(all(unix, feature = "signal-hook"))]
Luca Barbato's avatar
Luca Barbato committed
297
298
299
300
301
302
  let mut source = Source {
    limit: cli.limit,
    input: y4m_dec,
    count: 0,
    exit_requested
  };
303
  #[cfg(not(all(unix, feature = "signal-hook")))]
304
305
  let mut source = Source { limit: cli.limit, input: y4m_dec, count: 0 };

306
  if video_info.bit_depth == 8 {
307
    do_encode::<u8, y4m::Decoder<'_, Box<dyn Read>>>(
308
309
      cfg, cli.verbose, progress, &mut *cli.io.output, &mut source,
      cli.pass1file_name.as_ref(), cli.pass2file_name.as_ref(), y4m_enc
310
311
    )
  } else {
312
    do_encode::<u16, y4m::Decoder<'_, Box<dyn Read>>>(
313
314
      cfg, cli.verbose, progress, &mut *cli.io.output, &mut source,
      cli.pass1file_name.as_ref(), cli.pass2file_name.as_ref(), y4m_enc
315
    )
316
  }
Guillaume Martres's avatar
Guillaume Martres committed
317
}