thermalcam_decoder/thermaldecoder/src/lib.rs

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