src/main.rs

   1 extern crate nihav;
   2
   3 use std::io::prelude::*;
   4 use std::io::SeekFrom;
   5 use std::fs::File;
   6 use std::path::Path;
   7 use nihav::io::byteio::{FileReader, ByteReader};
   8 use nihav::frame::*;
   9 use nihav::codecs::*;
  10 use nihav::demuxers::*;
  11 use nihav::detect;
  12 use std::env;
  13
  14 fn write_pgmyuv(pfx: &str, strno: usize, num: u64, frmref: NAFrameRef) {
  15     let frm = frmref.borrow();
  16     if let NABufferType::None = frm.get_buffer() { return; }
  17     let name = format!("{}out{:02}_{:04}.pgm", pfx, strno, num);
  18     let mut ofile = File::create(name).unwrap();
  19     let buf = frm.get_buffer().get_vbuf().unwrap();
  20     let (w, h) = buf.get_dimensions(0);
  21     let (w2, h2) = buf.get_dimensions(1);
  22     let tot_h = h + h2;
  23     let hdr = format!("P5\n{} {}\n255\n", w, tot_h);
  24     ofile.write_all(hdr.as_bytes()).unwrap();
  25     let dta = buf.get_data();
  26     let ls = buf.get_stride(0);
  27     let mut idx = 0;
  28     let mut idx2 = w;
  29     let mut pad: Vec<u8> = Vec::with_capacity((w - w2 * 2) / 2);
  30     pad.resize((w - w2 * 2) / 2, 0xFF);
  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     let mut base1 = buf.get_offset(1);
  38     let stride1 = buf.get_stride(1);
  39     let mut base2 = buf.get_offset(2);
  40     let stride2 = buf.get_stride(2);
  41     for _ in 0..h2 {
  42         let bend1 = base1 + w2;
  43         let line = &dta[base1..bend1];
  44         ofile.write_all(line).unwrap();
  45         ofile.write_all(pad.as_slice()).unwrap();
  46
  47         let bend2 = base2 + w2;
  48         let line = &dta[base2..bend2];
  49         ofile.write_all(line).unwrap();
  50         ofile.write_all(pad.as_slice()).unwrap();
  51
  52         base1 += stride1;
  53         base2 += stride2;
  54     }
  55 }
  56
  57 fn write_palppm(pfx: &str, strno: usize, num: u64, frmref: NAFrameRef) {
  58     let frm = frmref.borrow();
  59     let name = format!("{}out{:02}_{:04}.ppm", pfx, strno, num);
  60     let mut ofile = File::create(name).unwrap();
  61     let buf = frm.get_buffer().get_vbuf().unwrap();
  62     let (w, h) = buf.get_dimensions(0);
  63     let paloff = buf.get_offset(1);
  64     let hdr = format!("P6\n{} {}\n255\n", w, h);
  65     ofile.write_all(hdr.as_bytes()).unwrap();
  66     let dta = buf.get_data();
  67     let ls = buf.get_stride(0);
  68     let mut idx  = 0;
  69     let mut line: Vec<u8> = Vec::with_capacity(w * 3);
  70     line.resize(w * 3, 0);
  71     for _ in 0..h {
  72         let src = &dta[idx..(idx+w)];
  73         for x in 0..w {
  74             let pix = src[x] as usize;
  75             line[x * 3 + 0] = dta[paloff + pix * 3 + 2];
  76             line[x * 3 + 1] = dta[paloff + pix * 3 + 1];
  77             line[x * 3 + 2] = dta[paloff + pix * 3 + 0];
  78         }
  79         ofile.write_all(line.as_slice()).unwrap();
  80         idx  += ls;
  81     }
  82 }
  83
  84 fn main() {
  85     let args: Vec<_> = env::args().collect();
  86
  87     if args.len() == 1 { return; }
  88     let lastpts = if args.len() > 2 {
  89                       Some(u64::from_str_radix(args[2].as_str(), 10).unwrap())
  90                   } else { None };
  91
  92     let name = args[1].as_str();
  93     let path = Path::new(name);
  94     let mut file = File::open(path).unwrap();
  95     let dmx_fact;
  96     let mut fr = FileReader::new_read(&mut file);
  97     let mut br = ByteReader::new(&mut fr);
  98     let res = detect::detect_format(name, &mut br);
  99     if res.is_none() {
 100         println!("cannot detect format for {}", name);
 101         return;
 102     }
 103     let (dmx_name, _) = res.unwrap();
 104 println!("trying demuxer {}", dmx_name);
 105     dmx_fact = find_demuxer(dmx_name).unwrap();
 106     br.seek(SeekFrom::Start(0)).unwrap();
 107     let mut dmx = dmx_fact.new_demuxer(&mut br);
 108     dmx.open().unwrap();
 109
 110     let mut decs: Vec<Option<Box<NADecoder>>> = Vec::new();
 111     for i in 0..dmx.get_num_streams() {
 112         let s = dmx.get_stream(i).unwrap();
 113         let info = s.get_info();
 114         let decfunc = find_decoder(info.get_name());
 115         if !info.is_video() {
 116             decs.push(None);
 117         } else if let Some(df) = decfunc {
 118             let mut dec = (df)();
 119             dec.init(info).unwrap();
 120             decs.push(Some(dec));
 121         } else {
 122 panic!("decoder {} not found", info.get_name());
 123         }
 124     }
 125
 126     loop {
 127         let pktres = dmx.get_frame();
 128         if let Err(e) = pktres {
 129             if e == DemuxerError::EOF { break; }
 130         }
 131         let pkt = pktres.unwrap();
 132         let streamno = pkt.get_stream().get_id() as usize;
 133         if let Some(ref mut dec) = decs[streamno] {
 134             let frm = dec.decode(&pkt).unwrap();
 135             if pkt.get_stream().get_info().is_video() {
 136                 if frm.borrow().get_frame_type() != FrameType::Skip {
 137                     if frm.borrow().get_buffer().get_vbuf().unwrap().get_info().get_format().is_paletted() {
 138                         write_palppm("", streamno, pkt.get_pts().unwrap(), frm);
 139                     } else {
 140                         write_pgmyuv("", streamno, pkt.get_pts().unwrap(), frm);
 141                     }
 142                 }
 143             }
 144         }
 145         if pkt.get_pts() != None && pkt.get_pts() == lastpts { break; }
 146     }
 147 //panic!("end");
 148 }