thermalcam_decoder/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::*;

struct PacketsIterator {
    cap: PcapNGReader<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 cap = PcapNGReader::new(65535, file)?;
        Ok(PacketsIterator { cap, i: 0, size: 0 })
    }
}

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

struct Frame {
    header: Header,
    pixels: Vec<u16>,
}

struct Decoder {}

impl Decoder {
    fn new(filename: &str) {}
}

impl Iterator for Decoder {
    type Item = Frame;

    fn next(&mut self) -> Option<Self::Item> {
        todo!()
    }
}

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);

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