rav1e.rs 6.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
extern crate clap;
Guillaume Martres's avatar
Guillaume Martres committed
11
12
extern crate rav1e;
extern crate y4m;
13
14
#[macro_use]
extern crate scan_fmt;
Guillaume Martres's avatar
Guillaume Martres committed
15

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

21
use std::slice;
22
23
use std::io;
use std::io::Write;
24
25
use std::io::Read;
use std::sync::Arc;
26
use decoder::Decoder;
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
use decoder::VideoDetails;

fn read_frame_batch<D: Decoder>(ctx: &mut Context, decoder: &mut D, video_info: VideoDetails) {
  loop {
    if ctx.needs_more_lookahead() {
      match decoder.read_frame(&video_info) {
        Ok(frame) => {
          match video_info.bit_depth {
            8 | 10 | 12 => {}
            _ => panic!("unknown input bit depth!")
          }

          let _ = ctx.send_frame(Some(Arc::new(frame)));
          continue;
        }
        _ => {
          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);
          ctx.flush();
        }
      }
    } else if !ctx.needs_more_frames(ctx.get_frame_count()) {
      ctx.flush();
    }
    break;
  }
}

// Encode and write a frame.
// Returns frame information in a `Result`.
fn process_frame(
  ctx: &mut Context, output_file: &mut dyn Write,
  y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
  mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> Result<Vec<FrameSummary>, ()> {
  let y4m_details = y4m_dec.get_video_details();
  let mut frame_summaries = Vec::new();
  read_frame_batch(ctx, y4m_dec, y4m_details);
  let pkt_wrapped = ctx.receive_packet();
  if let Ok(pkt) = pkt_wrapped {
    write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
69
70
71
72
73
74
75
    if let (Some(ref mut y4m_enc_uw), Some(ref rec)) = (y4m_enc.as_mut(), &pkt.rec) {
      let pitch_y = if y4m_details.bit_depth > 8 { y4m_details.width * 2 } else { y4m_details.width };
      let chroma_sampling_period = y4m_details.chroma_sampling.sampling_period();
      let (pitch_uv, height_uv) = (
        pitch_y / chroma_sampling_period.0,
        y4m_details.height / chroma_sampling_period.1
      );
76

77
78
79
80
81
      let (mut rec_y, mut rec_u, mut rec_v) = (
        vec![128u8; pitch_y * y4m_details.height],
        vec![128u8; pitch_uv * height_uv],
        vec![128u8; pitch_uv * height_uv]
      );
82

83
84
85
86
87
      let (stride_y, stride_u, stride_v) = (
        rec.planes[0].cfg.stride,
        rec.planes[1].cfg.stride,
        rec.planes[2].cfg.stride
      );
88

89
90
91
92
93
94
95
96
97
98
99
      for (line, line_out) in rec.planes[0]
        .data_origin()
        .chunks(stride_y)
        .zip(rec_y.chunks_mut(pitch_y))
      {
        if y4m_details.bit_depth > 8 {
          unsafe {
            line_out.copy_from_slice(slice::from_raw_parts::<u8>(
              line.as_ptr() as (*const u8),
              pitch_y
            ));
100
          }
101
102
103
104
        } else {
          line_out.copy_from_slice(
            &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]
          );
105
        }
106
107
108
109
110
111
112
113
114
115
116
117
      }
      for (line, line_out) in rec.planes[1]
        .data_origin()
        .chunks(stride_u)
        .zip(rec_u.chunks_mut(pitch_uv))
      {
        if y4m_details.bit_depth > 8 {
          unsafe {
            line_out.copy_from_slice(slice::from_raw_parts::<u8>(
              line.as_ptr() as (*const u8),
              pitch_uv
            ));
118
          }
119
120
121
122
        } else {
          line_out.copy_from_slice(
            &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
          );
123
        }
124
125
126
127
128
129
130
131
132
133
134
135
      }
      for (line, line_out) in rec.planes[2]
        .data_origin()
        .chunks(stride_v)
        .zip(rec_v.chunks_mut(pitch_uv))
      {
        if y4m_details.bit_depth > 8 {
          unsafe {
            line_out.copy_from_slice(slice::from_raw_parts::<u8>(
              line.as_ptr() as (*const u8),
              pitch_uv
            ));
136
          }
137
138
139
140
        } else {
          line_out.copy_from_slice(
            &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
          );
141
142
        }
      }
143
144
145

      let rec_frame = y4m::Frame::new([&rec_y, &rec_u, &rec_v], None);
      y4m_enc_uw.write_frame(&rec_frame).unwrap();
146
147
148
149
150
    }
    frame_summaries.push(pkt.into());
  }
  Ok(frame_summaries)
}
Guillaume Martres's avatar
Guillaume Martres committed
151
152

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

169
170
171
172
173
174
  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
175
  let cfg = Config {
176
    enc: cli.enc
Kyle Siefring's avatar
Kyle Siefring committed
177
  };
Luca Barbato's avatar
Luca Barbato committed
178
179
180

  let mut ctx = cfg.new_context();

181
182
183
  let stderr = io::stderr();
  let mut err = stderr.lock();

184
185
186
187
188
  let _ = writeln!(
    err,
    "{}x{} @ {}/{} fps",
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
189
190
    video_info.time_base.den,
    video_info.time_base.num
191
  );
192

Michael Bebenita's avatar
Michael Bebenita committed
193
  write_ivf_header(
194
    &mut cli.io.output,
195
196
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
197
198
    video_info.time_base.den as usize,
    video_info.time_base.num as usize
Michael Bebenita's avatar
Michael Bebenita committed
199
  );
Guillaume Martres's avatar
Guillaume Martres committed
200

201
  let mut progress = ProgressInfo::new(
Luca Barbato's avatar
Luca Barbato committed
202
    Rational { num: video_info.time_base.den, den: video_info.time_base.num },
203
204
    if cli.limit == 0 { None } else { Some(cli.limit) },
      cfg.enc.show_psnr
205
  );
Luca Barbato's avatar
Luca Barbato committed
206

207
  ctx.set_limit(cli.limit as u64);
208

209
210
  loop {
    match process_frame(&mut ctx, &mut cli.io.output, &mut y4m_dec, y4m_enc.as_mut()) {
211
212
213
214
215
216
217
218
219
      Ok(frame_info) => {
        for frame in frame_info {
          progress.add_frame(frame);
          let _ = if cli.verbose {
            writeln!(err, "{} - {}", frame, progress)
          } else {
            write!(err, "\r{}                    ", progress)
          };
        }
220
221
222
      },
      Err(_) => break,
    };
Luca Barbato's avatar
Luca Barbato committed
223

224
    if !ctx.needs_more_frames(progress.frames_encoded() as u64) {
Michael Bebenita's avatar
Michael Bebenita committed
225
      break;
Guillaume Martres's avatar
Guillaume Martres committed
226
    }
Luca Barbato's avatar
Luca Barbato committed
227

228
    cli.io.output.flush().unwrap();
Michael Bebenita's avatar
Michael Bebenita committed
229
  }
230
231

  let _ = write!(err, "\n{}\n", progress.print_stats());
Guillaume Martres's avatar
Guillaume Martres committed
232
}