thermalcam_decoder/src/lib.rs

use indicatif::ProgressBar;
use indicatif::ProgressBarIter;
use png;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use std::fs::{File, OpenOptions};
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 = OpenOptions::new()
            .read(true)
            .write(false)
            .create(false)
            .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 {
    #[allow(dead_code)]
    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())
}

/// writes to frames/<basename of filename>
pub fn decode_to_files(filename: &str) -> anyhow::Result<()> {
    let frameiter = Decoder::new(filename)?;
    let basename = std::path::Path::new(filename)
        .file_stem()
        .ok_or(anyhow::anyhow!("cannot get basename"))?
        .to_str()
        .ok_or(anyhow::anyhow!("cannot convert to utf-8 from os name"))?;
    let target_dir = format!("frames/{}", basename);
    let target_dir = std::path::Path::new(&target_dir);
    if !target_dir.exists() {
        std::fs::create_dir(target_dir)?;
    }
    for (i, frame) in frameiter.enumerate() {
        let name = format!("frames/{}/{:05}.png", basename, i);
        if let Err(_e) = write_raw_frame(&name, &frame.raw) {
            println!("skipping bad frame {}", i);
            continue;
        }
        let name = format!("{}/temp_{:05}.png", target_dir.display(), 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(())
}