diff --git a/.gitignore b/.gitignore index 8e5c6d2..85aecf6 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ /venv /frames +**/*.png diff --git a/thermalcamdecoder/src/lib.rs b/thermalcamdecoder/src/lib.rs new file mode 100644 index 0000000..ac054e5 --- /dev/null +++ b/thermalcamdecoder/src/lib.rs @@ -0,0 +1,186 @@ +use indicatif::{ProgressBar, ProgressIterator}; +use png; +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::*; + +fn get_packets_without_udp_header() -> anyhow::Result<(Vec>, usize, usize)> { + let file = File::open("in.pcap")?; + let mut cap = PcapNGReader::new(65535, file)?; + let mut i = 0; + let mut size = 0; + let mut data = vec![]; + loop { + match cap.next() { + Ok((offset, packet)) => { + i += 1; + size += offset; + match packet { + PcapBlockOwned::Legacy(_block) => { + println!("dunno"); + } + PcapBlockOwned::LegacyHeader(_block) => { + println!("dunnoheader"); + } + PcapBlockOwned::NG(block) => { + if block.is_data_block() { + match block { + Block::EnhancedPacket(ref epb) => { + if epb.origlen == 6972 { + // remove udp header + data.push(epb.data[0x2a..].to_vec()); + } + } + Block::SimplePacket(ref ep) => { + if ep.origlen == 6972 { + println!("found one regular"); + } + } + _ => { + println!("unsupported packet"); + } + } + } + } + } + cap.consume(offset) + } + Err(PcapError::Eof) => break, + Err(PcapError::Incomplete) => { + cap.refill().unwrap(); + } + Err(_) => { + println!("unexpected error"); + break; + } + } + } + Ok((data, i, size)) +} + +#[repr(C, packed)] +#[derive(Clone, Debug)] +struct Header { + c1: u32, + c2: u16, + part: u16, + a: u16, + ffaa: u16, + b: u16, + c: u16, + d: u16, +} + +impl Header { + fn read(data: &[u8]) -> anyhow::Result { + Ok(Header { + c1: u32::from_be_bytes([data[0], data[1], data[2], data[3]]), + c2: u16::from_be_bytes([data[4], data[5]]), + part: u16::from_be_bytes([data[6], data[7]]), + a: u16::from_be_bytes([data[8], data[9]]), + ffaa: u16::from_be_bytes([data[10], data[11]]), + b: u16::from_be_bytes([data[12], data[13]]), + c: u16::from_be_bytes([data[14], data[15]]), + d: u16::from_be_bytes([data[16], data[17]]), + }) + } + + #[allow(dead_code)] + fn read_via_cast(data: &[u8]) -> anyhow::Result<&Self> { + let size = std::mem::size_of::(); + if data.len() < size { + return Err(anyhow::anyhow!("not large enough")); + } + let mem = &data[..size].as_ptr(); + Ok(unsafe { &*mem.cast() }) + } +} + +const HDR_SIZE: usize = std::mem::size_of::
(); + +fn main() -> anyhow::Result<()> { + decode_raw("dump.bin")?; +} + +fn decode(filename: &str) -> PyResult>> { + let res = decode_raw(filename)?; + Ok(res.into()) +} + +fn decode_raw(filename: &str) -> anyhow::Result>> { + let (data, i, size) = get_packets_without_udp_header()?; + + println!("found {} packets, saved {}, {} size", i, data.len(), size); + let mut dump = File::create_new("dump.bin"); + 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); + } + 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) { + let name = format!("{:03}.png", 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) { + let name = format!("temp_{:03}.png", 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 p = frame.as_ptr(); + let samples: &[u16] = unsafe { std::slice::from_raw_parts(p.cast(), frame.len() / 2) }; + framesu16.push(samples.clone()); + let frame = samples + .iter() + .copied() + .map(|x| { + let x: f64 = x.into(); + ((-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(framesu16) +} + +#[pymodule] +fn thermaldecoder(_py: Python<'_>, m: &PyModule) -> PyResult<()> { + m.add_function(wrap_pyfunction!(decode, m)?)?; + Ok(()) +}