From ffe695773d38f16f1a6f04086f47c19c1c7b4a92 Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Sat, 16 Feb 2019 10:35:09 +0100 Subject: [PATCH] write RGB data properly and from more formats --- src/frmwriter.rs | 76 ++++++++++++++++++++++++++++++++++++++++++++++-- src/main.rs | 10 +++++-- 2 files changed, 80 insertions(+), 6 deletions(-) diff --git a/src/frmwriter.rs b/src/frmwriter.rs index a01d8b3..1a94fab 100644 --- a/src/frmwriter.rs +++ b/src/frmwriter.rs @@ -57,6 +57,11 @@ pub fn write_palppm(pfx: &str, strno: usize, num: u64, frm: Ref) { ofile.write_all(hdr.as_bytes()).unwrap(); let dta = buf.get_data(); let ls = buf.get_stride(0); + let offs: [usize; 3] = [ + buf.get_info().get_format().get_chromaton(0).unwrap().get_offset() as usize, + buf.get_info().get_format().get_chromaton(1).unwrap().get_offset() as usize, + buf.get_info().get_format().get_chromaton(2).unwrap().get_offset() as usize + ]; let mut idx = 0; let mut line: Vec = Vec::with_capacity(w * 3); line.resize(w * 3, 0); @@ -64,11 +69,76 @@ pub fn write_palppm(pfx: &str, strno: usize, num: u64, frm: Ref) { 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]; + line[x * 3 + 0] = dta[paloff + pix * 3 + offs[0]]; + line[x * 3 + 1] = dta[paloff + pix * 3 + offs[1]]; + line[x * 3 + 2] = dta[paloff + pix * 3 + offs[2]]; } ofile.write_all(line.as_slice()).unwrap(); idx += ls; } } + +pub fn write_ppm(pfx: &str, strno: usize, num: u64, frm: Ref) { + let name = format!("{}out{:02}_{:08}.ppm", pfx, strno, num); + let mut ofile = File::create(name).unwrap(); + if let NABufferType::VideoPacked(ref buf) = frm.get_buffer() { + let (w, h) = buf.get_dimensions(0); + let hdr = format!("P6\n{} {}\n255\n", w, h); + ofile.write_all(hdr.as_bytes()).unwrap(); + let dta = buf.get_data(); + let stride = buf.get_stride(0); + let offs: [usize; 3] = [ + buf.get_info().get_format().get_chromaton(0).unwrap().get_offset() as usize, + buf.get_info().get_format().get_chromaton(1).unwrap().get_offset() as usize, + buf.get_info().get_format().get_chromaton(2).unwrap().get_offset() as usize + ]; + let step = buf.get_info().get_format().get_elem_size() as usize; + let mut line: Vec = Vec::with_capacity(w * 3); + line.resize(w * 3, 0); + for src in dta.chunks(stride) { + for x in 0..w { + line[x * 3 + 0] = src[x * step + offs[0]]; + line[x * 3 + 1] = src[x * step + offs[1]]; + line[x * 3 + 2] = src[x * step + offs[2]]; + } + ofile.write_all(line.as_slice()).unwrap(); + } + } else if let NABufferType::Video16(ref buf) = frm.get_buffer() { + let (w, h) = buf.get_dimensions(0); + let hdr = format!("P6\n{} {}\n255\n", w, h); + ofile.write_all(hdr.as_bytes()).unwrap(); + let dta = buf.get_data(); + let stride = buf.get_stride(0); + let depths: [u8; 3] = [ + buf.get_info().get_format().get_chromaton(0).unwrap().get_depth(), + buf.get_info().get_format().get_chromaton(1).unwrap().get_depth(), + buf.get_info().get_format().get_chromaton(2).unwrap().get_depth() + ]; + let masks: [u16; 3] = [ + (1 << depths[0]) - 1, + (1 << depths[1]) - 1, + (1 << depths[2]) - 1 + ]; + let shifts: [u8; 3] = [ + buf.get_info().get_format().get_chromaton(0).unwrap().get_shift(), + buf.get_info().get_format().get_chromaton(1).unwrap().get_shift(), + buf.get_info().get_format().get_chromaton(2).unwrap().get_shift() + ]; + let mut line: Vec = Vec::with_capacity(w * 3); + line.resize(w * 3, 0); + for src in dta.chunks(stride) { + for x in 0..w { + let elem = src[x]; + let r = ((elem >> shifts[0]) & masks[0]) << (8 - depths[0]); + let g = ((elem >> shifts[1]) & masks[1]) << (8 - depths[1]); + let b = ((elem >> shifts[2]) & masks[2]) << (8 - depths[2]); + line[x * 3 + 0] = r as u8; + line[x * 3 + 1] = g as u8; + line[x * 3 + 2] = b as u8; + } + ofile.write_all(line.as_slice()).unwrap(); + } + } else { +panic!(" unhandled buf format"); + } +} diff --git a/src/main.rs b/src/main.rs index 214427e..d30dd87 100644 --- a/src/main.rs +++ b/src/main.rs @@ -41,11 +41,15 @@ impl FrameOutput { NumberMode::PktPTS => { pkt.get_pts().unwrap() }, NumberMode::FrmPTS => { 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() { + let vinfo = frm.get_buffer().get_video_info().unwrap(); + if vinfo.get_format().is_paletted() { write_palppm(self.prefix, self.streamno, pts, frm); - } else { + } else if vinfo.get_format().get_model().is_yuv() { write_pgmyuv(self.prefix, self.streamno, pts, frm); + } else if vinfo.get_format().get_model().is_rgb() { + write_ppm(self.prefix, self.streamno, pts, frm); + } else { +panic!(" unknown format"); } } self.frameno += 1; -- 2.30.2