diff --git a/examples/main.rs b/examples/main.rs index 5df7887..a8a5cad 100644 --- a/examples/main.rs +++ b/examples/main.rs @@ -1,11 +1,12 @@ -use thermaldecoder::decode_raw; +use thermaldecoder::decode_to_files; use std::env; fn main() -> anyhow::Result<()> { let mut arg = env::args(); + arg.next(); // skip executable let filename = arg.next().ok_or(anyhow::anyhow!("unexpected"))?; let frames = arg.next().unwrap_or("frames".into()); - decode_raw(&filename, &frames)?; + decode_to_files(&filename, &frames)?; Ok(()) } diff --git a/src/lib.rs b/src/lib.rs index 6619b89..2e02c3a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,17 +1,17 @@ -use indicatif::{ProgressBar, ProgressIterator}; +use indicatif::ProgressBar; +use indicatif::ProgressBarIter; use png; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use std::fs::File; use std::io::BufWriter; -use std::io::Write; use std::path::Path; use pcap_parser::traits::PcapReaderIterator; use pcap_parser::*; struct PacketsIterator { - cap: PcapNGReader, + cap: PcapNGReader>, i: usize, size: usize, } @@ -73,14 +73,16 @@ impl Iterator for PacketsIterator { impl PacketsIterator { fn new(filename: &str) -> anyhow::Result { let file = File::open(filename)?; - let cap = PcapNGReader::new(65535, file)?; + let pb = ProgressBar::new(file.metadata()?.len()); + let wrap = pb.wrap_read(file); + let cap = PcapNGReader::new(65535, wrap)?; Ok(PacketsIterator { cap, i: 0, size: 0 }) } } #[repr(C, packed)] #[derive(Clone, Debug)] -struct Header { +pub struct Header { c1: u32, c2: u16, part: u16, @@ -118,107 +120,145 @@ impl Header { const HDR_SIZE: usize = std::mem::size_of::
(); -#[pyfunction] -fn decode(filename: &str, frames_root: &str) -> PyResult>> { - let res = - decode_raw(filename, frames_root).map_err(|_| PyValueError::new_err("failed to read"))?; - Ok(res.into()) -} - -struct Frame { +pub struct Frame { header: Header, - pixels: Vec, + raw: Vec, } -struct Decoder {} +impl Frame { + fn pixels(&self) -> Vec { + (0..self.raw.len()) + .step_by(2) + .map(|i| u16::from_be_bytes([self.raw[i], self.raw[i + 1]])) + .collect() + } +} + +pub struct Decoder { + packiter: PacketsIterator, + parts: Vec>, +} impl Decoder { - fn new(filename: &str) {} + pub fn new(filename: &str) -> anyhow::Result { + let packiter = PacketsIterator::new(filename)?; + Ok(Decoder { + packiter, + parts: vec![], + }) + } } impl Iterator for Decoder { type Item = Frame; fn next(&mut self) -> Option { - todo!() + let mut ret = None; + while let Some(packet) = self.packiter.next() { + let hdr = match Header::read(&packet) { + Ok(header) => header, + Err(e) => { + println!("error reading header: {:?}", e); + return None; + } + }; + let data = packet[HDR_SIZE..].to_vec(); + if hdr.part == 0 && self.parts.len() > 0 { + ret = Some(Frame { + header: hdr, + raw: self.parts.concat(), + }); + self.parts.clear(); + } + self.parts.push(data); + if ret.is_some() { + return ret; + } + } + None } } -pub fn decode_raw(filename: &str, frames_root: &str) -> anyhow::Result>> { - let mut packiter = PacketsIterator::new(filename)?; - let mut data = vec![]; - while let Some(p) = packiter.next() { - data.push(p); - } - let (i, size) = (packiter.i, packiter.size); +fn write_raw_frame(name: &str, data: &[u8]) -> anyhow::Result<()> { + let path = Path::new(&name); + let file = File::create(path)?; + let ref mut w = BufWriter::new(file); + let mut encoder = png::Encoder::new(w, 288, 384); + encoder.set_color(png::ColorType::Grayscale); + encoder.set_depth(png::BitDepth::Sixteen); + let mut writer = encoder.write_header()?; + writer.write_image_data(data)?; + Ok(()) +} - println!("found {} packets, saved {}, {} size", i, data.len(), size); - let dump_filename = format!("{}/dump.bin", filename); - let mut dump = File::create_new(dump_filename); - let mut frames = vec![]; - let mut parts = vec![]; - for packet in data.iter() { - if let Ok(ref mut dump) = dump { - dump.write_all(&packet)?; - } - let hdr = Header::read(packet)?; - let data = packet[HDR_SIZE..].to_vec(); - if hdr.part == 0 && parts.len() > 0 { - frames.push(parts.concat()); - parts.clear(); - } - parts.push(data); +fn write_calibrated_frame(name: &str, data: &[u16]) -> anyhow::Result<()> { + let path = Path::new(&name); + let file = File::create(path).unwrap(); + let ref mut w = BufWriter::new(file); + let mut encoder = png::Encoder::new(w, 288, 384); + encoder.set_color(png::ColorType::Grayscale); + encoder.set_depth(png::BitDepth::Eight); + let mut writer = encoder.write_header()?; + //let samples: &[u16] = unsafe { std::slice::from_raw_parts(p.cast(), frame.len() / 2) }; + let frame = data + .iter() + .copied() + .map(|x| { + let x: f64 = x.into(); + let x = x / 256.0; + ((-1.665884e-08) * x.powf(4.) + + (1.347094e-05) * x.powf(3.) + + (-4.396264e-03) * x.powf(2.) + + (9.506939e-01) * x + + (-6.353247e+01)) as u8 + }) + .collect::>(); + writer.write_image_data(&frame)?; + Ok(()) +} + +#[pyclass] +struct PyFrameIterator { + iter: Decoder, +} + +#[pymethods] +impl PyFrameIterator { + fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf } - println!("found {} frames", frames.len()); - println!("writing raw pngs"); - let pb = ProgressBar::new(frames.len() as u64); - for (i, frame) in frames.iter().enumerate().progress_with(pb) { + fn __next__(mut slf: PyRefMut<'_, Self>) -> Option { + match slf.iter.next() { + Some(f) => { + let pixels = f.pixels(); + let p = Python::with_gil(|py| pixels.into_py(py)); + Some(p) + } + None => None.into(), + } + } +} + +#[pyfunction] +fn decode(filename: &str) -> PyResult { + let decoder = Decoder::new(filename).map_err(|_| PyValueError::new_err("failed to read"))?; + let iter = PyFrameIterator { iter: decoder }; + Ok(iter.into()) +} + +pub fn decode_to_files(filename: &str, frames_root: &str) -> anyhow::Result<()> { + let frameiter = Decoder::new(filename)?; + for (i, frame) in frameiter.enumerate() { let name = format!("{}/{:04}.png", frames_root, i); - let path = Path::new(&name); - let file = File::create(path).unwrap(); - let ref mut w = BufWriter::new(file); - let mut encoder = png::Encoder::new(w, 288, 384); - encoder.set_color(png::ColorType::Grayscale); - encoder.set_depth(png::BitDepth::Sixteen); - let mut writer = encoder.write_header()?; - writer.write_image_data(&frame)?; - } - - println!("writing calibrated (value is temperature)"); - let pb = ProgressBar::new(frames.len() as u64); - let mut framesu16 = vec![]; - for (i, frame) in frames.iter().enumerate().progress_with(pb) { + if let Err(_e) = write_raw_frame(&name, &frame.raw) { + println!("skipping bad frame {}", i); + continue; + } let name = format!("{}/temp_{:04}.png", frames_root, i); - let path = Path::new(&name); - let file = File::create(path).unwrap(); - let ref mut w = BufWriter::new(file); - let mut encoder = png::Encoder::new(w, 288, 384); - encoder.set_color(png::ColorType::Grayscale); - encoder.set_depth(png::BitDepth::Eight); - let mut writer = encoder.write_header()?; - //let samples: &[u16] = unsafe { std::slice::from_raw_parts(p.cast(), frame.len() / 2) }; - let samples: Vec = (0..frame.len()) - .step_by(2) - .map(|i| u16::from_be_bytes([frame[i], frame[i + 1]])) - .collect(); - framesu16.push(samples.clone()); - let frame = samples - .iter() - .copied() - .map(|x| { - let x: f64 = x.into(); - let x = x / 256.0; - ((-1.665884e-08) * x.powf(4.) - + (1.347094e-05) * x.powf(3.) - + (-4.396264e-03) * x.powf(2.) - + (9.506939e-01) * x - + (-6.353247e+01)) as u8 - }) - .collect::>(); - writer.write_image_data(&frame)?; + let pixels = frame.pixels(); + write_calibrated_frame(&name, &pixels)?; } - - Ok(framesu16) + Ok(()) } #[pymodule] diff --git a/test_rust.py b/test_rust.py index 6e56382..5d3a3d7 100755 --- a/test_rust.py +++ b/test_rust.py @@ -5,7 +5,7 @@ import numpy as np import matplotlib.pyplot as plt root = Path('frames') root.mkdir(exist_ok=True) -frames = decode('in.pcap', 'frames') +frames = list(decode('in.pcap')) f = np.array(frames[0]) f.shape = (384, 288) plt.imshow(f)