.gitignore

@@ -1,4 +1,4 @@
 /venv
 /frames
 **/*.png
-target
+thermaldecoder/target

examples/main.rs

@@ -1,12 +0,0 @@
-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_to_files(&filename, &frames)?;
-    Ok(())
-}

src/lib.rs

@@ -1,268 +0,0 @@
-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::path::Path;
-
-use pcap_parser::traits::PcapReaderIterator;
-use pcap_parser::*;
-
-struct PacketsIterator {
-    cap: PcapNGReader<ProgressBarIter<File>>,
-    i: usize,
-    size: usize,
-}
-
-impl Iterator for PacketsIterator {
-    type Item = Vec<u8>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let mut ret = None;
-        while ret.is_none() {
-            match self.cap.next() {
-                Ok((offset, packet)) => {
-                    self.i += 1;
-                    self.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
-                                            ret = Some(epb.data[0x2a..].to_vec());
-                                        }
-                                    }
-                                    Block::SimplePacket(ref ep) => {
-                                        if ep.origlen == 6972 {
-                                            println!("found one regular");
-                                        }
-                                    }
-                                    _ => {
-                                        println!("unsupported packet");
-                                    }
-                                }
-                            }
-                        }
-                    }
-                    self.cap.consume(offset)
-                }
-                Err(PcapError::Eof) => break,
-                Err(PcapError::Incomplete) => {
-                    self.cap.refill().unwrap();
-                }
-                Err(_) => {
-                    println!("unexpected error");
-                    break;
-                }
-            }
-        }
-        ret
-    }
-}
-
-impl PacketsIterator {
-    fn new(filename: &str) -> anyhow::Result<Self> {
-        let file = File::open(filename)?;
-        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)]
-pub 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<Self> {
-        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::<Self>();
-        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::<Header>();
-
-pub struct Frame {
-    header: Header,
-    raw: Vec<u8>,
-}
-
-impl Frame {
-    fn pixels(&self) -> Vec<u16> {
-        (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<Vec<u8>>,
-}
-
-impl Decoder {
-    pub fn new(filename: &str) -> anyhow::Result<Self> {
-        let packiter = PacketsIterator::new(filename)?;
-        Ok(Decoder {
-            packiter,
-            parts: vec![],
-        })
-    }
-}
-
-impl Iterator for Decoder {
-    type Item = Frame;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        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
-    }
-}
-
-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(())
-}
-
-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::<Vec<u8>>();
-    writer.write_image_data(&frame)?;
-    Ok(())
-}
-
-#[pyclass]
-struct PyFrameIterator {
-    iter: Decoder,
-}
-
-#[pymethods]
-impl PyFrameIterator {
-    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
-        slf
-    }
-    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<PyObject> {
-        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<PyFrameIterator> {
-    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);
-        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 pixels = frame.pixels();
-        write_calibrated_frame(&name, &pixels)?;
-    }
-    Ok(())
-}
-
-#[pymodule]
-fn thermaldecoder(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
-    m.add_function(wrap_pyfunction!(decode, m)?)?;
-    Ok(())
-}

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 = list(decode('in.pcap'))
+frames = decode('in.pcap', 'frames')
 f = np.array(frames[0])
 f.shape = (384, 288)
 plt.imshow(f)

thermaldecoder/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 
 [lib]
 name = "thermaldecoder"
-crate-type = ["rlib", "cdylib"]
+crate-type = ["cdylib"]
 
 [dependencies]
 anyhow = "1.0.77"

thermaldecoder/src/lib.rs

@@ -0,0 +1,190 @@
+use indicatif::{ProgressBar, ProgressIterator};
+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::*;
+
+fn get_packets_without_udp_header(filename: &str) -> anyhow::Result<(Vec<Vec<u8>>, usize, usize)> {
+    let file = File::open(filename)?;
+    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<Self> {
+        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::<Self>();
+        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::<Header>();
+
+#[pyfunction]
+fn decode(filename: &str, frames_root: &str) -> PyResult<Vec<Vec<u16>>> {
+    let res =
+        decode_raw(filename, frames_root).map_err(|_| PyValueError::new_err("failed to read"))?;
+    Ok(res.into())
+}
+
+fn decode_raw(filename: &str, frames_root: &str) -> anyhow::Result<Vec<Vec<u16>>> {
+    let (data, i, size) = get_packets_without_udp_header(filename)?;
+
+    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);
+    }
+    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!("{}/{: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) {
+        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<u16> = (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::<Vec<u8>>();
+        writer.write_image_data(&frame)?;
+    }
+
+    Ok(framesu16)
+}
+
+#[pymodule]
+fn thermaldecoder(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(decode, m)?)?;
+    Ok(())
+}

thermaldecoder/src/main.rs

@@ -0,0 +1,4 @@
+//fn main() -> anyhow::Result<()> {
+//    decode_raw("dump.bin")?;
+//    Ok(())
+//}