thermalcam_decoder/examples/live.rs

use clap::Parser;
use dotenv::dotenv;
use std::io::Write;
use thermaldecoder::{Header, HDR_SIZE};
use v4l::video::Output;

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    #[arg(short, long, default_value_t = false)]
    temperature: bool,
    #[arg(short, long, default_value = "/dev/video0")]
    device: String,
    #[arg(short, long)]
    red_cutoff: Option<f64>,
}

fn pixel_to_celcius(x: u16) -> u16 {
    let x: f64 = x.into();
    let x = x / 256.0;
    let t = (-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);
    (t * 256.0) as u16
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();
    dotenv().ok();
    let device = match std::env::var("THERMALCAM_IFACE=enp1s0f0") {
        Ok(d) => {
            let device = pcap::Device::list()
                .expect("device list failed")
                .into_iter()
                .find(|x| x.name == d)
                .expect(&format!("could not find device {}", d));
            device
        }
        Err(_) => pcap::Device::lookup()
            .expect("device lookup failed")
            .expect("no device available"),
    };
    // get the default Device

    println!("Using device {}", device.name);
    let output = args.device;
    println!("Using output v4l2loopback device {}", output);

    const WIDTH: usize = 288;
    const HEIGHT: usize = 384;
    let greyscale = !args.temperature || args.red_cutoff.is_none();
    let fourcc_repr = if greyscale {
        [
            b'Y', // | 0b10000000
            b'1', b'6',
            b' ', // Note: not using b' ' | 0x80, (V4L2_PIX_FMT_Y16_BE)
                  // because VID_S_FMT ioctl returns EINVAL, so just swap the bytes here
        ]
    } else {
        // RGB32 is 4 bytes R, G, B, A
        [b'R', b'G', b'B', b'4']
    };
    let bytes_per_pixel = if greyscale { 2 } else { 4 };
    let fourcc = v4l::format::FourCC { repr: fourcc_repr };
    let mut out = v4l::Device::with_path(output)?;
    // To find the fourcc code, use v4l2-ctl --list-formats-out /dev/video0
    // (or read the source :)
    // flip axes
    let format = v4l::Format::new(HEIGHT as u32, WIDTH as u32, fourcc);
    Output::set_format(&out, &format)?;

    // Setup Capture
    let mut cap = pcap::Capture::from_device(device)
        .unwrap()
        .immediate_mode(true)
        .open()
        .unwrap();

    // get a packet and print its bytes
    const PACKET_LEN: usize = 6972;
    // input is grayscale 16 bits per pixel
    const FRAME_LEN: usize = WIDTH * HEIGHT * 2;
    let mut frame = [0u8; FRAME_LEN];
    let mut len = 0;
    let output_frame_len = WIDTH * HEIGHT * bytes_per_pixel;
    let mut swapped_vec = vec![0u8; output_frame_len];
    let swapped = &mut swapped_vec;
    while let Ok(p) = cap.next_packet() {
        let data = p.data;
        if data.len() != PACKET_LEN {
            continue;
        }
        let data = &data[0x2a..];
        let header = match Header::read(data) {
            Ok(header) => header,
            Err(_) => continue,
        };
        let data = &data[HDR_SIZE..];
        if (header.part == 0 && len > 0)
            // do not write out of bounds - would panic, instead just skip
            || (data.len() + len > FRAME_LEN)
        {
            if len == FRAME_LEN {
                // swap the bytes, we are using LE, not BE, 16 bit grayscale
                // possibly limitation of current v4l2loopback or v4l rust wrapper or libv4l2
                for i in 0..FRAME_LEN / 2 {
                    let x = i % WIDTH;
                    let y = (i / WIDTH) % HEIGHT;
                    let mut pixel = u16::from_be_bytes([frame[i * 2], frame[i * 2 + 1]]);
                    if greyscale {
                        if args.temperature {
                            pixel = pixel_to_celcius(pixel);
                        }
                        let pixel_swapped = pixel.to_le_bytes();
                        let out_i = ((HEIGHT - 1 - y) + (WIDTH - 1 - x) * HEIGHT) * 2;
                        swapped[out_i..out_i + 2].copy_from_slice(&pixel_swapped);
                    } else {
                        pixel = pixel_to_celcius(pixel);
                        let cutoff = args.red_cutoff.unwrap();
                        let (r, g, b) = if pixel > (256.0 * cutoff) as u16 {
                            let p = pixel - (256.0 * cutoff) as u16;
                            let p = (p / 256).max(127);
                            ((128 + p) as u8, 0, 0)
                        } else {
                            let g = frame[i * 2];
                            let b = frame[i * 2 + 1];
                            (0, g, b)
                        };
                        let out_i = ((HEIGHT - 1 - y) + (WIDTH - 1 - x) * HEIGHT) * 4;
                        swapped[out_i..out_i + 4].copy_from_slice(&[0, r, g, b]);
                    }
                }
                out.write_all(&swapped[..])?;
            }
            len = 0;
        }
        frame[len..len + data.len()].copy_from_slice(data);
        len += data.len();
    }
    Ok(())
}