catprint-rs/src/image.rs

use crate::protocol::*;
use dither::prelude::*;
use std::path::Path;

pub struct Image {
    image: Img<f64>,
    mean: f64,
}

fn rle_bytes(val: u8, mut counter: u32) -> Vec<u8> {
    let mut compressed = vec![];
    if counter > 0 {
        while counter > 127 {
            let code = (val << 7) | 127;
            compressed.push(code);
            counter -= 127;
        }
        let code = (val << 7) | (counter as u8);
        compressed.push(code);
    }
    compressed
}

impl Image {
    pub fn load(
        path: &Path,
        rotate: bool,
    ) -> std::result::Result<Self, Box<dyn std::error::Error>> {
        let image = image::open(path)?;
        let image = if rotate { image.rotate90() } else { image };
        let image = image
            .resize(
                crate::protocol::PIXELS_PER_LINE as u32,
                !0_u32,
                image::imageops::FilterType::CatmullRom,
            )
            .into_rgb8();

        let image: Img<RGB<u8>> = unsafe {
            Img::from_raw_buf(
                image.pixels().map(|p| RGB::from(p.0)).collect(),
                image.width(),
            )
        };

        let image: Img<RGB<f64>> = image.convert_with(|rgb| rgb.convert_with(f64::from));

        let image = image.convert_with(|rgb| rgb.to_chroma_corrected_black_and_white());
        Ok(Self { image, mean: 0.5 })
    }

    pub fn kmean(self) -> Self {
        let mean = self.image.iter().sum::<f64>() / self.image.len() as f64;
        Self {
            image: self.image,
            mean,
        }
    }

    pub fn dither(self, ditherer: &Ditherer) -> Self {
        let image = ditherer.dither(self.image, dither::create_quantize_n_bits_func(1).unwrap());
        Self { image, mean: 0.5 }
    }

    /// Returns a line, bool is true if compressed, false if uncompressed
    pub fn line(
        &self,
        y: u32,
    ) -> Option<(bool, usize, [u8; crate::protocol::PIXELS_PER_LINE / 8])> {
        if y > self.image.height() {
            None
        } else {
            let mut compressed = Vec::<u8>::new();

            let mut counter = 0_u32;
            let mut last_val = 2;
            for x in 0..crate::protocol::PIXELS_PER_LINE {
                let x = x as u32;
                let pixel = self.image.get((x, y)).unwrap();
                let val = if pixel > &self.mean { 0 } else { 1 };
                if val == last_val {
                    counter = counter + 1
                } else if counter > 0 {
                    compressed.extend(rle_bytes(last_val, counter));
                    counter = 1;
                }
                last_val = val;
            }
            compressed.extend(rle_bytes(last_val, counter));

            if compressed.len() > crate::protocol::PIXELS_PER_LINE / 8 {
                let mut data = [0_u8; crate::protocol::PIXELS_PER_LINE / 8];
                for x in 0..crate::protocol::PIXELS_PER_LINE {
                    let x = x as u32;
                    let pixel = self.image.get((x, y)).unwrap();
                    let val = if pixel > &self.mean { 0 } else { 1 };
                    let i = (x / 8) as usize;
                    let j = x % 8;
                    let current = data[i];
                    data[i] = current | (val << j);
                }

                Some((false, crate::protocol::PIXELS_PER_LINE / 8, data))
            } else {
                use std::convert::TryInto;
                let len = compressed.len();
                compressed.resize(crate::protocol::PIXELS_PER_LINE / 8, 0);
                Some((true, len, compressed.try_into().unwrap()))
            }
        }
    }

    pub fn line_count(&self) -> u32 {
        self.image.height()
    }

    pub fn print(&self, mode: DrawingMode, quality: Quality, energy: u16) -> Vec<Command> {
        let mut commands = vec![
            Command::SetQuality(quality),
            Command::SetEnergy(energy),
            Command::SetDrawingMode(mode),
        ];

        commands.push(Command::MagicLattice(LatticeType::Start));
        for y in 0..self.line_count() {
            let (compressed, len, pixels) = self.line(y).unwrap();
            commands.push(Command::Print(compressed, len, pixels));
        }
        commands.push(Command::MagicLattice(LatticeType::End));

        commands
    }
}