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

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(())
 }

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<File>,
+    cap: PcapNGReader<ProgressBarIter<File>>,
     i: usize,
     size: usize,
 }
@@ -73,14 +73,16 @@ impl Iterator for PacketsIterator {
 impl PacketsIterator {
     fn new(filename: &str) -> anyhow::Result<Self> {
         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::<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())
-}
-
-struct Frame {
+pub struct Frame {
     header: Header,
-    pixels: Vec<u16>,
+    raw: Vec<u8>,
 }
 
-struct Decoder {}
+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 {
-    fn new(filename: &str) {}
+    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> {
-        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<Vec<Vec<u16>>> {
-    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::<Vec<u8>>();
+    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<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);
-        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<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)?;
+        let pixels = frame.pixels();
+        write_calibrated_frame(&name, &pixels)?;
     }
-
-    Ok(framesu16)
+    Ok(())
 }
 
 #[pymodule]

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)