rust version iterator based, 3.5 MB memory turning 5 GB frames to files on disk, 10MiB/second

.gitignore

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

Cargo.toml

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

examples/main.rs

@@ -0,0 +1,12 @@
+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

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

thermaldecoder/src/lib.rs

@@ -1,190 +0,0 @@
-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

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