typographer/

lib.rs

//! Advanced ASCII art generator. Key features:
//! - Converts pixel luminance to ASCII characters.
//! - Uses smart Canny edge-detection to improve feature readability.
//! - Has customizable colour modes. 
//! - Compensates for the character aspect-ratio. 
//! - Highly configurable. 
//!
//! Also available as a CLI on [GitHub](https://github.com/user-simon/typographer). 
//! 
//! # Basic Usage
//!
//! The [`Typographer`] algorithm is composed of three sub-algorithms:
//! - [`Fill`]: computes fill symbols from the gradient using pixel lumas. 
//! - [`Edges`]: computes symbols along the feature edges of the image. 
//! - [`Colours`]: computes symbol colours using one of several modes. 
//!
//! The algorithm is then ran with [`Typographer::render`], which returns an [`AsciiImage`] represented as a
//! 2D array of (possibly) styled symbols. 
//! 
//! Run with default settings: 
//! ```no_run
//! use typographer::{Typographer, Size, RgbImage};
//!
//! let typographer = Typographer::new_coloured();
//! // let img: RgbImage
//! # let img = RgbImage::default();
//! let ascii: AsciiImage = typographer.render(img, Size::Width(100));
//!
//! for symbol in ascii.iter() {
//!     print!(symbol);
//! }
//! ```
//!
//! Run with custom settings:
//! ```no_run
//! use typographer::*;
//!
//! let fill = Fill {
//!     gradient: "   .oO",
//!     ..Fill::default()
//! };
//! let edges = Edges {
//!     angles: ['|', '/', '-', '\\'], 
//!     sigma: 0.2,
//!     ..Edges::default()
//! };
//! let typographer = Typographer {
//!     luma_threshold: 0.4, 
//!     symbol_aspect_ratio: 2.2, 
//!     fill: Some(fill), 
//!     edges: Some(edges), 
//!     colours: None, // disable colours
//!     ..Typographer::default()
//! };
//! 
//! # let img = RgbImage::default();
//! // let img: RgbImage
//! let ascii: AsciiImage = typographer.render(img, Size::Stretched(100, 200));
//!
//! for symbol in ascii.iter() {
//!     print!(symbol);
//! }
//! ```

use std::{cell::LazyCell};
use crossterm::style::{StyledContent, Stylize};
use image::{imageops, Luma, Pixel};
use itertools::Itertools;
use palette::{Hsv, IntoColor};

/// A "pixel" of the ASCII image. This has an aspect ratio defined by [`Typographer::symbol_aspect_ratio`]. 
pub type Styled = StyledContent<char>;

/// Input image with RGB pixels. 
pub type RgbImage = image::RgbImage;

/// Algorithm used to downsample images. 
pub type Sampler = imageops::FilterType;

/// Typographer entry-point.
///
/// By default, this always produces an empty image. To add output, define [`Typographer::fill`] and/or
/// [`Typographer::edges`]. For full configurability, this could be done with the struct update syntax:
/// 
/// ```
/// let typographer = Typographer {
///     char_aspect_ratio: 1.5, 
///     fill: Some(Fill {
///         gradient: " `.:-=;*?+#%@".toString(), 
///         brightness_gain: 0.1, 
///         ..Default::default()
///     }), 
///     edges: Some(Edges::default()), 
///     ..Default::default()
/// };
/// ```
///
/// ... or using the using the [`Typographer::new`] and [`Typographer::new_coloured`] short-hands.
///
/// Images are then rendered using [`Typographer::render`], provided an input image and output size in
/// symbols. 
#[derive(Clone)]
pub struct Typographer {
    /// Algorithm used to downsample the image. 
    pub sampler: Sampler, 
    /// The aspect ratio `height / width` of the symbols. This is usually around `2.0`. 
    pub symbol_aspect_ratio: f32, 
    /// Luma threshold to reduce noise. Defined in `[0, 1]`. 
    pub luma_threshold: f32, 
    /// Settings used to render luma. 
    pub fill: Option<Fill>,
    /// Settings used to render edges. 
    pub edges: Option<Edges>,
    /// Settings used to colourise the symbols. 
    pub colours: Option<Colours>, 
}

impl Default for Typographer {
    fn default() -> Self {
        Typographer {
            sampler: Sampler::Nearest, 
            symbol_aspect_ratio: 2.0, 
            luma_threshold: 0.05, 
            fill: Some(Fill::default()),
            edges: Some(Edges::default()), 
            colours: None, 
        }
    }
}

impl Typographer {
    /// Constructs with only fill and edges enabled. Each have default settings. 
    pub fn new() -> Typographer {
        Typographer {
            fill: Some(Fill::default()),
            edges: Some(Edges::default()), 
            colours: None, 
            ..Default::default()
        }
    }

    /// Constructs with fill, edges, and colours enabled. Each have default settings. 
    pub fn new_coloured() -> Typographer {
        Typographer {
            fill: Some(Fill::default()),
            edges: Some(Edges::default()), 
            colours: Some(Colours::default()), 
            ..Default::default()
        }
    }
    
    /// Renders an RGB image to an ASCII image of given size in symbols. 
    pub fn render(&self, image: &RgbImage, size: Size) -> AsciiImage {
        // get aspect ratio preserving output width and height. this also corrects for the symbol
        // aspect ratio by stretching such that the image is unstretched when the symbols are rendered in e.g.
        // the terminal
        let (width, height) = {
            let aspect_ratio = image.height() as f32 / image.width() as f32;
            let scale = |a, b| ((a as f32) * b) as u32;

            match size {
                Size::Width(w) => (w, scale(w, aspect_ratio / self.symbol_aspect_ratio)), 
                Size::Height(h) => (scale(h, self.symbol_aspect_ratio / aspect_ratio), h), 
                Size::Stretched(w, h) => (w, h), 
            }
        };

        // downscale image and apply the threshold
        let image = {
            // we don't want to perform the expensive edge-detection on the full-sized image, so we downscale
            // it to a multiple of the output size, clamping to 8 multiples (yielding 64:1 input:output
            // pixels)
            let clamped_multiples = |a, b| a * u32::clamp(b / a, 1, 8);
            let width = clamped_multiples(width, image.width());
            let height = clamped_multiples(height, image.height());
            let mut image = imageops::resize(image, width, height, self.sampler);

            for rgb in image.pixels_mut() {
                const STRENGTH: i32 = 4;
                *rgb = transform_hsv(rgb, |mut hsv| {
                    hsv.value = match hsv.value < self.luma_threshold {
                        true => hsv.value.powi(STRENGTH) / self.luma_threshold.powi(STRENGTH - 1),
                        false => hsv.value,
                    };
                    hsv
                });
            }
            image
        };

        let output_sized = LazyCell::new(|| imageops::resize(&image, width, height, self.sampler));

        let mut ascii = AsciiImage::new(width, height);

        if let Some(fill) = &self.fill {
            ascii = fill.apply(LazyCell::force(&output_sized), ascii);
        }
        if let Some(edges) = &self.edges {
            ascii = edges.apply(&image, ascii, self.symbol_aspect_ratio); 
        }
        if let Some(colour) = &self.colours {
            ascii = colour.apply(LazyCell::force(&output_sized), ascii);
        }
        ascii
    }
}

/// Represents a rendered ASCII image. 
///
/// Consists of a matrix of [`Styled`]. 
#[derive(Clone)]
pub struct AsciiImage {
    symbols: Vec<Styled>, 
    width: u32, 
    height: u32, 
}

impl AsciiImage {
    fn new(width: u32, height: u32) -> AsciiImage {
        AsciiImage {
            symbols: vec![' '.stylize(); (width * height) as usize], 
            width, 
            height, 
        }
    }
    
    /// Applies a generic pass over each symbol in the image. 
    fn with_pass<T>(mut self, iter: impl Iterator<Item = T>, mut apply: impl FnMut(Styled, T) -> Styled) -> Self {
        for (symbol, x) in std::iter::zip(self.symbols.iter_mut(), iter) {
            *symbol = apply(*symbol, x);
        }
        self
    }

    /// Returns an iterator of image rows. 
    pub fn iter_rows(&self) -> impl Iterator<Item = &[Styled]> {
        self.symbols.chunks(self.width as usize)
    }

    /// Returns a flat iterator of symbols for printing. Each row is appended a newline. 
    pub fn iter(&self) -> impl Iterator<Item = Styled> {
        self.iter_rows()
            .map(|row| row.iter()
                .cloned()
                .chain(std::iter::once('\n'.stylize()))
            )
            .flatten()
    }

    /// Returns a flat iterator of symbols for printing, with discarded colour information. Each row is
    /// appended a newline. 
    pub fn iter_unstyled(&self) -> impl Iterator<Item = char> {
        self.iter().map(|symbol| symbol.content().clone())
    }

    /// Returns a flat iterator of symbols. 
    pub fn iter_symbols(&self) -> impl Iterator<Item = Styled> {
        self.symbols.iter().cloned()
    }

    /// Returns a flat iterator of symbols, with discarded colour information. 
    pub fn iter_symbols_unstyled(&self) -> impl Iterator<Item = char> {
        self.symbols.iter().map(Styled::content).cloned()
    }

    /// Returns the image width. 
    pub fn width(&self) -> u32 {
        self.width
    }

    /// Returns the image height. 
    pub fn height(&self) -> u32 {
        self.height
    }
    
    /// Returns the internal flat buffer used to store the symbols. 
    pub fn into_inner(self) -> Vec<Styled> {
        self.symbols
    }
}

/// Specifies the width/height of the output image, in symbols. 
#[derive(Clone, Copy)]
pub enum Size {
    /// Specifies ASCII image width. The height is computed to preserve aspect ratio. 
    Width(u32), 
    /// Specifies ASCII image height. The width is computed to preserve aspect ratio. 
    Height(u32), 
    /// Specifies ASCII image width and height. Note that the output image may have a stretched aspect ratio. 
    Stretched(u32, u32), 
}

/// Settings used when computing fill symbols from lumas. 
#[derive(Clone)]
pub struct Fill {
    /// Symbols mapped from low to high luma. 
    pub gradient: String, 
    /// Reverses the order of the gradient for light backgrounds. 
    pub light_mode: bool, 
}

impl Default for Fill {
    fn default() -> Fill {
        Fill {
            gradient: " .:coPO?@■".to_string(), 
            light_mode: false, 
        }
    }
}

impl Fill {
    fn apply(&self, input: &RgbImage, ascii: AsciiImage) -> AsciiImage {
        debug_assert!(input.width() == ascii.width && input.height() == ascii.height);

        if self.gradient.is_empty() {
            return ascii
        }
        
        let gradient: Vec<char> = if self.light_mode {
            self.gradient.chars().rev().collect()
        } else {
            self.gradient.chars().collect()
        };
        let iter = input
            .pixels()
            .map(|rgb| {
                let Luma([luma]) = rgb.to_luma();
                let luma = luma as f32 / 256.0;
                let index = gradient.len() as f32 * luma;
                let symbol = gradient[index as usize].stylize();
                symbol
            });
        ascii.with_pass(iter, |_, symbol| symbol)
    }
}

/// Settings used when computing edge symbols using edge detection. The default settings work well enough
/// in most cases, and are exposed for advanced usage only. 
#[derive(Clone, Copy)]
pub struct Edges {
    /// Angle symbols given in the order `| / — \`. 
    pub angles: [char; 4], 
    /// Blur size in pixels, used to reduce noise. 
    pub sigma: f32,
    /// How large a change in luma must be to register as an edge. Computed from (fallible) image analysis if
    /// not given, which works best if the image has a clear background and foreground. Defined in [0, 1]. 
    pub threshold: Option<f32>, 
    /// The maximum fraction of opposing angles allowed when downsampling eges. E.g. if `|` is the most
    /// common angle with population `p`, then the opposing angle `—` can have a population of no more than
    /// `unanimity * p` for the edge to count. This serves to reduce noise. 
    pub unanimity: f32, 
    /// Sensitivity of the edge filtering performed to reduce noise. Defined in [0, ..]. 
    pub sensitivity: f32, 
}

impl Default for Edges {
    fn default() -> Self {
        Edges {
            angles: ['|', '/', '─', '\\'], 
            sigma: 1.0, 
            threshold: None, 
            unanimity: 0.95, 
            sensitivity: 1.0, 
        }
    }
}

impl Edges {
    fn apply(&self, input: &RgbImage, ascii: AsciiImage, symbol_aspect_ratio: f32) -> AsciiImage {
        let sigma_scale = u32::min(input.width(), input.height()) as f32 / 300.0;
        let sigma = self.sigma * sigma_scale;

        let edges = {
            let input = imageops::grayscale(input);
            let (weak_threshold, strong_threshold) = smart_thresholds(&input);

            edge_detection::canny(
                input,
                sigma, 
                strong_threshold, 
                weak_threshold, 
            )
        };

        let window_w = input.width() / ascii.width;
        let window_h = input.height() / ascii.height;

        // divide the image into windows and get the edge angles in each window
        let windows = Itertools::cartesian_product(0..ascii.height, 0..ascii.width)
            .map(|(y, x)| (y * window_h, x * window_w))
            .map(|(y, x)| Itertools::cartesian_product(0..window_h, 0..window_w)
                .map(move |(dy, dx)| (y + dy, x + dx))
                .map(|(py, px)| edges.interpolate(px as f32, py as f32))
            );

        // downscale by computing a histogram for each window and choosing the most common angle kind among
        // { | / — \ }. this preserves fine angle details compared to e.g. gaussian downsampling, which would
        // average non-angles with angles, producing all non-angles
        let window_angles = windows
            .map(|window| window.map(|edge| (edge.magnitude() != 0.0)
                .then(|| {
                    let (dx, dy) = edge.dir_norm();
                    f32::atan2(dy, symbol_aspect_ratio * dx)
                })
                .map(|angle| angle / std::f32::consts::PI * 0.5 + 0.5) // remap [-π, π] to [0, 1]
                .map(|angle| (angle * 8.0).round() as u8) // quantise to 0..8
            ));

        let sum_threshold = match self.sensitivity {
            0.0 => usize::MAX,
            _ => {
                let scale = u32::min(window_w, window_h);
                usize::max(1, (scale as f32 / self.sensitivity) as usize)
            }
        };
        let downsampled = window_angles
            .map(|window| window
                .flatten()
                .map(|angle| angle % 4)
            )
            .map(histogram::<4>)
            .map(|histogram| match histogram.iter().sum::<usize>() >= sum_threshold {
                true => histogram.iter().position_max().and_then(|max_edge| {
                    let max = histogram[max_edge] as f32;
                    let opposite = histogram[(max_edge + 2) % 4] as f32;
                    let unanimity = (max - opposite) / max;
                    (unanimity >= self.unanimity).then_some(max_edge)
                }),
                false => None, 
            });

        // convert each angle kind to a symbol
        let iter = downsampled.map(|angle| angle
            .map(|angle| self.angles[angle])
            .map(Stylize::stylize)
        );
        ascii.with_pass(iter, |prev, symbol| symbol.unwrap_or(prev))
    }
}

/// Settings used when computing symbol colours. 
#[derive(Clone, Copy, PartialEq)]
pub enum Colours {
    /// Render only hue and saturation, with the luma encoded by the gradient symbol. 
    Chromatic, 
    /// Render only luma. 
    Greyscale, 
    /// Render colours with RGB channels each quantised into the given number of buckets. 
    Retro(u8), 
    /// Render full colours. 
    Full, 
}

impl Default for Colours {
    fn default() -> Self {
        Colours::Chromatic
    }
}

impl Colours {
    fn apply(&self, input: &RgbImage, ascii: AsciiImage) -> AsciiImage {
        debug_assert!(input.width() == ascii.width && input.height() == ascii.height);

        let iter = input
            .pixels()
            .map(|&rgb| match *self {
                Colours::Chromatic => transform_hsv(&rgb, |hsv| {
                    // models perceived saturation close to absolute black
                    let saturation = f32::min(hsv.saturation, 2.0 * hsv.value);
                    let value = 1.0;
                    Hsv::new(hsv.hue, saturation, value)
                }), 
                Colours::Greyscale => {
                    let image::Luma([luma]) = rgb.to_luma();
                    image::Rgb([luma, luma, luma])
                }, 
                Colours::Retro(buckets) => {
                    debug_assert!(buckets >= 1);

                    let bucket_size = u8::MAX / buckets;
                    let quantise = |x| bucket_size * (x / bucket_size);
                    let image::Rgb(rgb) = rgb;
                    image::Rgb(rgb.map(quantise))
                }, 
                Colours::Full => rgb, 
            })
            .map(|image::Rgb([r, g, b])| crossterm::style::Color::Rgb{ r, g, b });
        ascii.with_pass(iter, StyledContent::with)
    }
}

/// Computes a fixed-size histogram of a sequence of integers. 
fn histogram<const N: usize>(data: impl IntoIterator<Item = u8>) -> [usize; N] {
    data.into_iter().fold([0; N], |mut acc, x| {
        acc[x as usize] += 1;
        acc
    })
}

/// Finds the best thresholds using Otsu's optimization-based thresholding algorithm. 
fn smart_thresholds(image: &image::GrayImage) -> (f32, f32){
    const VARIANCE_THRESHOLD: f32 = 4e8;

    let histogram = histogram::<256>(image.pixels().map(|luma| luma[0]));
    let histogram_iter = || histogram.iter()
        .enumerate()
        .map(|(luma, &bin_size)| (luma as f32 / 255.0, bin_size as f32));
    let luma_sum = histogram_iter()
        .map(|(luma, bin_size)| bin_size * luma)
        .sum::<f32>();
    let pixel_count = (image.width() * image.height()) as f32;
    let fallback_thresholds = || {
        let mean_luma = luma_sum / pixel_count;
        let strong_threshold = mean_luma * 0.4;
        let weak_threshold = strong_threshold * 0.2;
        eprintln!("Using fallback thresholds");
        (weak_threshold, strong_threshold)
    };

    let (min, max) = {
        let non_zero = |&x| x > 0;
        let min = histogram.iter().position(non_zero);
        let max = histogram.iter().rposition(non_zero);

        match Option::zip(min, max) {
            Some((min, max)) => (min + 1, max - 1),
            None => return fallback_thresholds(),
        }
    };
    // TODO investigate if we can do a 2D search to optimize both the lower and upper threshold by maximizing
    // the inter-class variance between 3 classes
    histogram_iter()
        // compute cumulative sum of pixel count and luma values
        .scan((0.0, 0.0), |(pixel_sum, luma_sum), (luma, bin_size)| {
            *pixel_sum += bin_size;
            *luma_sum += bin_size * luma;
            Some((*pixel_sum, *luma_sum, luma))
        })
        // ignore lumas outside the dynamic range of the image
        .skip(min)
        .take(max - min)
        // compute the intra-class variance of the luma threshold
        .map(|(pixel_cumsum, luma_cumsum, threshold)| {
            let lo_pixel_cumsum = pixel_cumsum;
            let hi_pixel_cumsum = pixel_count - lo_pixel_cumsum;
            let lo_luma_cumsum = luma_cumsum;
            let hi_luma_cumsum = luma_sum - lo_luma_cumsum;

            let lo_mean = lo_luma_cumsum / lo_pixel_cumsum;
            let hi_mean = hi_luma_cumsum / hi_pixel_cumsum;
            let mean_diff = lo_mean - hi_mean;

            let intra_class_variance = lo_pixel_cumsum * hi_pixel_cumsum * mean_diff * mean_diff;

            (intra_class_variance, threshold)
        })
        .max_by(|x, y| PartialOrd::partial_cmp(x, y).unwrap())
        .and_then(|(variance, threshold)| match variance >= VARIANCE_THRESHOLD {
            true => Some((0.5 * threshold, threshold)),
            false => None, 
        })
        .unwrap_or_else(fallback_thresholds)
}

/// Applies a function over the HSV representation of an RGB colour. 
fn transform_hsv(rgb: &image::Rgb<u8>, function: impl Fn(palette::Hsv) -> palette::Hsv) -> image::Rgb<u8> {
    // into HSV
    let [r, g, b] = rgb.0.map(|x| x as f32 / (u8::MAX as f32));
    let hsv = palette::Srgb::new(r, g, b).into_color();

    // apply transformation
    let palette::Hsv{ hue, saturation, value, .. } = function(hsv);
    let hsv = palette::Hsv::new(hue, f32::clamp(saturation, 0.0, 1.0), f32::clamp(value, 0.0, 1.0));

    // back into SRGB
    let palette::Srgb{ red, green, blue, .. } = hsv.into_color();
    let [r, g, b] = [red, green, blue].map(|x| (x * u8::MAX as f32) as u8);
    image::Rgb([r, g, b])
}