io.rs

7.9 kB · rust · 242 lines

1use crate::cell::Cell;2use crate::colors::Color;3use crate::errors::{value_error, MrlyError, Result};4use crate::ramp::Colorizer;5use std::collections::HashMap;6use std::fs;7use std::path::Path;89pub use crate::codec::png;1011/// Maps a slice of type bytes through a colorizer into rgba pixels.12pub fn colorize(types: &[u8], colorizer: &Colorizer, max: usize) -> Vec<[u8; 4]> {13    let values: Vec<usize> = types.iter().map(|&v| v as usize).collect();14    colorizer.colors(&values, max)15}1617/// Snaps every color in a grid to its nearest palette entry, or an error on an empty palette.18pub fn recolor(grid: &Cell, palette: &[Color]) -> Result<Cell> {19    if palette.is_empty() {20        return value_error("palette must not be empty.");21    }22    let colors = grid23        .colors24        .clone()25        .unwrap_or_else(|| vec![[0, 0, 0, 0]; grid.size()]);26    let mut out = grid.clone();27    out.colors = Some(colors.iter().map(|&c| nearest(c, palette)).collect());28    Ok(out)29}3031fn nearest(c: [u8; 4], palette: &[Color]) -> [u8; 4] {32    let mut best = palette[0];33    let mut best_dist = i32::MAX;34    for &p in palette {35        let dr = p.r as i32 - c[0] as i32;36        let dg = p.g as i32 - c[1] as i32;37        let db = p.b as i32 - c[2] as i32;38        let dist = dr * dr + dg * dg + db * db;39        if dist < best_dist {40            best_dist = dist;41            best = p;42        }43    }44    [best.r, best.g, best.b, c[3]]45}4647/// The visual class of a grid's pixels.48#[derive(Clone, Copy, Debug, PartialEq, Eq)]49pub enum Kind {50    /// Pure black and white.51    Binary,52    /// Equal channels throughout.53    Grayscale,54    /// At least one colored pixel.55    Color,56}5758/// The visual summary of one grid.59#[derive(Clone, Debug, PartialEq)]60pub struct Analysis {61    /// The visual class.62    pub kind: Kind,63    /// The most frequent color.64    pub dominant: Color,65    /// The mean luminance across all pixels.66    pub mean_luminance: f64,67    /// The darkest pixel's luminance.68    pub min_luminance: f64,69    /// The brightest pixel's luminance.70    pub max_luminance: f64,71}7273fn luminance(c: [u8; 4]) -> f64 {74    0.299 * c[0] as f64 + 0.587 * c[1] as f64 + 0.114 * c[2] as f6475}7677fn dominant(colors: &[[u8; 4]]) -> [u8; 4] {78    let mut counts: HashMap<[u8; 4], usize> = HashMap::new();79    for &c in colors {80        *counts.entry(c).or_insert(0) += 1;81    }82    counts83        .into_iter()84        .max_by_key(|&(_, n)| n)85        .map(|(c, _)| c)86        .unwrap_or([0, 0, 0, 0])87}8889fn luminance_stats(colors: &[[u8; 4]]) -> (f64, f64, f64) {90    if colors.is_empty() {91        return (0.0, 0.0, 0.0);92    }93    let values: Vec<f64> = colors.iter().map(|&c| luminance(c)).collect();94    let mean = values.iter().sum::<f64>() / values.len() as f64;95    let min = values.iter().cloned().fold(f64::INFINITY, f64::min);96    let max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);97    (mean, min, max)98}99100/// Classifies a grid's pixels and reports the dominant color and luminance bounds.101pub fn analyze(grid: &Cell) -> Analysis {102    match &grid.colors {103        Some(colors) => {104            let binary = colors105                .iter()106                .all(|&c| (c[0], c[1], c[2]) == (0, 0, 0) || (c[0], c[1], c[2]) == (255, 255, 255));107            let grayscale = !binary && colors.iter().all(|&c| c[0] == c[1] && c[1] == c[2]);108            let kind = if binary {109                Kind::Binary110            } else if grayscale {111                Kind::Grayscale112            } else {113                Kind::Color114            };115            let rgb = dominant(colors);116            let (mean, min, max) = luminance_stats(colors);117            Analysis {118                kind,119                dominant: Color::rgba(rgb[0], rgb[1], rgb[2], rgb[3]),120                mean_luminance: mean,121                min_luminance: min,122                max_luminance: max,123            }124        }125        None => {126            let bytes = grid.types.bytes();127            let binary = bytes.iter().all(|&v| v == 0 || v == 1);128            let kind = if binary {129                Kind::Binary130            } else {131                Kind::Grayscale132            };133            let colors: Vec<[u8; 4]> = bytes.iter().map(|&v| [v, v, v, 255]).collect();134            let rgb = dominant(&colors);135            let (mean, min, max) = luminance_stats(&colors);136            Analysis {137                kind,138                dominant: Color::rgb(rgb[0], rgb[1], rgb[2]),139                mean_luminance: mean,140                min_luminance: min,141                max_luminance: max,142            }143        }144    }145}146147/// Creates the directory and every missing parent, or an error naming the path.148pub fn make(dir: &Path) -> Result<()> {149    fs::create_dir_all(dir).map_err(|error| broke("make", dir, &error))150}151152/// Writes the bytes to the path, or an error naming it.153pub fn write(path: &Path, bytes: &[u8]) -> Result<()> {154    fs::write(path, bytes).map_err(|error| broke("write", path, &error))155}156157/// Names the act, the path and the io error that stopped it.158pub fn broke(act: &str, path: &Path, error: &std::io::Error) -> MrlyError {159    MrlyError::Value(format!("could not {act} {}: {error}", path.display()))160}161162#[cfg(test)]163mod tests {164    use super::*;165    use crate::atoms;166    use crate::colors::{BLACK, BLUE, RED, WHITE};167    #[test]168    fn colorize_maps_types_through_colorizer() {169        let types = [0u8, 1, 0, 1];170        let colorizer = Colorizer::two_tone(WHITE, BLACK);171        let colors = colorize(&types, &colorizer, 1);172        assert_eq!(173            colors,174            vec![175                [255, 255, 255, 255],176                [0, 0, 0, 255],177                [255, 255, 255, 255],178                [0, 0, 0, 255]179            ]180        );181    }182    #[test]183    fn recolor_snaps_to_nearest_palette_entry() {184        let mut grid = Cell::new(atoms::carpet_2d(2));185        grid.colors = Some(vec![186            [10, 10, 10, 255],187            [240, 20, 20, 255],188            [20, 20, 240, 255],189            [250, 250, 250, 255],190        ]);191        let out = recolor(&grid, &[BLACK, RED, BLUE, WHITE]).unwrap();192        let colors = out.colors.unwrap();193        assert_eq!(colors[0], [BLACK.r, BLACK.g, BLACK.b, 255]);194        assert_eq!(colors[1], [RED.r, RED.g, RED.b, 255]);195        assert_eq!(colors[2], [BLUE.r, BLUE.g, BLUE.b, 255]);196        assert_eq!(colors[3], [WHITE.r, WHITE.g, WHITE.b, 255]);197    }198    #[test]199    fn recolor_rejects_empty_palette() {200        let grid = Cell::new(atoms::carpet_2d(2));201        assert!(recolor(&grid, &[]).is_err());202    }203    #[test]204    fn analyze_detects_binary_types_without_colors() {205        let grid = Cell::new(atoms::carpet_2d(3));206        let a = analyze(&grid);207        assert_eq!(a.kind, Kind::Binary);208    }209    #[test]210    fn analyze_detects_grayscale_and_color() {211        let mut gray = Cell::new(atoms::carpet_2d(2));212        gray.colors = Some(vec![[10, 10, 10, 255], [200, 200, 200, 255]]);213        assert_eq!(analyze(&gray).kind, Kind::Grayscale);214        let mut color = Cell::new(atoms::carpet_2d(2));215        color.colors = Some(vec![[10, 200, 30, 255], [200, 10, 30, 255]]);216        assert_eq!(analyze(&color).kind, Kind::Color);217    }218    #[test]219    fn analyze_reports_luminance_bounds() {220        let mut grid = Cell::new(atoms::carpet_2d(2));221        grid.colors = Some(vec![222            [0, 0, 0, 255],223            [255, 255, 255, 255],224            [0, 0, 0, 255],225            [0, 0, 0, 255],226        ]);227        let a = analyze(&grid);228        assert_eq!(a.min_luminance, 0.0);229        assert!((a.max_luminance - 255.0).abs() < f64::EPSILON);230        assert_eq!(a.dominant, Color::rgba(0, 0, 0, 255));231    }232    #[test]233    fn make_and_write_land_on_disk() {234        let dir = std::env::temp_dir().join("mrlycore_io_make_write");235        fs::remove_dir_all(&dir).ok();236        make(&dir).unwrap();237        write(&dir.join("note.txt"), b"honk").unwrap();238        assert_eq!(fs::read(dir.join("note.txt")).unwrap(), b"honk");239        assert!(write(&dir, b"honk").is_err());240        fs::remove_dir_all(&dir).ok();241    }242}