color_blinder/

lib.rs

pub mod gimp;
pub mod web_algorithms;

#[cfg(feature = "labels")]
pub mod labels;

use bitflags::bitflags;
use cfg_if::cfg_if;
use image::{GenericImage, ImageBuffer, Rgba, RgbaImage};
use std::iter::{IntoIterator, Iterator};
use std::path::PathBuf;
use std::sync::{
    mpsc::{channel, Receiver, Sender},
    Arc,
};
use threadpool::ThreadPool;
use web_algorithms::{anomylize, blindMK, monochrome};

pub type RgbaBuf = ImageBuffer<Rgba<u8>, Vec<u8>>;
pub type NamedBuf = (&'static str, RgbaBuf);
pub type FilterFunc = Box<dyn Fn(&Rgba<u8>) -> Rgba<u8> + Send + Sync>;

pub const COMBINED_MARGIN: u32 = 25;
const ORIGINAL_TEXT: &'static str = "Original";

#[derive(Clone)]
pub enum RevBlind {
    Protan,
    Deutan,
    Tritan,
}

bitflags! {
    #[allow(non_camal_case)]
    pub struct FilterKind: usize {
        const ORIGINAL          = 0b0000_0000_0001;
        const ACHROMATOMALY     = 0b0000_0000_0010;
        const ACHROMATOPSIA     = 0b0000_0000_0100;
        const DEUTERANOMALY     = 0b0000_0000_1000;
        const DEUTERANOPIA      = 0b0000_0001_0000;
        const DEUTERANOPIABVM97 = 0b0000_0010_0000;
        const PROTANOMALY       = 0b0000_0100_0000;
        const PROTANOPIA        = 0b0000_1000_0000;
        const PROTANOPIABVM97   = 0b0001_0000_0000;
        const TRITANOMALY       = 0b0010_0000_0000;
        const TRITANOPIA        = 0b0100_0000_0000;
        const TRITANOPIABVM97   = 0b1000_0000_0000;

        const MONOCHROME = Self::ACHROMATOMALY.bits | Self::ACHROMATOPSIA.bits;
        const DEUTAN = Self::DEUTERANOPIA.bits | Self::DEUTERANOPIABVM97.bits;
        const PROTAN = Self::PROTANOMALY.bits | Self::PROTANOPIA.bits | Self::PROTANOPIABVM97.bits;
        const TRITAN = Self::TRITANOMALY.bits | Self::TRITANOPIA.bits | Self::TRITANOPIABVM97.bits;
        const BVM97 = Self::DEUTERANOPIABVM97.bits | Self::PROTANOPIABVM97.bits | Self::TRITANOPIABVM97.bits;
        // All flags but the ORIGINAL
        const ALL_FILTERS       = 0b1111_1111_1110;
    }
}

impl FilterKind {
    pub fn to_str(&self) -> &'static str {
        match *self {
            FilterKind::ORIGINAL => "Original",
            FilterKind::ACHROMATOMALY => "Achromatomaly",
            FilterKind::ACHROMATOPSIA => "Achromatopsia",
            FilterKind::DEUTERANOMALY => "Deuteranomaly",
            FilterKind::DEUTERANOPIA => "Deuteranopia",
            FilterKind::DEUTERANOPIABVM97 => "DeuteranopiaBVM97",
            FilterKind::PROTANOMALY => "Protanomaly",
            FilterKind::PROTANOPIA => "Protanopia",
            FilterKind::PROTANOPIABVM97 => "ProtanopiaBVM97",
            FilterKind::TRITANOMALY => "Tritanomaly",
            FilterKind::TRITANOPIA => "Tritanopia",
            FilterKind::TRITANOPIABVM97 => "TritanopiaBVM97",
            // Groups
            FilterKind::MONOCHROME => "MONOCHROME",
            FilterKind::DEUTAN => "DEUTAN",
            FilterKind::PROTAN => "PROTAN",
            FilterKind::TRITAN => "TRITAN",
            FilterKind::BVM97 => "BVM97",
            FilterKind::ALL_FILTERS => "ALL_FILTERS",
            unknown => panic!("FilterKind::to_str(0b{:b})", unknown.bits),
        }
    }

    pub fn from_strs(filters: &[&str]) -> Result<FilterKind, FilterKindParseError> {
        let mut kind = FilterKind::empty();
        let mut unused_filters = vec![];

        for filter in filters {
            let filter = filter.trim();
            let matches = FilterKind::all()
                .iter()
                .filter(|name| name.to_str().eq_ignore_ascii_case(filter))
                .fold(FilterKind::empty(), |b, c| b | c);

            if matches.bits().count_ones() > 0 {
                kind |= matches;
            } else {
                let mut used = false;

                if "MONOCHROME".eq_ignore_ascii_case(filter) {
                    kind |= FilterKind::MONOCHROME;
                    used = true;
                }
                if "DEUTAN".eq_ignore_ascii_case(filter) {
                    kind |= FilterKind::DEUTAN;
                    used = true;
                }
                if "PROTAN".eq_ignore_ascii_case(filter) {
                    kind |= FilterKind::PROTAN;
                    used = true;
                }
                if "TRITAN".eq_ignore_ascii_case(filter) {
                    kind |= FilterKind::TRITAN;
                    used = true;
                }
                if "BVM97".eq_ignore_ascii_case(filter) {
                    kind |= FilterKind::BVM97;
                    used = true;
                }

                if used == false {
                    unused_filters.push(filter.to_string());
                }
            }
        }

        if unused_filters.is_empty() {
            Ok(kind)
        } else {
            Err(FilterKindParseError::UnknownMatches(unused_filters))
        }
    }

    pub fn iter(&self) -> FilterKindIter {
        FilterKindIter {
            filter: *self,
            state: 1,
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
pub enum FilterKindParseError {
    UnknownMatches(Vec<String>),
}

pub struct FilterKindIter {
    filter: FilterKind,
    state: usize,
}
impl Iterator for FilterKindIter {
    type Item = FilterKind;

    fn next(&mut self) -> Option<Self::Item> {
        let mut current = FilterKind::from_bits_truncate(self.state);

        // Check to see if we've finished counting or not.
        while current.is_empty() == false {
            self.state <<= 1;
            if self.filter.intersects(current) {
                return Some(current);
            }
            current = FilterKind::from_bits_truncate(self.state)
        }
        None
    }
}

#[derive(Debug, PartialEq)]
pub struct CombineInfo {
    total_height: u32,
    total_width: u32,
    positions: Vec<(u32, u32)>,
}

pub struct Config {
    pub combine_output: bool,
    pub processing: ProcessingStyle,
    #[cfg(feature = "labels")]
    pub render_label: bool,
}

#[derive(Clone)]
pub struct RuntimeConfig {
    pub combine_output: bool,
    #[cfg(feature = "labels")]
    pub render_label: bool,
}

#[derive(Clone)]
pub enum ProcessingStyle {
    /// Use the main thread for processing (single threaded)
    Inline,
    /// Use a threadpool for the work (autodectecion for how many workers)
    MustOffload(Option<ThreadPool>),
    /// Use a threadpool if there are more than one images to process
    MayOffload(Option<ThreadPool>),
}

pub struct ProcessingContext {
    filter_functions: Vec<(FilterKind, FilterFunc)>,
    config: RuntimeConfig,
    processing: ProcessingStyle,
    pool: Option<ThreadPool>,
}

pub enum ProcessingResult {
    Inline(Vec<NamedBuf>),
    Offload {
        pool: ThreadPool,
        sender: Sender<NamedBuf>,
        results: Receiver<NamedBuf>,
    },
}

impl ProcessingResult {
    fn new(pool: &Option<ThreadPool>) -> ProcessingResult {
        let (sender, results) = channel();
        use ProcessingResult::*;
        match pool {
            Some(pool) => Offload {
                pool: pool.clone(),
                sender,
                results,
            },
            None => Inline(Vec::with_capacity(12)),
        }
    }
    fn push(&mut self, buf: NamedBuf) {
        use ProcessingResult::*;
        match self {
            Inline(processed_images) => processed_images.push(buf),
            Offload { sender, .. } => sender.send(buf).unwrap(),
        }
    }

    fn submit_task<F>(&mut self, func: Box<F>)
    where
        F: FnOnce() -> NamedBuf + Send,
    {
        use ProcessingResult::*;
        match self {
            Inline(_) => {
                self.push(func());
            }
            Offload { pool, sender, .. } => {
                let sender = sender.clone();
                let func = unsafe {
                    std::mem::transmute::<
                        Box<dyn FnOnce() -> NamedBuf + Send>,
                        Box<dyn FnOnce() -> NamedBuf + Send + 'static>,
                    >(func)
                };
                pool.execute(move || sender.send(func()).expect("submit_task failed"));
            }
        }
    }
}
impl IntoIterator for ProcessingResult {
    type Item = NamedBuf;
    type IntoIter = Box<dyn Iterator<Item = NamedBuf>>;

    fn into_iter(self) -> Self::IntoIter {
        use ProcessingResult::*;
        match self {
            Inline(bufs) => {
                //let i: Box<dyn IntoIterator<Item = NamedBuf, IntoIter = dyn Iterator<Item = NamedBuf>>> =
                Box::new(bufs.into_iter())
                //; i
            }
            Offload { results, .. } => {
                //let i: Box<dyn IntoIterator<Item = NamedBuf, IntoIter = dyn Iterator<Item = NamedBuf>>> =
                Box::new(results.into_iter())
                //; i
            }
        }
        //self.0.into_iter()
    }
}

impl Config {
    pub fn into_context(self) -> ProcessingContext {
        let filter_functions: Vec<(FilterKind, FilterFunc)> = vec![
            (
                FilterKind::ACHROMATOMALY,
                Box::new(move |p| anomylize(p, monochrome(p))),
            ),
            (FilterKind::ACHROMATOPSIA, Box::new(move |p| monochrome(p))),
            (
                FilterKind::DEUTERANOMALY,
                Box::new(move |p| anomylize(p, blindMK(p, RevBlind::Deutan))),
            ),
            (
                FilterKind::DEUTERANOPIA,
                Box::new(move |p| blindMK(p, RevBlind::Deutan)),
            ),
            (
                FilterKind::DEUTERANOPIABVM97,
                Box::new(move |p| gimp::bvm97(RevBlind::Deutan)(p)),
            ),
            (
                FilterKind::PROTANOMALY,
                Box::new(move |p| anomylize(p, blindMK(p, RevBlind::Protan))),
            ),
            (
                FilterKind::PROTANOPIA,
                Box::new(move |p| blindMK(p, RevBlind::Protan)),
            ),
            (
                FilterKind::PROTANOPIABVM97,
                Box::new(move |p| gimp::bvm97(RevBlind::Protan)(p)),
            ),
            (
                FilterKind::TRITANOMALY,
                Box::new(move |p| anomylize(p, blindMK(p, RevBlind::Tritan))),
            ),
            (
                FilterKind::TRITANOPIA,
                Box::new(move |p| blindMK(p, RevBlind::Tritan)),
            ),
            (
                FilterKind::TRITANOPIABVM97,
                Box::new(move |p| gimp::bvm97(RevBlind::Tritan)(p)),
            ),
        ];

        let pool = match self.processing {
            ProcessingStyle::Inline => None,
            ProcessingStyle::MustOffload(ref pool) | ProcessingStyle::MayOffload(ref pool) => {
                Some(match pool {
                    None => threadpool::Builder::new().build(),
                    Some(pool) => pool.clone(),
                })
            }
        };

        ProcessingContext {
            config: RuntimeConfig {
                combine_output: self.combine_output,
                #[cfg(feature = "labels")]
                render_label: self.render_label,
            },
            processing: self.processing,
            filter_functions,
            pool,
        }
    }
}

impl ProcessingContext {
    pub fn combine_output(&mut self, combine_output: bool) {
        self.config.combine_output = combine_output;
    }
    pub fn process_file(
        &self,
        input: &PathBuf,
        filter_filter: FilterKind,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let path = format!("{}", input.display());
        let img: RgbaBuf = image::open(&input)?.to_rgba8();
        let (name, extension) = path.split_at(
            path.rfind(".")
                .expect("path does not contain a detectable extension"),
        );

        let mut images = self.process(img, filter_filter)?.into_iter();

        if self.config.combine_output {
            let out_filename = format!("{}_combined{}", name, extension);

            print!("\n    saving {} ... ", out_filename);
            flush()?;
            let (_label, buffer) = images.next().unwrap();
            buffer.save(out_filename)?;
            print!("done");
            flush()?;
        } else {
            for (label, buffer) in images {
                let out_filename = format!("{}_{}{}", name, label, extension);
                print!("\n    saving {} ... ", out_filename);
                flush()?;
                buffer.save(out_filename)?;
                print!("done");
                flush()?;
            }
        }
        Ok(())
    }

    pub fn process(
        &self,
        img: RgbaBuf,
        filter_filter: FilterKind,
    ) -> Result<ProcessingResult, Box<dyn std::error::Error>> {
        let ref config = self.config;
        let state = Arc::new((img, config.clone()));

        let mut results = ProcessingResult::new(match self.processing {
            ProcessingStyle::MayOffload(_) if filter_filter.bits().count_ones() <= 1 => &None,
            _ => &self.pool,
        });
        if config.combine_output {
            let state = state.clone();
            results.submit_task(Box::new(move || {
                let state = state.clone();
                let mut buffer = setup_buffer(&state, ORIGINAL_TEXT);
                buffer
                    .copy_from(&state.0, 0, 0)
                    .expect("unable to allocate buffer");
                (ORIGINAL_TEXT, buffer)
            }));
        }

        for (kind, func) in self
            .filter_functions
            .iter()
            .filter(|(kind, _)| filter_filter.contains(*kind))
        {
            let label = kind.to_str();

            let state = state.clone();
            let func = func.clone();
            results.submit_task(Box::new(move || {
                let mut buffer = setup_buffer(&state, &label);

                for (x, y, pixel) in state.0.enumerate_pixels() {
                    buffer.put_pixel(x, y, func(pixel));
                }

                (label, buffer)
            }));
        }

        if config.combine_output {
            let mut combined_results = ProcessingResult::new(&self.pool);

            combined_results.submit_task(Box::new(move || {
                let mut processed_images = results.into_iter().collect::<Vec<_>>();
                assert!(processed_images.len() > 0, "no partial images found");

                // Reorder the partial images if they come out of order from the ThreadPool
                // while keeping the ORIGINAL in the front
                processed_images.sort_by(|a, b| {
                    use std::cmp::Ordering::*;
                    if a.0 == ORIGINAL_TEXT {
                        return Less;
                    }
                    if b.0 == ORIGINAL_TEXT {
                        return Greater;
                    }
                    a.0.cmp(b.0)
                });

                let max_dimensions = processed_images
                    .iter()
                    .map(|t: &NamedBuf| &t.1)
                    .fold((0, 0), |b, i| b.max(i.dimensions()));
                let positions =
                    calculate_combined_positions(max_dimensions, processed_images.len());

                let mut buffer: RgbaBuf =
                    ImageBuffer::new(positions.total_width, positions.total_height);

                for ((_label, img), (x, y)) in
                    processed_images.iter().zip(positions.positions.iter())
                {
                    buffer
                        .copy_from(img, *x, *y)
                        .expect("combined_results copy_from failed");
                }

                ("Combined", buffer)
            }));

            Ok(combined_results)
        } else {
            Ok(results)
        }
    }
}

cfg_if! {
    if #[cfg(feature = "labels")] {
        fn setup_buffer(img: &Arc<(RgbaImage, RuntimeConfig)>, label: &str, ) -> RgbaImage {
            let (img, config) = &**img;
            let mut width = img.width();
            let mut height = img.height();
                let label = if config.render_label {
                    let label = labels::render(label);
                    height += label.height();
                    width = width.max(label.width());
                    Some(label)
                } else {
                    None
                };
            let mut buffer = ImageBuffer::new(width, height);
                if let Some(label) = label {
                    // centering text
                    let x = width.saturating_sub(label.width()) / 2;
                    buffer.copy_from(&label, x, img.height()).expect("setup_buffer::feature=labels");
            }
            buffer
        }
    } else {
        fn setup_buffer(img: &Arc<(RgbaImage, RuntimeConfig)>, _label: &str, ) -> RgbaImage {
            ImageBuffer::new(img.0.width(), img.0.height())
        }
    }
}

fn flush() -> std::io::Result<()> {
    use std::io::Write;
    std::io::stdout().flush()
}

/// Calculate a grid with collapsed margins
fn calculate_combined_positions((width, height): (u32, u32), n_pictures: usize) -> CombineInfo {
    assert!(n_pictures > 0, "n_pictures must be non-zero");

    let sqrt = (n_pictures as f64).sqrt();
    let columns = sqrt.ceil() as u32;
    let mut rows = sqrt.floor() as u32;

    if columns * rows < n_pictures as u32 {
        rows += 1;
    }

    let mut positions = vec![];
    for y in 0..rows {
        for x in 0..columns {
            positions.push((
                COMBINED_MARGIN + x * (width + COMBINED_MARGIN),
                COMBINED_MARGIN + y * (height + COMBINED_MARGIN),
            ));
        }
    }
    assert!(
        n_pictures <= positions.len(),
        "only {} positions generated, need {}",
        positions.len(),
        n_pictures
    );

    CombineInfo {
        total_width: COMBINED_MARGIN + (COMBINED_MARGIN + width) * columns,
        total_height: COMBINED_MARGIN + (COMBINED_MARGIN + height) * rows,
        positions,
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use pretty_assertions::assert_eq;

    const MARGIN: u32 = 25;

    #[test]
    fn single() {
        let positions = calculate_combined_positions((20, 10), 1);
        assert_eq!(positions.total_width, MARGIN + 20 + MARGIN);
        assert_eq!(positions.total_height, MARGIN + 10 + MARGIN);
        assert_eq!(positions.positions, [(MARGIN, MARGIN)]);
    }

    #[test]
    fn three() {
        let positions = calculate_combined_positions((20, 10), 3);
        assert_eq!(
            CombineInfo {
                total_width: MARGIN + 2 * (20 + MARGIN),
                total_height: MARGIN + 2 * (10 + MARGIN),
                positions: vec![
                    (MARGIN, MARGIN),
                    (MARGIN + 20 + MARGIN, MARGIN), // expect a 2x2 matrix
                    (MARGIN, MARGIN + 10 + MARGIN),
                    (MARGIN + 20 + MARGIN, MARGIN + 10 + MARGIN),
                ],
            },
            positions
        );
    }

    #[test]
    fn four() {
        let positions = calculate_combined_positions((20, 10), 4);
        assert_eq!(
            CombineInfo {
                total_width: MARGIN + 2 * (20 + MARGIN),
                total_height: MARGIN + 2 * (10 + MARGIN),
                positions: vec![
                    (MARGIN, MARGIN),
                    (MARGIN + 20 + MARGIN, MARGIN), // expect a 2x2 matrix
                    (MARGIN, MARGIN + 10 + MARGIN),
                    (MARGIN + 20 + MARGIN, MARGIN + 10 + MARGIN),
                ],
            },
            positions
        );
    }

    #[test]
    fn twelfe() {
        let positions = calculate_combined_positions((20, 10), 12);
        let row1 = MARGIN + 10 + MARGIN;
        let row2 = MARGIN + 2 * (10 + MARGIN);
        assert_eq!(
            CombineInfo {
                total_width: MARGIN + 4 * (20 + MARGIN),
                total_height: MARGIN + 3 * (10 + MARGIN),
                positions: vec![
                    (MARGIN + 0 * 45, MARGIN),
                    (MARGIN + 1 * 45, MARGIN),
                    (MARGIN + 2 * 45, MARGIN),
                    (MARGIN + 3 * 45, MARGIN), // expect a 4x3 matrix
                    (MARGIN + 0 * 45, row1),
                    (MARGIN + 1 * 45, row1),
                    (MARGIN + 2 * 45, row1),
                    (MARGIN + 3 * 45, row1), // row 2
                    (MARGIN + 0 * 45, row2),
                    (MARGIN + 1 * 45, row2),
                    (MARGIN + 2 * 45, row2),
                    (MARGIN + 3 * 45, row2)
                ],
            },
            positions
        );
    }

    #[test]
    fn iter_empty() {
        let mut iter = FilterKind::empty().iter();
        assert_eq!(None, iter.next());
    }

    #[test]
    fn iter_last() {
        let mut iter = FilterKind::TRITANOPIABVM97.iter();
        assert_eq!(Some(FilterKind::TRITANOPIABVM97), iter.next());
        assert_eq!(None, iter.next());
    }

    #[test]
    fn iter_first_two() {
        let mut iter = (FilterKind::ORIGINAL | FilterKind::ACHROMATOMALY).iter();
        assert_eq!(Some(FilterKind::ORIGINAL), iter.next());
        assert_eq!(Some(FilterKind::ACHROMATOMALY), iter.next());
        assert_eq!(None, iter.next());
    }

    #[test]
    fn from_strs() {
        let expect = Ok(FilterKind::PROTANOPIABVM97 | FilterKind::ACHROMATOMALY);
        let input = ["ACHROMATOMALY", "  PROTANOPIABVM97  "];
        assert_eq!(expect, FilterKind::from_strs(&input));
    }

    #[test]
    fn from_strs_group() {
        let expect = Ok(FilterKind::MONOCHROME);
        let input = ["MONOCHROME"];
        assert_eq!(expect, FilterKind::from_strs(&input));
    }

    #[test]
    fn from_strs_not_found() {
        let expect = Err(FilterKindParseError::UnknownMatches(vec![
            "something else".to_string(),
            "something different".to_string(),
        ]));
        let input = ["MONOCHROME", "something else", "something different"];
        assert_eq!(expect, FilterKind::from_strs(&input));
    }

    #[test]
    fn to_str() {
        assert_eq!("MONOCHROME", FilterKind::MONOCHROME.to_str());
    }
}