diff options
author | Kostya Shishkov <kostya.shishkov@gmail.com> | 2019-02-16 10:35:09 +0100 |
---|---|---|
committer | Kostya Shishkov <kostya.shishkov@gmail.com> | 2019-02-16 10:35:09 +0100 |
commit | ffe695773d38f16f1a6f04086f47c19c1c7b4a92 (patch) | |
tree | c5a53665c36c0064a7df506da054a4a5537eaab9 | |
parent | 98b31ef77be7fdff7d8b209358b37626a5cb7be7 (diff) | |
download | nihav-tool-ffe695773d38f16f1a6f04086f47c19c1c7b4a92.tar.gz |
write RGB data properly and from more formats
-rw-r--r-- | src/frmwriter.rs | 76 | ||||
-rw-r--r-- | 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<NAFrame>) { 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<u8> = 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<NAFrame>) { 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<NAFrame>) { + 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<u8> = 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<u8> = 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; |