Commit f21d3d66 authored by Xu Guangxin's avatar Xu Guangxin Committed by Luca Barbato
Browse files

refact: add Muxer trait for encode output

parent 0adf438b
......@@ -7,6 +7,7 @@
// 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.
use crate::muxer::{create_muxer, Muxer};
use crate::{ColorPrimaries, MatrixCoefficients, TransferCharacteristics};
use clap::{App, AppSettings, Arg, ArgMatches, SubCommand, Shell};
use rav1e::partition::BlockSize;
......@@ -20,7 +21,7 @@ use std::{fmt, io};
pub struct EncoderIO {
pub input: Box<dyn Read>,
pub output: Box<dyn Write>,
pub output: Box<dyn Muxer>,
pub rec: Option<Box<dyn Write>>
}
......@@ -296,10 +297,7 @@ pub fn parse_cli() -> CliOptions {
"-" => Box::new(io::stdin()) as Box<dyn Read>,
f => Box::new(File::open(&f).unwrap()) as Box<dyn Read>
},
output: match matches.value_of("OUTPUT").unwrap() {
"-" => Box::new(io::stdout()) as Box<dyn Write>,
f => Box::new(File::create(&f).unwrap()) as Box<dyn Write>
},
output: create_muxer(matches.value_of("OUTPUT").unwrap()),
rec: matches
.value_of("RECONSTRUCTION")
.map(|f| Box::new(File::create(&f).unwrap()) as Box<dyn Write>)
......
use super::Muxer;
use ivf::*;
use std::fs::File;
use std::io;
use std::io::Write;
pub struct IvfMuxer {
output: Box<dyn Write>,
}
impl Muxer for IvfMuxer {
fn write_header(
&mut self,
width: usize,
height: usize,
framerate_num: usize,
framerate_den: usize
) {
write_ivf_header(
&mut self.output,
width,
height,
framerate_num,
framerate_den,
);
}
fn write_frame(&mut self, pts: u64, data: &[u8]) {
write_ivf_frame(&mut self.output, pts, data);
}
fn flush(&mut self) -> io::Result<()> {
self.output.flush()
}
}
impl IvfMuxer {
pub fn open(path: &str) -> Box<dyn Muxer> {
let ivf = IvfMuxer {
output: match path {
"-" => Box::new(std::io::stdout()),
f => Box::new(File::create(&f).unwrap())
}
};
Box::new(ivf)
}
}
// Copyright (c) 2017-2019, 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.
mod muxer;
pub use muxer::*;
mod ivf;
use self::ivf::IvfMuxer;
use std::io;
pub trait Muxer {
fn write_header(
&mut self, width: usize, height: usize, framerate_num: usize,
framerate_den: usize
);
fn write_frame(&mut self, pts: u64, data: &[u8]);
fn flush(&mut self) -> io::Result<()>;
}
pub fn create_muxer(path: &str) -> Box<dyn Muxer> {
IvfMuxer::open(path)
}
......@@ -8,11 +8,9 @@
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
use crate::decoder::VideoDetails;
use rav1e::*;
use std::io::Write;
use std::slice;
use rav1e::*;
pub use ivf::*;
pub fn write_y4m_frame<T: Pixel>(y4m_enc: &mut y4m::Encoder<'_, Box<dyn Write>>, rec: &rav1e::Frame<T>, y4m_details: VideoDetails) {
let pitch_y = if y4m_details.bit_depth > 8 { y4m_details.width * 2 } else { y4m_details.width };
......
......@@ -7,8 +7,6 @@
// 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.
use y4m;
#[macro_use]
extern crate scan_fmt;
......@@ -16,7 +14,6 @@ mod common;
mod decoder;
mod muxer;
use crate::common::*;
use crate::muxer::*;
use rav1e::*;
use std::io;
......@@ -26,6 +23,7 @@ use std::path::Path;
use std::sync::Arc;
use crate::decoder::Decoder;
use crate::decoder::VideoDetails;
use crate::muxer::*;
use std::fs::File;
use std::io::BufWriter;
......@@ -60,17 +58,16 @@ impl<D: Decoder> Source<D> {
// Encode and write a frame.
// Returns frame information in a `Result`.
fn process_frame<T: Pixel, D: Decoder>(
ctx: &mut Context<T>,
output_file: &mut dyn Write,
ctx: &mut Context<T>, output_file: &mut dyn Muxer,
source: &mut Source<D>,
mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>,
mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> Option<Vec<FrameSummary>> {
let y4m_details = source.input.get_video_details();
let mut frame_summaries = Vec::new();
let pkt_wrapped = ctx.receive_packet();
match pkt_wrapped {
Ok(pkt) => {
write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
output_file.write_frame(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);
}
......@@ -102,7 +99,7 @@ fn write_stats_file<T: Pixel>(ctx: &Context<T>, filename: &Path) -> Result<(), i
fn do_encode<T: Pixel, D: Decoder>(
cfg: Config, verbose: bool, mut progress: ProgressInfo,
mut output: &mut dyn Write,
output: &mut Muxer,
source: &mut Source<D>,
mut y4m_enc: Option<y4m::Encoder<'_, Box<dyn Write>>>
) {
......@@ -110,7 +107,7 @@ fn do_encode<T: Pixel, D: Decoder>(
while let Some(frame_info) =
process_frame(&mut ctx, &mut output, source, y4m_enc.as_mut())
process_frame(&mut ctx, &mut *output, source, y4m_enc.as_mut())
{
for frame in frame_info {
progress.add_frame(frame);
......@@ -169,8 +166,7 @@ fn main() {
video_info.time_base.num
);
write_ivf_header(
&mut cli.io.output,
cli.io.output.write_header(
video_info.width,
video_info.height,
video_info.time_base.den as usize,
......@@ -191,11 +187,11 @@ fn main() {
if video_info.bit_depth == 8 {
do_encode::<u8, y4m::Decoder<'_, Box<dyn Read>>>(
cfg, cli.verbose, progress, &mut cli.io.output, &mut source, y4m_enc
cfg, cli.verbose, progress, &mut *cli.io.output, &mut source, y4m_enc
)
} else {
do_encode::<u16, y4m::Decoder<'_, Box<dyn Read>>>(
cfg, cli.verbose, progress, &mut cli.io.output, &mut source, y4m_enc
cfg, cli.verbose, progress, &mut *cli.io.output, &mut source, y4m_enc
)
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment