rav1e.rs 5.09 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.

10
11
use y4m;

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
19
use crate::common::*;
use crate::muxer::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
20
use rav1e::*;
21

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

32
33
34
35
36
37
fn read_frame<T: Pixel, D: Decoder>(ctx: &mut Context<T>, decoder: &mut D, video_info: VideoDetails) {
  match decoder.read_frame(&video_info) {
    Ok(frame) => {
      match video_info.bit_depth {
        8 | 10 | 12 => {}
        _ => panic!("unknown input bit depth!")
38
      }
39
40
41
42
43
44
45

      let _ = ctx.send_frame(Some(Arc::new(frame)));
    }
    _ => {
      let frames_to_be_coded = ctx.get_frame_count();
      // This is a hack, instead when EOF is reached simply "close" the encoder to input (flag)
      ctx.set_limit(frames_to_be_coded);
46
47
      ctx.flush();
    }
48
  };
49
50
51
52
}

// Encode and write a frame.
// Returns frame information in a `Result`.
53
54
fn process_frame<T: Pixel>(
  ctx: &mut Context<T>,
55
  output_file: &mut dyn Write,
56
  y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
57
  mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>,
58
59
60
61
) -> Result<Vec<FrameSummary>, ()> {
  let y4m_details = y4m_dec.get_video_details();
  let mut frame_summaries = Vec::new();
  let pkt_wrapped = ctx.receive_packet();
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
  match pkt_wrapped {
    Ok(pkt) => {
      write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
      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());
    }
    Err(EncoderStatus::NeedMoreFrames) => {
      read_frame(ctx, y4m_dec, y4m_details);
    }
    Err(EncoderStatus::NeedMoreData) | Err(EncoderStatus::EnoughData) => {
      // Expected statuses, continue without error
    }
    Err(EncoderStatus::Failure) | Err(EncoderStatus::InvalidKey) | Err(EncoderStatus::ParseError) => {
      panic!("Failed to encode video");
78
79
80
81
    }
  }
  Ok(frame_summaries)
}
Guillaume Martres's avatar
Guillaume Martres committed
82

83
fn write_stats_file<T: Pixel>(ctx: &Context<T>, filename: &Path) -> Result<(), io::Error> {
84
85
86
87
88
89
  let file = File::create(filename)?;
  let writer = BufWriter::new(file);
  serde_json::to_writer(writer, &ctx.first_pass_data).expect("Serialization should not fail");
  Ok(())
}

Guillaume Martres's avatar
Guillaume Martres committed
90
fn main() {
91
  let mut cli = parse_cli();
92
  let mut y4m_dec = y4m::decode(&mut cli.io.input).expect("input is not a y4m file");
93
  let video_info = y4m_dec.get_video_details();
94
  let mut y4m_enc = match cli.io.rec.as_mut() {
95
    Some(rec) => Some(
96
97
98
      y4m::encode(
        video_info.width,
        video_info.height,
Luca Barbato's avatar
Luca Barbato committed
99
        y4m::Ratio::new(video_info.time_base.den as usize, video_info.time_base.num as usize)
100
      ).with_colorspace(y4m_dec.get_colorspace())
101
102
103
        .write_header(rec)
        .unwrap()
    ),
Michael Bebenita's avatar
Michael Bebenita committed
104
105
    None => None
  };
106

107
108
109
110
111
112
  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
113
  let cfg = Config {
114
    enc: cli.enc
Kyle Siefring's avatar
Kyle Siefring committed
115
  };
Luca Barbato's avatar
Luca Barbato committed
116

117
118
  // FIXME for now, unconditionally create Context<u16>
  let mut ctx: Context<u16> = cfg.new_context();
Luca Barbato's avatar
Luca Barbato committed
119

120
121
122
  let stderr = io::stderr();
  let mut err = stderr.lock();

123
124
125
126
127
  let _ = writeln!(
    err,
    "{}x{} @ {}/{} fps",
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
128
129
    video_info.time_base.den,
    video_info.time_base.num
130
  );
131

Michael Bebenita's avatar
Michael Bebenita committed
132
  write_ivf_header(
133
    &mut cli.io.output,
134
135
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
136
137
    video_info.time_base.den as usize,
    video_info.time_base.num as usize
Michael Bebenita's avatar
Michael Bebenita committed
138
  );
Guillaume Martres's avatar
Guillaume Martres committed
139

140
  let mut progress = ProgressInfo::new(
Luca Barbato's avatar
Luca Barbato committed
141
    Rational { num: video_info.time_base.den, den: video_info.time_base.num },
142
143
    if cli.limit == 0 { None } else { Some(cli.limit) },
      cfg.enc.show_psnr
144
  );
Luca Barbato's avatar
Luca Barbato committed
145

146
  ctx.set_limit(cli.limit as u64);
147

148
149
150
151
152
153
154
155
156
  while let Ok(frame_info) = process_frame(&mut ctx, &mut cli.io.output, &mut y4m_dec, y4m_enc.as_mut()) {
    for frame in frame_info {
      progress.add_frame(frame);
      let _ = if cli.verbose {
        writeln!(err, "{} - {}", frame, progress)
      } else {
        write!(err, "\r{}                    ", progress)
      };
    }
Luca Barbato's avatar
Luca Barbato committed
157

158
    if !ctx.needs_more_frames(progress.frames_encoded() as u64) {
Michael Bebenita's avatar
Michael Bebenita committed
159
      break;
Guillaume Martres's avatar
Guillaume Martres committed
160
    }
Luca Barbato's avatar
Luca Barbato committed
161

162
    cli.io.output.flush().unwrap();
Michael Bebenita's avatar
Michael Bebenita committed
163
  }
164

165
166
167
168
169
170
  if cfg.enc.pass == Some(1) {
    if let Err(e) = write_stats_file(&ctx, cfg.enc.stats_file.as_ref().unwrap()) {
      let _ = writeln!(err, "\nError: Failed to write stats file! {}\n", e);
    }
  }
  let _ = write!(err, "\n{}\n", progress.print_summary());
Guillaume Martres's avatar
Guillaume Martres committed
171
}