nihav-core/src/test/dec_video.rs

   1 use std::fs::File;
   2 use std::io::prelude::*;
   3 use crate::frame::*;
   4 use crate::codecs::*;
   5 use crate::demuxers::*;
   6 //use crate::io::byteio::*;
   7 use super::wavwriter::WavWriter;
   8
   9 fn write_pgmyuv(pfx: &str, strno: usize, num: u64, frmref: NAFrameRef) {
  10     let frm = frmref.borrow();
  11     if let NABufferType::None = frm.get_buffer() { return; }
  12     let name = format!("assets/{}out{:02}_{:06}.pgm", pfx, strno, num);
  13     let mut ofile = File::create(name).unwrap();
  14     let buf = frm.get_buffer().get_vbuf().unwrap();
  15     let (w, h) = buf.get_dimensions(0);
  16     let (w2, h2) = buf.get_dimensions(1);
  17     let has_alpha = buf.get_info().get_format().has_alpha();
  18     let mut tot_h = h + h2;
  19     if has_alpha {
  20         tot_h += h;
  21     }
  22     if w2 > w/2 {
  23         tot_h += h2;
  24     }
  25     let hdr = format!("P5\n{} {}\n255\n", w, tot_h);
  26     ofile.write_all(hdr.as_bytes()).unwrap();
  27     let dta = buf.get_data();
  28     let ls = buf.get_stride(0);
  29     let mut idx = 0;
  30     let mut idx2 = w;
  31     for _ in 0..h {
  32         let line = &dta[idx..idx2];
  33         ofile.write_all(line).unwrap();
  34         idx  += ls;
  35         idx2 += ls;
  36     }
  37     if w2 <= w/2 {
  38         let mut pad: Vec<u8> = Vec::with_capacity((w - w2 * 2) / 2);
  39         pad.resize((w - w2 * 2) / 2, 0xFF);
  40         let mut base1 = buf.get_offset(1);
  41         let stride1 = buf.get_stride(1);
  42         let mut base2 = buf.get_offset(2);
  43         let stride2 = buf.get_stride(2);
  44         for _ in 0..h2 {
  45             let bend1 = base1 + w2;
  46             let line = &dta[base1..bend1];
  47             ofile.write_all(line).unwrap();
  48             ofile.write_all(pad.as_slice()).unwrap();
  49
  50             let bend2 = base2 + w2;
  51             let line = &dta[base2..bend2];
  52             ofile.write_all(line).unwrap();
  53             ofile.write_all(pad.as_slice()).unwrap();
  54
  55             base1 += stride1;
  56             base2 += stride2;
  57         }
  58     } else {
  59         let mut pad: Vec<u8> = Vec::with_capacity(w - w2);
  60         pad.resize(w - w2, 0xFF);
  61         let mut base1 = buf.get_offset(1);
  62         let stride1 = buf.get_stride(1);
  63         for _ in 0..h2 {
  64             let bend1 = base1 + w2;
  65             let line = &dta[base1..bend1];
  66             ofile.write_all(line).unwrap();
  67             ofile.write_all(pad.as_slice()).unwrap();
  68             base1 += stride1;
  69         }
  70         let mut base2 = buf.get_offset(2);
  71         let stride2 = buf.get_stride(2);
  72         for _ in 0..h2 {
  73             let bend2 = base2 + w2;
  74             let line = &dta[base2..bend2];
  75             ofile.write_all(line).unwrap();
  76             ofile.write_all(pad.as_slice()).unwrap();
  77             base2 += stride2;
  78         }
  79     }
  80     if has_alpha {
  81         let ls = buf.get_stride(3);
  82         let mut idx = buf.get_offset(3);
  83         let mut idx2 = idx + w;
  84         for _ in 0..h {
  85             let line = &dta[idx..idx2];
  86             ofile.write_all(line).unwrap();
  87             idx  += ls;
  88             idx2 += ls;
  89         }
  90     }
  91 }
  92
  93 fn write_palppm(pfx: &str, strno: usize, num: u64, frmref: NAFrameRef) {
  94     let frm = frmref.borrow();
  95     let name = format!("assets/{}out{:02}_{:06}.ppm", pfx, strno, num);
  96     let mut ofile = File::create(name).unwrap();
  97     let buf = frm.get_buffer().get_vbuf().unwrap();
  98     let (w, h) = buf.get_dimensions(0);
  99     let paloff = buf.get_offset(1);
 100     let hdr = format!("P6\n{} {}\n255\n", w, h);
 101     ofile.write_all(hdr.as_bytes()).unwrap();
 102     let dta = buf.get_data();
 103     let ls = buf.get_stride(0);
 104     let offs: [usize; 3] = [
 105             buf.get_info().get_format().get_chromaton(0).unwrap().get_offset() as usize,
 106             buf.get_info().get_format().get_chromaton(1).unwrap().get_offset() as usize,
 107             buf.get_info().get_format().get_chromaton(2).unwrap().get_offset() as usize
 108         ];
 109     let mut idx  = 0;
 110     let mut line: Vec<u8> = Vec::with_capacity(w * 3);
 111     line.resize(w * 3, 0);
 112     for _ in 0..h {
 113         let src = &dta[idx..(idx+w)];
 114         for x in 0..w {
 115             let pix = src[x] as usize;
 116             line[x * 3 + 0] = dta[paloff + pix * 3 + offs[0]];
 117             line[x * 3 + 1] = dta[paloff + pix * 3 + offs[1]];
 118             line[x * 3 + 2] = dta[paloff + pix * 3 + offs[2]];
 119         }
 120         ofile.write_all(line.as_slice()).unwrap();
 121         idx  += ls;
 122     }
 123 }
 124
 125 fn write_ppm(pfx: &str, strno: usize, num: u64, frmref: NAFrameRef) {
 126     let frm = frmref.borrow();
 127     let name = format!("assets/{}out{:02}_{:06}.ppm", pfx, strno, num);
 128     let mut ofile = File::create(name).unwrap();
 129     if let NABufferType::VideoPacked(ref buf) = frm.get_buffer() {
 130         let (w, h) = buf.get_dimensions(0);
 131         let hdr = format!("P6\n{} {}\n255\n", w, h);
 132         ofile.write_all(hdr.as_bytes()).unwrap();
 133         let dta = buf.get_data();
 134         let stride = buf.get_stride(0);
 135         let offs: [usize; 3] = [
 136                 buf.get_info().get_format().get_chromaton(0).unwrap().get_offset() as usize,
 137                 buf.get_info().get_format().get_chromaton(1).unwrap().get_offset() as usize,
 138                 buf.get_info().get_format().get_chromaton(2).unwrap().get_offset() as usize
 139             ];
 140         let step = buf.get_info().get_format().get_elem_size() as usize;
 141         let mut line: Vec<u8> = Vec::with_capacity(w * 3);
 142         line.resize(w * 3, 0);
 143         for src in dta.chunks(stride) {
 144             for x in 0..w {
 145                 line[x * 3 + 0] = src[x * step + offs[0]];
 146                 line[x * 3 + 1] = src[x * step + offs[1]];
 147                 line[x * 3 + 2] = src[x * step + offs[2]];
 148             }
 149             ofile.write_all(line.as_slice()).unwrap();
 150         }
 151     } else if let NABufferType::Video16(ref buf) = frm.get_buffer() {
 152         let (w, h) = buf.get_dimensions(0);
 153         let hdr = format!("P6\n{} {}\n255\n", w, h);
 154         ofile.write_all(hdr.as_bytes()).unwrap();
 155         let dta = buf.get_data();
 156         let stride = buf.get_stride(0);
 157         let depths: [u8; 3] = [
 158                 buf.get_info().get_format().get_chromaton(0).unwrap().get_depth(),
 159                 buf.get_info().get_format().get_chromaton(1).unwrap().get_depth(),
 160                 buf.get_info().get_format().get_chromaton(2).unwrap().get_depth()
 161             ];
 162         let masks: [u16; 3] = [
 163                 (1 << depths[0]) - 1,
 164                 (1 << depths[1]) - 1,
 165                 (1 << depths[2]) - 1
 166             ];
 167         let shifts: [u8; 3] = [
 168                 buf.get_info().get_format().get_chromaton(0).unwrap().get_shift(),
 169                 buf.get_info().get_format().get_chromaton(1).unwrap().get_shift(),
 170                 buf.get_info().get_format().get_chromaton(2).unwrap().get_shift()
 171             ];
 172         let mut line: Vec<u8> = Vec::with_capacity(w * 3);
 173         line.resize(w * 3, 0);
 174         for src in dta.chunks(stride) {
 175             for x in 0..w {
 176                 let elem = src[x];
 177                 let r = ((elem >> shifts[0]) & masks[0]) << (8 - depths[0]);
 178                 let g = ((elem >> shifts[1]) & masks[1]) << (8 - depths[1]);
 179                 let b = ((elem >> shifts[2]) & masks[2]) << (8 - depths[2]);
 180                 line[x * 3 + 0] = r as u8;
 181                 line[x * 3 + 1] = g as u8;
 182                 line[x * 3 + 2] = b as u8;
 183             }
 184             ofile.write_all(line.as_slice()).unwrap();
 185         }
 186     } else {
 187 panic!(" unhandled buf format");
 188     }
 189 }
 190
 191 /*fn open_wav_out(pfx: &str, strno: usize) -> WavWriter {
 192     let name = format!("assets/{}out{:02}.wav", pfx, strno);
 193     let mut file = File::create(name).unwrap();
 194     let mut fw = FileWriter::new_write(&mut file);
 195     let mut wr = ByteWriter::new(&mut fw);
 196     WavWriter::new(&mut wr)
 197 }*/
 198
 199 pub fn test_file_decoding(demuxer: &str, name: &str, limit: Option<u64>,
 200                           decode_video: bool, decode_audio: bool,
 201                           video_pfx: Option<&str>,
 202                           dmx_reg: &RegisteredDemuxers, dec_reg: &RegisteredDecoders) {
 203     let dmx_f = dmx_reg.find_demuxer(demuxer).unwrap();
 204     let mut file = File::open(name).unwrap();
 205     let mut fr = FileReader::new_read(&mut file);
 206     let mut br = ByteReader::new(&mut fr);
 207     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 208
 209     let mut decs: Vec<Option<Box<NADecoder>>> = Vec::new();
 210     for i in 0..dmx.get_num_streams() {
 211         let s = dmx.get_stream(i).unwrap();
 212         let info = s.get_info();
 213         let decfunc = dec_reg.find_decoder(info.get_name());
 214         if let Some(df) = decfunc {
 215             if (decode_video && info.is_video()) || (decode_audio && info.is_audio()) {
 216                 let mut dec = (df)();
 217                 dec.init(info).unwrap();
 218                 decs.push(Some(dec));
 219             } else {
 220                 decs.push(None);
 221             }
 222         } else {
 223             decs.push(None);
 224         }
 225     }
 226
 227     loop {
 228         let pktres = dmx.get_frame();
 229         if let Err(e) = pktres {
 230             if e == DemuxerError::EOF { break; }
 231             panic!("error");
 232         }
 233         let pkt = pktres.unwrap();
 234         if limit.is_some() && pkt.get_pts().is_some() {
 235             if pkt.get_pts().unwrap() > limit.unwrap() { break; }
 236         }
 237         let streamno = pkt.get_stream().get_id() as usize;
 238         if let Some(ref mut dec) = decs[streamno] {
 239             let frm = dec.decode(&pkt).unwrap();
 240             if pkt.get_stream().get_info().is_video() && video_pfx.is_some() && frm.borrow().get_frame_type() != FrameType::Skip {
 241                 let pfx = video_pfx.unwrap();
 242                 let pts = if let Some(fpts) = frm.borrow().get_pts() { fpts } else { pkt.get_pts().unwrap() };
 243                 let vinfo = frm.borrow().get_buffer().get_video_info().unwrap();
 244                 if vinfo.get_format().is_paletted() {
 245                     write_palppm(pfx, streamno, pts, frm);
 246                 } else if vinfo.get_format().get_model().is_yuv() {
 247                     write_pgmyuv(pfx, streamno, pts, frm);
 248                 } else if vinfo.get_format().get_model().is_rgb() {
 249                     write_ppm(pfx, streamno, pts, frm);
 250                 } else {
 251 panic!(" unknown format");
 252                 }
 253             }
 254         }
 255     }
 256 }
 257
 258 pub fn test_decode_audio(demuxer: &str, name: &str, limit: Option<u64>, audio_pfx: &str,
 259                          dmx_reg: &RegisteredDemuxers, dec_reg: &RegisteredDecoders) {
 260     let dmx_f = dmx_reg.find_demuxer(demuxer).unwrap();
 261     let mut file = File::open(name).unwrap();
 262     let mut fr = FileReader::new_read(&mut file);
 263     let mut br = ByteReader::new(&mut fr);
 264     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 265
 266     let mut decs: Vec<Option<Box<NADecoder>>> = Vec::new();
 267     for i in 0..dmx.get_num_streams() {
 268         let s = dmx.get_stream(i).unwrap();
 269         let info = s.get_info();
 270         let decfunc = dec_reg.find_decoder(info.get_name());
 271         if let Some(df) = decfunc {
 272             if info.is_audio() {
 273                 let mut dec = (df)();
 274                 dec.init(info).unwrap();
 275                 decs.push(Some(dec));
 276             } else {
 277                 decs.push(None);
 278             }
 279         } else {
 280             decs.push(None);
 281         }
 282     }
 283
 284     let name = format!("assets/{}out.wav", audio_pfx);
 285     let file = File::create(name).unwrap();
 286     let mut fw = FileWriter::new_write(file);
 287     let mut wr = ByteWriter::new(&mut fw);
 288     let mut wwr = WavWriter::new(&mut wr);
 289     let mut wrote_header = false;
 290
 291     loop {
 292         let pktres = dmx.get_frame();
 293         if let Err(e) = pktres {
 294             if e == DemuxerError::EOF { break; }
 295             panic!("error");
 296         }
 297         let pkt = pktres.unwrap();
 298         if limit.is_some() && pkt.get_pts().is_some() {
 299             if pkt.get_pts().unwrap() > limit.unwrap() { break; }
 300         }
 301         let streamno = pkt.get_stream().get_id() as usize;
 302         if let Some(ref mut dec) = decs[streamno] {
 303             let frm_ = dec.decode(&pkt).unwrap();
 304             let frm = frm_.borrow();
 305             if frm.get_info().is_audio() {
 306                 if !wrote_header {
 307                     wwr.write_header(frm.get_info().as_ref().get_properties().get_audio_info().unwrap()).unwrap();
 308                     wrote_header = true;
 309                 }
 310                 wwr.write_frame(frm.get_buffer()).unwrap();
 311             }
 312         }
 313     }
 314 }