rav1e.rs 5.7 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
  let file = File::create(filename)?;
  let writer = BufWriter::new(file);
Luca Barbato's avatar
Luca Barbato committed
86
  serde_json::to_writer(writer, ctx.get_first_pass_data()).expect("Serialization should not fail");
87
88
89
  Ok(())
}

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
118
119
120
121
122
123
124
125
126
127
128
fn do_encode<T: Pixel>(
  cfg: Config, limit: usize, verbose: bool, mut progress: ProgressInfo,
  mut err: std::io::StderrLock, mut output: &mut dyn Write,
  mut y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
  mut y4m_enc: Option<y4m::Encoder<'_, Box<dyn Write>>>
) {
  let mut ctx: Context<T> = cfg.new_context();

  ctx.set_limit(limit as u64);

  while let Ok(frame_info) =
    process_frame(&mut ctx, &mut output, &mut y4m_dec, y4m_enc.as_mut())
  {
    for frame in frame_info {
      progress.add_frame(frame);
      let _ = if verbose {
        writeln!(err, "{} - {}", frame, progress)
      } else {
        write!(err, "\r{}                    ", progress)
      };
    }

    if !ctx.needs_more_frames(progress.frames_encoded() as u64) {
      break;
    }

    output.flush().unwrap();
  }

  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
129
fn main() {
130
  let mut cli = parse_cli();
131
  let mut y4m_dec = y4m::decode(&mut cli.io.input).expect("input is not a y4m file");
132
  let video_info = y4m_dec.get_video_details();
133
  let y4m_enc = match cli.io.rec.as_mut() {
134
    Some(rec) => Some(
135
136
137
      y4m::encode(
        video_info.width,
        video_info.height,
Luca Barbato's avatar
Luca Barbato committed
138
        y4m::Ratio::new(video_info.time_base.den as usize, video_info.time_base.num as usize)
139
      ).with_colorspace(y4m_dec.get_colorspace())
140
141
142
        .write_header(rec)
        .unwrap()
    ),
Michael Bebenita's avatar
Michael Bebenita committed
143
144
    None => None
  };
145

146
147
148
149
150
151
  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
152
  let cfg = Config {
Luca Barbato's avatar
Luca Barbato committed
153
154
    enc: cli.enc,
    threads: cli.threads,
Kyle Siefring's avatar
Kyle Siefring committed
155
  };
Luca Barbato's avatar
Luca Barbato committed
156

157
158
159
  let stderr = io::stderr();
  let mut err = stderr.lock();

160
161
162
163
164
  let _ = writeln!(
    err,
    "{}x{} @ {}/{} fps",
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
165
166
    video_info.time_base.den,
    video_info.time_base.num
167
  );
168

Michael Bebenita's avatar
Michael Bebenita committed
169
  write_ivf_header(
170
    &mut cli.io.output,
171
172
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
173
174
    video_info.time_base.den as usize,
    video_info.time_base.num as usize
Michael Bebenita's avatar
Michael Bebenita committed
175
  );
Guillaume Martres's avatar
Guillaume Martres committed
176

177
  let progress = ProgressInfo::new(
Luca Barbato's avatar
Luca Barbato committed
178
    Rational { num: video_info.time_base.den, den: video_info.time_base.num },
179
180
    if cli.limit == 0 { None } else { Some(cli.limit) },
      cfg.enc.show_psnr
181
  );
Luca Barbato's avatar
Luca Barbato committed
182

183
184
  for _ in 0..cli.skip {
    y4m_dec.read_frame().expect("Skipped more frames than in the input");
Vibhoothi's avatar
Vibhoothi committed
185
  }
186

187
188
189
190
191
192
193
194
  if video_info.bit_depth == 8 {
    do_encode::<u8>(
      cfg, cli.limit, cli.verbose, progress, err, &mut cli.io.output, &mut y4m_dec, y4m_enc
    )
  } else {
    do_encode::<u16>(
      cfg, cli.limit, cli.verbose, progress, err, &mut cli.io.output, &mut y4m_dec, y4m_enc
    )
195
  }
Guillaume Martres's avatar
Guillaume Martres committed
196
}