common.rs 7.21 KB
Newer Older
1
use clap::{App, Arg};
2
use rav1e::*;
3 4
use std::fs::File;
use std::io;
5
use std::io::prelude::*;
Luca Barbato's avatar
Luca Barbato committed
6
use std::sync::Arc;
7
use std::slice;
8
use y4m;
9 10

pub struct EncoderIO {
gibix's avatar
gibix committed
11 12 13
    pub input: Box<dyn Read>,
    pub output: Box<dyn Write>,
    pub rec: Option<Box<dyn Write>>,
14 15 16 17 18 19 20 21 22 23 24 25 26 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
}

pub trait FromCli {
    fn from_cli() -> (EncoderIO, EncoderConfig);
}

impl FromCli for EncoderConfig {
    fn from_cli() -> (EncoderIO, EncoderConfig) {
        let matches = App::new("rav1e")
            .version("0.1.0")
            .about("AV1 video encoder")
           .arg(Arg::with_name("INPUT")
                .help("Uncompressed YUV4MPEG2 video input")
                .required(true)
                .index(1))
            .arg(Arg::with_name("OUTPUT")
                .help("Compressed AV1 in IVF video output")
                .short("o")
                .long("output")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("RECONSTRUCTION")
                .short("r")
                .takes_value(true))
            .arg(Arg::with_name("LIMIT")
                .help("Maximum number of frames to encode")
                .short("l")
                .long("limit")
                .takes_value(true)
                .default_value("0"))
            .arg(Arg::with_name("QP")
                .help("Quantizer (0-255)")
                .long("quantizer")
                .takes_value(true)
                .default_value("100"))
            .arg(Arg::with_name("SPEED")
                .help("Speed level (0(slow)-10(fast))")
                .short("s")
                .long("speed")
                .takes_value(true)
                .default_value("3"))
            .arg(Arg::with_name("TUNE")
                .help("Quality tuning (Will enforce partition sizes >= 8x8)")
                .long("tune")
                .possible_values(&Tune::variants())
                .default_value("psnr")
                .case_insensitive(true))
            .get_matches();


        let io = EncoderIO {
            input: match matches.value_of("INPUT").unwrap() {
gibix's avatar
gibix committed
66 67
                "-" => Box::new(io::stdin()) as Box<dyn Read>,
                f => Box::new(File::open(&f).unwrap()) as Box<dyn Read>
68 69
            },
            output: match matches.value_of("OUTPUT").unwrap() {
gibix's avatar
gibix committed
70 71
                "-" => Box::new(io::stdout()) as Box<dyn Write>,
                f => Box::new(File::create(&f).unwrap()) as Box<dyn Write>
72 73
            },
            rec: matches.value_of("RECONSTRUCTION").map(|f| {
gibix's avatar
gibix committed
74
                Box::new(File::create(&f).unwrap()) as Box<dyn Write>
75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94
            })
        };

        let config = EncoderConfig {
            limit: matches.value_of("LIMIT").unwrap().parse().unwrap(),
            quantizer: matches.value_of("QP").unwrap().parse().unwrap(),
            speed: matches.value_of("SPEED").unwrap().parse().unwrap(),
            tune: matches.value_of("TUNE").unwrap().parse().unwrap()
        };

        // Validate arguments
        if config.quantizer == 0 {
            unimplemented!();
        } else if config.quantizer > 255 || config.speed > 10 {
            panic!("argument out of range");
        }

        (io, config)
    }
}
95 96

/// Encode and write a frame.
Luca Barbato's avatar
Luca Barbato committed
97
pub fn process_frame(ctx: &mut Context,
gibix's avatar
gibix committed
98 99 100
                     output_file: &mut dyn Write,
                     y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
                     y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>) -> bool {
101 102 103 104
    unsafe {
        av1_rtcd();
        aom_dsp_rtcd();
    }
Luca Barbato's avatar
Luca Barbato committed
105 106
    let width = y4m_dec.get_width();
    let height = y4m_dec.get_height();
107 108
    let y4m_bits = y4m_dec.get_bit_depth();
    let y4m_bytes = y4m_dec.get_bytes_per_sample();
Luca Barbato's avatar
Luca Barbato committed
109 110
    let bit_depth = y4m_dec.get_colorspace().get_bit_depth();

111 112 113 114 115
    match y4m_dec.read_frame() {
        Ok(y4m_frame) => {
            let y4m_y = y4m_frame.get_y_plane();
            let y4m_u = y4m_frame.get_u_plane();
            let y4m_v = y4m_frame.get_v_plane();
Luca Barbato's avatar
Luca Barbato committed
116
            let mut input = ctx.new_frame();
Luca Barbato's avatar
Luca Barbato committed
117
            {
Luca Barbato's avatar
Luca Barbato committed
118
                let input = Arc::get_mut(&mut input).unwrap();
Luca Barbato's avatar
Luca Barbato committed
119 120 121 122
                input.planes[0].copy_from_raw_u8(&y4m_y, width * y4m_bytes, y4m_bytes);
                input.planes[1].copy_from_raw_u8(&y4m_u, width * y4m_bytes / 2, y4m_bytes);
                input.planes[2].copy_from_raw_u8(&y4m_v, width * y4m_bytes / 2, y4m_bytes);
            }
123 124 125 126 127 128

            match y4m_bits {
                8 | 10 | 12 => {},
                _ => panic! ("unknown input bit depth!"),
            }

Luca Barbato's avatar
Luca Barbato committed
129 130
            let _ = ctx.send_frame(input);
            let pkt = ctx.receive_packet().unwrap();
131
            eprintln!("{}", pkt);
Luca Barbato's avatar
Luca Barbato committed
132
            write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
133
            if let Some(mut y4m_enc) = y4m_enc {
Luca Barbato's avatar
Luca Barbato committed
134
                let pitch_y = if bit_depth > 8 {
135 136 137 138 139 140 141 142 143 144 145 146
                    width * 2
                } else {
                    width
                };
                let pitch_uv = pitch_y / 2;

                let (mut rec_y, mut rec_u, mut rec_v) = (
                    vec![128u8; pitch_y * height],
                    vec![128u8; pitch_uv * (height / 2)],
                    vec![128u8; pitch_uv * (height / 2)]);

                let (stride_y, stride_u, stride_v) = (
Luca Barbato's avatar
Luca Barbato committed
147 148 149
                    pkt.rec.planes[0].cfg.stride,
                    pkt.rec.planes[1].cfg.stride,
                    pkt.rec.planes[2].cfg.stride);
150

Luca Barbato's avatar
Luca Barbato committed
151 152
                for (line, line_out) in pkt.rec.planes[0].data_origin().chunks(stride_y).zip(rec_y.chunks_mut(pitch_y)) {
                    if bit_depth > 8 {
153 154 155 156 157 158 159 160 161
                        unsafe {
                            line_out.copy_from_slice(
                                slice::from_raw_parts::<u8>(line.as_ptr() as (*const u8), pitch_y));
                        }
                    } else {
                        line_out.copy_from_slice(
                            &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]);
                    }
                }
Luca Barbato's avatar
Luca Barbato committed
162 163
                for (line, line_out) in pkt.rec.planes[1].data_origin().chunks(stride_u).zip(rec_u.chunks_mut(pitch_uv)) {
                    if bit_depth > 8 {
164 165 166 167 168 169 170 171 172
                        unsafe {
                            line_out.copy_from_slice(
                                slice::from_raw_parts::<u8>(line.as_ptr() as (*const u8), pitch_uv));
                        }
                    } else {
                        line_out.copy_from_slice(
                            &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]);
                    }
                }
Luca Barbato's avatar
Luca Barbato committed
173 174
                for (line, line_out) in pkt.rec.planes[2].data_origin().chunks(stride_v).zip(rec_v.chunks_mut(pitch_uv)) {
                    if bit_depth > 8 {
175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193
                        unsafe {
                            line_out.copy_from_slice(
                                slice::from_raw_parts::<u8>(line.as_ptr() as (*const u8), pitch_uv));
                        }
                    } else {
                        line_out.copy_from_slice(
                            &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]);
                    }
                }

                let rec_frame = y4m::Frame::new([&rec_y, &rec_u, &rec_v], None);
                y4m_enc.write_frame(&rec_frame).unwrap();
            }

            true
        },
        _ => false
    }
}