From 3660c127f52a92862248464767a3d92cf63aa182 Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Wed, 5 Sep 2018 13:58:25 +0200 Subject: [PATCH] update --- .gitignore | 2 + Cargo.toml | 2 +- src/frmwriter.rs | 74 +++++++++++ src/main.rs | 317 +++++++++++++++++++++++++---------------------- src/wavwriter.rs | 121 ++++++++++++++++++ 5 files changed, 367 insertions(+), 149 deletions(-) create mode 100644 .gitignore create mode 100644 src/frmwriter.rs rewrite src/main.rs (67%) create mode 100644 src/wavwriter.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..324c57f --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +target/ +**/*.rs.bk diff --git a/Cargo.toml b/Cargo.toml index 9ffafdd..31e9517 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,4 +4,4 @@ version = "0.1.0" authors = ["Kostya Shishkov "] [dependencies] -nihav = { path="/home/kst/devel/nihav" } \ No newline at end of file +nihav = { path="/home/kst/devel/nihav" } diff --git a/src/frmwriter.rs b/src/frmwriter.rs new file mode 100644 index 0000000..ac0c7b5 --- /dev/null +++ b/src/frmwriter.rs @@ -0,0 +1,74 @@ +extern crate nihav; + +use nihav::frame::*; +use std::io::prelude::*; +use std::fs::File; +use std::cell::Ref; + +pub fn write_pgmyuv(pfx: &str, strno: usize, num: u64, frm: Ref) { + if let NABufferType::None = frm.get_buffer() { return; } + let name = format!("{}out{:02}_{:04}.pgm", pfx, strno, num); + let mut ofile = File::create(name).unwrap(); + let buf = frm.get_buffer().get_vbuf().unwrap(); + let (w, h) = buf.get_dimensions(0); + let (w2, h2) = buf.get_dimensions(1); + let tot_h = h + h2; + let hdr = format!("P5\n{} {}\n255\n", w, tot_h); + ofile.write_all(hdr.as_bytes()).unwrap(); + let dta = buf.get_data(); + let ls = buf.get_stride(0); + let mut idx = 0; + let mut idx2 = w; + let mut pad: Vec = Vec::with_capacity((w - w2 * 2) / 2); + pad.resize((w - w2 * 2) / 2, 0xFF); + for _ in 0..h { + let line = &dta[idx..idx2]; + ofile.write_all(line).unwrap(); + idx += ls; + idx2 += ls; + } + let mut base1 = buf.get_offset(1); + let stride1 = buf.get_stride(1); + let mut base2 = buf.get_offset(2); + let stride2 = buf.get_stride(2); + for _ in 0..h2 { + let bend1 = base1 + w2; + let line = &dta[base1..bend1]; + ofile.write_all(line).unwrap(); + ofile.write_all(pad.as_slice()).unwrap(); + + let bend2 = base2 + w2; + let line = &dta[base2..bend2]; + ofile.write_all(line).unwrap(); + ofile.write_all(pad.as_slice()).unwrap(); + + base1 += stride1; + base2 += stride2; + } +} + +pub fn write_palppm(pfx: &str, strno: usize, num: u64, frm: Ref) { + let name = format!("{}out{:02}_{:04}.ppm", pfx, strno, num); + let mut ofile = File::create(name).unwrap(); + let buf = frm.get_buffer().get_vbuf().unwrap(); + let (w, h) = buf.get_dimensions(0); + let paloff = buf.get_offset(1); + let hdr = format!("P6\n{} {}\n255\n", w, h); + ofile.write_all(hdr.as_bytes()).unwrap(); + let dta = buf.get_data(); + let ls = buf.get_stride(0); + let mut idx = 0; + let mut line: Vec = Vec::with_capacity(w * 3); + line.resize(w * 3, 0); + for _ in 0..h { + let src = &dta[idx..(idx+w)]; + for x in 0..w { + let pix = src[x] as usize; + line[x * 3 + 0] = dta[paloff + pix * 3 + 2]; + line[x * 3 + 1] = dta[paloff + pix * 3 + 1]; + line[x * 3 + 2] = dta[paloff + pix * 3 + 0]; + } + ofile.write_all(line.as_slice()).unwrap(); + idx += ls; + } +} diff --git a/src/main.rs b/src/main.rs dissimilarity index 67% index 2517b53..0200590 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,148 +1,169 @@ -extern crate nihav; - -use std::io::prelude::*; -use std::io::SeekFrom; -use std::fs::File; -use std::path::Path; -use nihav::io::byteio::{FileReader, ByteReader}; -use nihav::frame::*; -use nihav::codecs::*; -use nihav::demuxers::*; -use nihav::detect; -use std::env; - -fn write_pgmyuv(pfx: &str, strno: usize, num: u64, frmref: NAFrameRef) { - let frm = frmref.borrow(); - if let NABufferType::None = frm.get_buffer() { return; } - let name = format!("{}out{:02}_{:04}.pgm", pfx, strno, num); - let mut ofile = File::create(name).unwrap(); - let buf = frm.get_buffer().get_vbuf().unwrap(); - let (w, h) = buf.get_dimensions(0); - let (w2, h2) = buf.get_dimensions(1); - let tot_h = h + h2; - let hdr = format!("P5\n{} {}\n255\n", w, tot_h); - ofile.write_all(hdr.as_bytes()).unwrap(); - let dta = buf.get_data(); - let ls = buf.get_stride(0); - let mut idx = 0; - let mut idx2 = w; - let mut pad: Vec = Vec::with_capacity((w - w2 * 2) / 2); - pad.resize((w - w2 * 2) / 2, 0xFF); - for _ in 0..h { - let line = &dta[idx..idx2]; - ofile.write_all(line).unwrap(); - idx += ls; - idx2 += ls; - } - let mut base1 = buf.get_offset(1); - let stride1 = buf.get_stride(1); - let mut base2 = buf.get_offset(2); - let stride2 = buf.get_stride(2); - for _ in 0..h2 { - let bend1 = base1 + w2; - let line = &dta[base1..bend1]; - ofile.write_all(line).unwrap(); - ofile.write_all(pad.as_slice()).unwrap(); - - let bend2 = base2 + w2; - let line = &dta[base2..bend2]; - ofile.write_all(line).unwrap(); - ofile.write_all(pad.as_slice()).unwrap(); - - base1 += stride1; - base2 += stride2; - } -} - -fn write_palppm(pfx: &str, strno: usize, num: u64, frmref: NAFrameRef) { - let frm = frmref.borrow(); - let name = format!("{}out{:02}_{:04}.ppm", pfx, strno, num); - let mut ofile = File::create(name).unwrap(); - let buf = frm.get_buffer().get_vbuf().unwrap(); - let (w, h) = buf.get_dimensions(0); - let paloff = buf.get_offset(1); - let hdr = format!("P6\n{} {}\n255\n", w, h); - ofile.write_all(hdr.as_bytes()).unwrap(); - let dta = buf.get_data(); - let ls = buf.get_stride(0); - let mut idx = 0; - let mut line: Vec = Vec::with_capacity(w * 3); - line.resize(w * 3, 0); - for _ in 0..h { - let src = &dta[idx..(idx+w)]; - for x in 0..w { - let pix = src[x] as usize; - line[x * 3 + 0] = dta[paloff + pix * 3 + 2]; - line[x * 3 + 1] = dta[paloff + pix * 3 + 1]; - line[x * 3 + 2] = dta[paloff + pix * 3 + 0]; - } - ofile.write_all(line.as_slice()).unwrap(); - idx += ls; - } -} - -fn main() { - let args: Vec<_> = env::args().collect(); - - if args.len() == 1 { return; } - let lastpts = if args.len() > 2 { - Some(u64::from_str_radix(args[2].as_str(), 10).unwrap()) - } else { None }; - - let name = args[1].as_str(); - let path = Path::new(name); - let mut file = File::open(path).unwrap(); - let dmx_fact; - let mut fr = FileReader::new_read(&mut file); - let mut br = ByteReader::new(&mut fr); - let res = detect::detect_format(name, &mut br); - if res.is_none() { - println!("cannot detect format for {}", name); - return; - } - let (dmx_name, _) = res.unwrap(); -println!("trying demuxer {}", dmx_name); - dmx_fact = find_demuxer(dmx_name).unwrap(); - br.seek(SeekFrom::Start(0)).unwrap(); - let mut dmx = dmx_fact.new_demuxer(&mut br); - dmx.open().unwrap(); - - let mut decs: Vec>> = Vec::new(); - for i in 0..dmx.get_num_streams() { - let s = dmx.get_stream(i).unwrap(); - let info = s.get_info(); - let decfunc = find_decoder(info.get_name()); - if !info.is_video() { - decs.push(None); - } else if let Some(df) = decfunc { - let mut dec = (df)(); - dec.init(info).unwrap(); - decs.push(Some(dec)); - } else { -panic!("decoder {} not found", info.get_name()); - } - } - - loop { - let pktres = dmx.get_frame(); - if let Err(e) = pktres { - if e == DemuxerError::EOF { break; } - } - let pkt = pktres.unwrap(); - let streamno = pkt.get_stream().get_id() as usize; - if let Some(ref mut dec) = decs[streamno] { - let frm = dec.decode(&pkt).unwrap(); - if pkt.get_stream().get_info().is_video() { - if frm.borrow().get_frame_type() != FrameType::Skip { - if frm.borrow().get_buffer().get_vbuf().unwrap().get_info().get_format().is_paletted() { - write_palppm("", streamno, pkt.get_pts().unwrap(), frm); - } else { - write_pgmyuv("", streamno, pkt.get_pts().unwrap(), frm); - } - } - } - } - if pkt.get_pts() != None && pkt.get_pts() == lastpts { break; } - } -//panic!("end"); -} +extern crate nihav; + +use std::io::SeekFrom; +use std::fs::File; +use std::path::Path; +use nihav::io::byteio::{FileReader, ByteReader}; +use nihav::frame::*; +use nihav::codecs::*; +use nihav::demuxers::*; +use nihav::detect; +use std::env; + +mod frmwriter; +use frmwriter::*; + +mod wavwriter; +use wavwriter::WavWriter; + +struct FrameOutput { + prefix: &'static str, + streamno: usize, +} + +impl FrameOutput { + fn output_frame(&mut self, pkt: &NAPacket, frmref: NAFrameRef) { + let frm = frmref.borrow(); + if frm.get_frame_type() != FrameType::Skip { + let pts = if let Some(pt) = frm.get_pts() { pt } else { pkt.get_pts().unwrap() }; + let frmbuf = frm.get_buffer().get_vbuf().unwrap(); + if frmbuf.get_info().get_format().is_paletted() { + write_palppm(self.prefix, self.streamno, pts, frm); + } else { + write_pgmyuv(self.prefix, self.streamno, pts, frm); + } + } + } +} + +struct AudioOutput { + wwr: WavWriter<'static>, + wrote_header: bool, +} + +impl AudioOutput { + fn new(name: &String) -> Self { Self { wwr: WavWriter::new(name), wrote_header: false } } + fn output_frame(&mut self, _pkt: &NAPacket, frm: NAFrameRef) { + if !self.wrote_header { + self.wwr.write_header(frm.borrow().get_info().as_ref().get_properties().get_audio_info().unwrap()).unwrap(); + self.wrote_header = true; + } + self.wwr.write_frame(frm.borrow().get_buffer()).unwrap(); + } +} + +enum Outputter { + Video(FrameOutput), + Audio(AudioOutput), + None, +} + +fn main() { + let args: Vec<_> = env::args().collect(); + + if args.len() == 1 { + println!("usage: nihav-tool [-noout] [-vn] [-an] input [lastpts]"); + return; + } + let mut lastpts: Option = None; + let mut cur_arg: usize = 1; + let mut noout = false; + let mut decode_video = true; + let mut decode_audio = true; + + while (cur_arg < args.len()) && (args[cur_arg].chars().next().unwrap() == '-') { + match args[cur_arg].as_str() { + "--" => { break; }, + "-noout" => { noout = true; }, + "-an" => { decode_audio = false; }, + "-vn" => { decode_video = false; }, + _ => { println!("unknown option {}", args[cur_arg]); return; }, + } + cur_arg += 1; + } + let name = args[cur_arg].as_str(); + cur_arg += 1; +println!(" name = {} arg pos {}/{}", name, cur_arg, args.len()); + if cur_arg < args.len() { + lastpts = Some(u64::from_str_radix(args[cur_arg].as_str(), 10).unwrap()); + } + + let path = Path::new(name); + let mut file = File::open(path).unwrap(); + let dmx_fact; + let mut fr = FileReader::new_read(&mut file); + let mut br = ByteReader::new(&mut fr); + let res = detect::detect_format(name, &mut br); + if res.is_none() { + println!("cannot detect format for {}", name); + return; + } + let (dmx_name, _) = res.unwrap(); +println!("trying demuxer {} on {}", dmx_name, name); + dmx_fact = find_demuxer(dmx_name).unwrap(); + br.seek(SeekFrom::Start(0)).unwrap(); + let mut dmx = create_demuxer(dmx_fact, &mut br).unwrap(); + + let mut decs: Vec>> = Vec::new(); + let mut writers: Vec = Vec::new(); + for i in 0..dmx.get_num_streams() { + let s = dmx.get_stream(i).unwrap(); + let info = s.get_info(); + let decfunc = find_decoder(info.get_name()); +println!("stream {} - {} {}", i, s, info.get_name()); + let mut has_out = false; + if info.is_video() { + if decode_video { + let mut dec = (decfunc.unwrap())(); + dec.init(info).unwrap(); + decs.push(Some(dec)); + if !noout { + writers.push(Outputter::Video(FrameOutput{prefix: "", streamno: i})); + has_out = true; + } + } else { + decs.push(None); + } + } else if info.is_audio() { + if decode_audio { + let mut dec = (decfunc.unwrap())(); + dec.init(info).unwrap(); + decs.push(Some(dec)); + if !noout { + let name = format!("out{:02}.wav", i); + writers.push(Outputter::Audio(AudioOutput::new(&name))); + has_out = true; + } + } else { + decs.push(None); + } + } else { + decs.push(None); +panic!("decoder {} not found", info.get_name()); + } + if !has_out { + writers.push(Outputter::None); + } + } + + loop { + let pktres = dmx.get_frame(); + if let Err(e) = pktres { + if e == DemuxerError::EOF { break; } + } + let pkt = pktres.unwrap(); + let streamno = pkt.get_stream().get_id() as usize; + if let Some(ref mut dec) = decs[streamno] { + let frm = dec.decode(&pkt).unwrap(); + if !noout { + match writers[streamno] { + Outputter::Video(ref mut wr) => { wr.output_frame(&pkt, frm); }, + Outputter::Audio(ref mut wr) => { wr.output_frame(&pkt, frm); }, + _ => {}, + }; + } + } + if pkt.get_pts() != None && pkt.get_pts() == lastpts { break; } + } +//panic!("end"); +} diff --git a/src/wavwriter.rs b/src/wavwriter.rs new file mode 100644 index 0000000..2296f6f --- /dev/null +++ b/src/wavwriter.rs @@ -0,0 +1,121 @@ +extern crate nihav; + +use nihav::io::byteio::*; +use nihav::frame::*; +use std::fs::File; +use std::io::SeekFrom; + +pub struct WavWriter<'a> { + io: Box>, + data_pos: u64, +} + +fn write_byte(wr: &mut ByteWriter, sample: u8) -> ByteIOResult<()> { + wr.write_byte(sample) +} + +fn write_s16(wr: &mut ByteWriter, sample: i16) -> ByteIOResult<()> { + wr.write_u16le(sample as u16) +} + +fn write_s32(wr: &mut ByteWriter, sample: i32) -> ByteIOResult<()> { + wr.write_u16le((sample >> 16) as u16) +} + +fn write_f32(wr: &mut ByteWriter, sample: f32) -> ByteIOResult<()> { + let mut out = (sample * 32768.0) as i32; + if out < -32768 { out = -32768; } + if out > 32767 { out = 32767; } + if out < 0 { out += 65536; } + wr.write_u16le(out as u16) +} + +macro_rules! write_data { + ($wr:expr, $buf:expr, $write:ident) => ({ + let len = $buf.get_length(); + let ainfo = $buf.get_info(); + let nch = ainfo.get_channels() as usize; + let mut offs: Vec = Vec::with_capacity(nch); + for ch in 0..nch { offs.push($buf.get_offset(ch)); } + let data = $buf.get_data(); + + for i in 0..len { + for ch in 0..nch { + let sample = data[offs[ch] + i]; + $write($wr, sample)?; + } + } + }) +} + +impl<'a> WavWriter<'a> { + pub fn new(name: &String) -> Self { + let file = File::create(name).unwrap(); + let fw = Box::new(FileWriter::new_write(file)); + let io = ByteWriter::new(Box::leak(fw)); + WavWriter { io: Box::new(io), data_pos: 0 } + } + pub fn write_header(&mut self, ainfo: NAAudioInfo) -> ByteIOResult<()> { + let bits = ainfo.get_format().get_bits() as usize; + + self.io.write_buf(b"RIFF")?; + self.io.write_u32le(0)?; + self.io.write_buf(b"WAVE")?; + + self.io.write_buf(b"fmt ")?; + self.io.write_u32le(16)?; + self.io.write_u16le(0x0001)?; // PCM + self.io.write_u16le(ainfo.get_channels() as u16)?; + self.io.write_u32le(ainfo.get_sample_rate() as u32)?; + + if bits < 16 { + self.io.write_u32le((ainfo.get_channels() as u32) * (ainfo.get_sample_rate() as u32))?; + self.io.write_u16le(ainfo.get_channels() as u16)?; // block align + self.io.write_u16le(8)?; + } else { + self.io.write_u32le(2 * (ainfo.get_channels() as u32) * (ainfo.get_sample_rate() as u32))?; + self.io.write_u16le((2 * ainfo.get_channels()) as u16)?; // block align + self.io.write_u16le(16)?; + } + + self.io.write_buf(b"data")?; + self.io.write_u32le(0)?; + + self.data_pos = self.io.tell(); + Ok(()) + } + pub fn write_frame(&mut self, abuf: NABufferType) -> ByteIOResult<()> { + match abuf { + NABufferType::AudioU8(ref buf) => { + write_data!(&mut self.io, buf, write_byte); + } + NABufferType::AudioI16(ref buf) => { + write_data!(&mut self.io, buf, write_s16); + } + NABufferType::AudioI32(ref buf) => { + write_data!(&mut self.io, buf, write_s32); + } + NABufferType::AudioF32(ref buf) => { + write_data!(&mut self.io, buf, write_f32); + } + NABufferType::AudioPacked(ref buf) => { + self.io.write_buf(buf.get_data().as_slice())?; + } + _ => {}, + }; + Ok(()) + } +} + +impl<'a> Drop for WavWriter<'a> { + #[allow(unused_variables)] + fn drop(&mut self) { + let size = self.io.tell(); + if (self.data_pos > 0) && (size >= self.data_pos) { + let res = self.io.seek(SeekFrom::Start(4)); + let res = self.io.write_u32le((size - 8) as u32); + let res = self.io.seek(SeekFrom::Start(self.data_pos - 4)); + let res = self.io.write_u32le(((size as u64) - self.data_pos) as u32); + } + } +} -- 2.30.2