rav1e.rs 4.43 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
25
use std::io::Read;
use std::sync::Arc;
Raphaël Zumer's avatar
Raphaël Zumer committed
26
27
use crate::decoder::Decoder;
use crate::decoder::VideoDetails;
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

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
    if let (Some(ref mut y4m_enc_uw), Some(ref rec)) = (y4m_enc.as_mut(), &pkt.rec) {
70
      write_y4m_frame(y4m_enc_uw, rec, y4m_details);
71
72
73
74
75
    }
    frame_summaries.push(pkt.into());
  }
  Ok(frame_summaries)
}
Guillaume Martres's avatar
Guillaume Martres committed
76
77

fn main() {
78
  let mut cli = parse_cli();
79
  let mut y4m_dec = y4m::decode(&mut cli.io.input).expect("input is not a y4m file");
80
  let video_info = y4m_dec.get_video_details();
81
  let mut y4m_enc = match cli.io.rec.as_mut() {
82
    Some(rec) => Some(
83
84
85
      y4m::encode(
        video_info.width,
        video_info.height,
Luca Barbato's avatar
Luca Barbato committed
86
        y4m::Ratio::new(video_info.time_base.den as usize, video_info.time_base.num as usize)
87
      ).with_colorspace(y4m_dec.get_colorspace())
88
89
90
        .write_header(rec)
        .unwrap()
    ),
Michael Bebenita's avatar
Michael Bebenita committed
91
92
    None => None
  };
93

94
95
96
97
98
99
  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
100
  let cfg = Config {
101
    enc: cli.enc
Kyle Siefring's avatar
Kyle Siefring committed
102
  };
Luca Barbato's avatar
Luca Barbato committed
103
104
105

  let mut ctx = cfg.new_context();

106
107
108
  let stderr = io::stderr();
  let mut err = stderr.lock();

109
110
111
112
113
  let _ = writeln!(
    err,
    "{}x{} @ {}/{} fps",
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
114
115
    video_info.time_base.den,
    video_info.time_base.num
116
  );
117

Michael Bebenita's avatar
Michael Bebenita committed
118
  write_ivf_header(
119
    &mut cli.io.output,
120
121
    video_info.width,
    video_info.height,
Luca Barbato's avatar
Luca Barbato committed
122
123
    video_info.time_base.den as usize,
    video_info.time_base.num as usize
Michael Bebenita's avatar
Michael Bebenita committed
124
  );
Guillaume Martres's avatar
Guillaume Martres committed
125

126
  let mut progress = ProgressInfo::new(
Luca Barbato's avatar
Luca Barbato committed
127
    Rational { num: video_info.time_base.den, den: video_info.time_base.num },
128
129
    if cli.limit == 0 { None } else { Some(cli.limit) },
      cfg.enc.show_psnr
130
  );
Luca Barbato's avatar
Luca Barbato committed
131

132
  ctx.set_limit(cli.limit as u64);
133

134
135
  loop {
    match process_frame(&mut ctx, &mut cli.io.output, &mut y4m_dec, y4m_enc.as_mut()) {
136
137
138
139
140
141
142
143
144
      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)
          };
        }
145
146
147
      },
      Err(_) => break,
    };
Luca Barbato's avatar
Luca Barbato committed
148

149
    if !ctx.needs_more_frames(progress.frames_encoded() as u64) {
Michael Bebenita's avatar
Michael Bebenita committed
150
      break;
Guillaume Martres's avatar
Guillaume Martres committed
151
    }
Luca Barbato's avatar
Luca Barbato committed
152

153
    cli.io.output.flush().unwrap();
Michael Bebenita's avatar
Michael Bebenita committed
154
  }
155
156

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