retroimg/
color.rs

1//! Color depth manipulation module
2use exoquant::ditherer::FloydSteinberg;
3use exoquant::optimizer::{KMeans, Optimizer};
4use exoquant::{Color, Histogram, Quantizer, Remapper, SimpleColorSpace};
5use image::{ImageBuffer, Rgb, RgbImage};
6use itertools::Itertools;
7use num_integer::Roots;
8use std::str::FromStr;
9
10pub mod cga;
11pub mod ega;
12
13/// Enumeration of supported color distance algorithms
14/// for loss calculation.
15///
16/// The use of one algorithm or the other may affect slightly
17/// which palette colors are chosen,
18/// especially in modes such as CGA.
19#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
20pub enum LossAlgorithm {
21    /// L2, Euclidean distance
22    #[default]
23    L2,
24    /// L1, Manhattan distance
25    L1,
26}
27
28impl std::fmt::Display for LossAlgorithm {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        match self {
31            LossAlgorithm::L1 => f.write_str("L1"),
32            LossAlgorithm::L2 => f.write_str("L2"),
33        }
34    }
35}
36
37/// An error returned by a failed attempt at
38/// creating a [`LossAlgorithm`] from a string.
39#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
40pub struct LossAlgorithmParseError;
41
42impl std::fmt::Display for LossAlgorithmParseError {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        f.write_str("invalid distance/loss algorithm, should be \"L1\" or \"L2\"")
45    }
46}
47
48impl std::error::Error for LossAlgorithmParseError {}
49
50impl FromStr for LossAlgorithm {
51    type Err = LossAlgorithmParseError;
52
53    fn from_str(s: &str) -> Result<Self, Self::Err> {
54        match s {
55            "L1" | "l1" => Ok(LossAlgorithm::L1),
56            "L2" | "l2" => Ok(LossAlgorithm::L2),
57            _ => Err(LossAlgorithmParseError),
58        }
59    }
60}
61
62impl LossAlgorithm {
63    /// calculate the difference between 2 colors
64    /// using the given loss algorithm
65    #[inline]
66    pub fn color_diff(self, c1: Color, c2: Color) -> u64 {
67        match self {
68            LossAlgorithm::L1 => color_diff_l1(c1, c2),
69            LossAlgorithm::L2 => color_diff_l2(c1, c2),
70        }
71    }
72
73    /// calculate the difference between 2 images
74    /// using the default loss
75    ///
76    /// # Panic
77    ///
78    /// Panics if the two slices of colors do not have the same length.
79    pub fn image_diff(self, a: &[Color], b: &[Color]) -> u64 {
80        assert_eq!(a.len(), b.len());
81        Iterator::zip(a.iter(), b.iter())
82            .map(|(a, b)| self.color_diff(*a, *b))
83            .sum()
84    }
85}
86
87/// calculate the L1 difference between 2 colors
88fn color_diff_l1(c1: Color, c2: Color) -> u64 {
89    let Color {
90        r: r1,
91        g: g1,
92        b: b1,
93        ..
94    } = c1;
95    let Color {
96        r: r2,
97        g: g2,
98        b: b2,
99        ..
100    } = c2;
101    let (r1, r2) = (i64::from(r1), i64::from(r2));
102    let (g1, g2) = (i64::from(g1), i64::from(g2));
103    let (b1, b2) = (i64::from(b1), i64::from(b2));
104    (r1 - r2).abs() as u64 + (g1 - g2).abs() as u64 + (b1 - b2).abs() as u64
105}
106
107/// calculate the L2 difference between 2 colors
108fn color_diff_l2(c1: Color, c2: Color) -> u64 {
109    let Color {
110        r: r1,
111        g: g1,
112        b: b1,
113        ..
114    } = c1;
115    let Color {
116        r: r2,
117        g: g2,
118        b: b2,
119        ..
120    } = c2;
121    let (r1, r2) = (i64::from(r1), i64::from(r2));
122    let (g1, g2) = (i64::from(g1), i64::from(g2));
123    let (b1, b2) = (i64::from(b1), i64::from(b2));
124    let dr = (r1 - r2) as u64;
125    let dg = (g1 - g2) as u64;
126    let db = (b1 - b2) as u64;
127
128    (dr.saturating_mul(dr)
129        .saturating_add(dg.saturating_mul(dg))
130        .saturating_add(db.saturating_mul(db)))
131    .sqrt()
132}
133
134/// calculate the median RGB color of the given buffer
135fn color_median(colors: &[Color]) -> Color {
136    let mut colors_r = colors.iter().map(|c| c.r).collect_vec();
137    let mut colors_g = colors.iter().map(|c| c.g).collect_vec();
138    let mut colors_b = colors.iter().map(|c| c.b).collect_vec();
139    colors_r.sort_unstable();
140    colors_g.sort_unstable();
141    colors_b.sort_unstable();
142    let r = colors_r[colors_r.len() / 2];
143    let g = colors_r[colors_g.len() / 2];
144    let b = colors_r[colors_b.len() / 2];
145
146    Color { r, g, b, a: 255 }
147}
148
149/// The options for transforming an image to have a different color depth.
150#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
151pub struct ColorOptions {
152    /// The maximum number of colors to admit.
153    /// `None` means no limit
154    pub num_colors: Option<u32>,
155
156    /// The distance algorithm to use for calculating the loss between colors.
157    ///
158    /// The default is L2.
159    pub loss: LossAlgorithm,
160}
161
162/// Color depth image converter.
163pub trait ColorDepth {
164    /// Convert and retrieve the loss from converting an image.
165    fn convert_image_with_loss(&self, image: &RgbImage, options: ColorOptions)
166        -> (Vec<Color>, u64);
167
168    /// Convert an RGB image to this color depth.
169    fn convert_image(&self, image: &RgbImage, options: ColorOptions) -> Vec<Color> {
170        self.convert_image_with_loss(image, options).0
171    }
172
173    /// Estimate the loss obtained from converting an image.
174    /// For the best results, greater discrepancies should result in higher
175    /// loss values.
176    fn loss(&self, image: &RgbImage, options: ColorOptions) -> u64 {
177        self.convert_image_with_loss(image, options).1
178    }
179}
180
181impl<'a, T: ColorDepth> ColorDepth for &'a T {
182    fn convert_image_with_loss(
183        &self,
184        image: &RgbImage,
185        options: ColorOptions,
186    ) -> (Vec<Color>, u64) {
187        (**self).convert_image_with_loss(image, options)
188    }
189
190    /// Convert an RGB image to this color depth.
191    fn convert_image(&self, image: &RgbImage, options: ColorOptions) -> Vec<Color> {
192        (**self).convert_image(image, options)
193    }
194
195    /// Estimate the loss obtained from converting an image.
196    /// For the best results, greater discrepancies should result in higher
197    /// loss values.
198    fn loss(&self, image: &RgbImage, options: ColorOptions) -> u64 {
199        (**self).loss(image, options)
200    }
201}
202
203/// Trait for anything which maps one color to another.
204pub trait ColorMapper {
205    /// Convert a single color
206    fn convert_color(&self, c: Color) -> Color;
207}
208
209impl<'a, T: ColorMapper> ColorMapper for &'a T {
210    fn convert_color(&self, c: Color) -> Color {
211        (**self).convert_color(c)
212    }
213}
214
215impl ColorMapper for fn(Color) -> Color {
216    fn convert_color(&self, c: Color) -> Color {
217        self(c)
218    }
219}
220
221/// A color depth implementation with color mapping.
222#[derive(Debug, Default, Copy, Clone)]
223pub struct MappingColorDepth<M>(M);
224
225impl<M> ColorMapper for MappingColorDepth<M>
226where
227    M: ColorMapper,
228{
229    fn convert_color(&self, c: Color) -> Color {
230        self.0.convert_color(c)
231    }
232}
233
234impl<M> ColorDepth for MappingColorDepth<M>
235where
236    M: ColorMapper,
237{
238    fn convert_image_with_loss(
239        &self,
240        image: &RgbImage,
241        options: ColorOptions,
242    ) -> (Vec<Color>, u64) {
243        let original = image
244            .pixels()
245            .map(|&p| {
246                let Rgb([r, g, b]) = p;
247                Color { r, g, b, a: 255 }
248            })
249            .collect_vec();
250        let pixels = image
251            .pixels()
252            .map(|&p| {
253                let Rgb([r, g, b]) = p;
254                self.0.convert_color(Color { r, g, b, a: 255 })
255            })
256            .collect_vec();
257
258        // optimize palette and dither
259        let converted_pixels = if let Some(num_colors) = options.num_colors {
260            let mut palette = build_palette(&pixels, num_colors);
261
262            // reduce palette's color depth
263            for c in &mut palette {
264                *c = self.convert_color(*c);
265            }
266
267            let colorspace = SimpleColorSpace::default();
268            let ditherer = FloydSteinberg::new();
269            let remapper = Remapper::new(&palette, &colorspace, &ditherer);
270            let indexed_data = remapper.remap(&pixels, image.width() as usize);
271            indexed_data
272                .into_iter()
273                .map(|i| palette[i as usize])
274                .collect_vec()
275        } else {
276            pixels
277        };
278        let loss = options.loss.image_diff(&original, &converted_pixels);
279        (converted_pixels, loss)
280    }
281}
282
283/// True 24-bit color, 8 bits per channel, virtually no limit in color depth.
284#[derive(Debug, Default, Copy, Clone)]
285pub struct TrueColor24BitMapper;
286
287impl ColorMapper for TrueColor24BitMapper {
288    fn convert_color(&self, pixel: Color) -> Color {
289        pixel
290    }
291}
292
293/// A color depth implementation which retains all 24 bits per pixel.
294pub type TrueColor24Bit = MappingColorDepth<TrueColor24BitMapper>;
295
296impl TrueColor24Bit {
297    pub fn new() -> Self {
298        MappingColorDepth::default()
299    }
300}
301
302/// A color mapper that reduces sample precision to 6 bits per sample
303/// (18 bits per pixel).
304#[derive(Debug, Default, Copy, Clone)]
305pub struct Vga18BitMapper;
306
307impl ColorMapper for Vga18BitMapper {
308    fn convert_color(&self, pixel: Color) -> Color {
309        let Color { r, g, b, a } = pixel;
310        Color {
311            r: (r & !0x03) | r >> 6,
312            g: (g & !0x03) | g >> 6,
313            b: (b & !0x03) | b >> 6,
314            a,
315        }
316    }
317}
318
319/// A color mapper that reduces pixel precision to 16 bits per pixel
320/// (5 bits for red, 6 bits for green, 6 bits for blue).
321#[derive(Debug, Default, Copy, Clone)]
322pub struct Vga16BitMapper;
323
324impl ColorMapper for Vga16BitMapper {
325    fn convert_color(&self, pixel: Color) -> Color {
326        let Color { r, g, b, a } = pixel;
327        Color {
328            r: (r & !0x07) | r >> 5,
329            g: (g & !0x03) | g >> 6,
330            b: (b & !0x07) | b >> 5,
331            a,
332        }
333    }
334}
335
336/// VGA 18-bit color (64 levels per channel)
337pub type Vga18Bit = MappingColorDepth<Vga18BitMapper>;
338
339impl Vga18Bit {
340    pub fn new() -> Self {
341        MappingColorDepth::default()
342    }
343}
344
345/// VGA 16-bit color, also called High color mode in legacy systems
346/// (5 bits for red and blue channels, 6 bits for green)
347pub type Vga16Bit = MappingColorDepth<Vga16BitMapper>;
348
349impl Vga16Bit {
350    pub fn new() -> Self {
351        MappingColorDepth::default()
352    }
353}
354
355/// Color depth defined by a hardware-level palette of RGB colors.
356#[derive(Debug, Copy, Clone)]
357pub struct FixedPalette<T>(T);
358
359impl<T> FixedPalette<T>
360where
361    T: AsRef<[[u8; 3]]>,
362{
363    fn convert_color(&self, pixel: Color) -> Color {
364        let Color {
365            r: sr,
366            g: sg,
367            b: sb,
368            a: _,
369        } = pixel;
370        let (sr, sg, sb) = (i32::from(sr), i32::from(sg), i32::from(sb));
371        let [r, g, b] = *self
372            .0
373            .as_ref()
374            .iter()
375            .min_by_key(|[pr, pg, pb]| {
376                let (pr, pg, pb) = (i32::from(*pr), i32::from(*pg), i32::from(*pb));
377                let rd = sr - pr;
378                let rg = sg - pg;
379                let rb = sb - pb;
380                rd * rd + rg * rg + rb * rb
381            })
382            .unwrap();
383        Color { r, g, b, a: 255 }
384    }
385}
386
387impl<T> ColorDepth for FixedPalette<T>
388where
389    T: AsRef<[[u8; 3]]>,
390{
391    fn convert_image_with_loss(
392        &self,
393        image: &RgbImage,
394        options: ColorOptions,
395    ) -> (Vec<Color>, u64) {
396        let original = image
397            .pixels()
398            .map(|&p| {
399                let Rgb([r, g, b]) = p;
400                Color { r, g, b, a: 255 }
401            })
402            .collect_vec();
403
404        // optimize palette and dither
405        let converted_pixels = if let Some(num_colors) = options.num_colors {
406            let mut palette = build_palette(&original, num_colors);
407
408            // reduce palette's color depth
409            for c in &mut palette {
410                *c = self.convert_color(*c);
411            }
412
413            let colorspace = SimpleColorSpace::default();
414            let ditherer = FloydSteinberg::new();
415            let remapper = Remapper::new(&palette, &colorspace, &ditherer);
416            let indexed_data = remapper.remap(&original, image.width() as usize);
417            indexed_data
418                .into_iter()
419                .map(|i| palette[i as usize])
420                .collect_vec()
421        } else {
422            original.clone()
423        };
424        let loss = options.loss.image_diff(&original, &converted_pixels);
425        (converted_pixels, loss)
426    }
427}
428
429fn build_palette(pixels: &[Color], num_colors: u32) -> Vec<Color> {
430    // optimize palette and dither
431    let mut histogram = Histogram::new();
432    histogram.extend(pixels.iter().cloned());
433    let colorspace = SimpleColorSpace::default();
434    let optimizer = KMeans;
435    let mut quantizer = Quantizer::new(&histogram, &colorspace);
436    while quantizer.num_colors() < num_colors as usize {
437        quantizer.step();
438        // very optional optimization, !very slow!
439        // you probably only want to do this every N steps, if at all.
440        if quantizer.num_colors() % 256 == 0 {
441            quantizer = quantizer.optimize(&optimizer, 16);
442        }
443    }
444
445    let palette = quantizer.colors(&colorspace);
446    // this optimization is more useful than the above and a lot less slow
447    optimizer.optimize_palette(&colorspace, &palette, &histogram, 8)
448}
449
450/// Color depth emulating a combination of one freely selectable
451/// background color (`B`) with any of the other colors (`F`).
452#[derive(Debug, Copy, Clone)]
453pub struct BackForePalette<B, F>(B, F);
454
455impl<B, F> BackForePalette<B, F>
456where
457    B: AsRef<[[u8; 3]]>,
458    F: AsRef<[[u8; 3]]>,
459{
460    fn convert_color<T>(pixel: Color, palette: T) -> Color
461    where
462        T: AsRef<[[u8; 3]]>,
463    {
464        let Color {
465            r: sr,
466            g: sg,
467            b: sb,
468            a: _,
469        } = pixel;
470        let (sr, sg, sb) = (i32::from(sr), i32::from(sg), i32::from(sb));
471        let [r, g, b] = *palette
472            .as_ref()
473            .iter()
474            .min_by_key(|[pr, pg, pb]| {
475                let (pr, pg, pb) = (i32::from(*pr), i32::from(*pg), i32::from(*pb));
476                let rd = sr - pr;
477                let rg = sg - pg;
478                let rb = sb - pb;
479                rd * rd + rg * rg + rb * rb
480            })
481            .unwrap();
482        Color { r, g, b, a: 255 }
483    }
484
485    fn convert_color_back(&self, pixel: Color) -> Color {
486        BackForePalette::<B, F>::convert_color(pixel, &self.0)
487    }
488
489    /// Identify the best background color
490    fn background_color(&self, image: &RgbImage) -> Color {
491        // we'll fetch the median color of the image for the time being
492        let original = image
493            .pixels()
494            .map(|&p| {
495                let Rgb([r, g, b]) = p;
496                Color { r, g, b, a: 255 }
497            })
498            .collect_vec();
499        color_median(&original)
500    }
501}
502
503impl<B, F> ColorDepth for BackForePalette<B, F>
504where
505    B: AsRef<[[u8; 3]]>,
506    F: AsRef<[[u8; 3]]>,
507{
508    fn convert_image_with_loss(
509        &self,
510        image: &RgbImage,
511        options: ColorOptions,
512    ) -> (Vec<Color>, u64) {
513        // first try to identify the background color
514        let bkg_color = self.background_color(image);
515        let bkg_color = self.convert_color_back(bkg_color);
516
517        // then build a palette with the extra color
518        let mut fixed = self.1.as_ref().to_vec();
519        fixed.push([bkg_color.r, bkg_color.g, bkg_color.b]);
520        let fixed = FixedPalette(fixed);
521
522        let original = image
523            .pixels()
524            .map(|&p| {
525                let Rgb([r, g, b]) = p;
526                Color { r, g, b, a: 255 }
527            })
528            .collect_vec();
529
530        // optimize palette and dither
531        let converted_pixels = if let Some(num_colors) = options.num_colors {
532            let mut palette = build_palette(&original, num_colors);
533
534            // reduce palette's color depth
535            for c in &mut palette {
536                *c = fixed.convert_color(*c);
537            }
538
539            let colorspace = SimpleColorSpace::default();
540            let ditherer = FloydSteinberg::new();
541            let remapper = Remapper::new(&palette, &colorspace, &ditherer);
542            let indexed_data = remapper.remap(&original, image.width() as usize);
543            indexed_data
544                .into_iter()
545                .map(|i| palette[i as usize])
546                .collect_vec()
547        } else {
548            original.clone()
549        };
550        let loss = options.loss.image_diff(&original, &converted_pixels);
551
552        (converted_pixels, loss)
553    }
554}
555
556/// A collection of palettes, the one yielding the lowest loss is used.
557#[derive(Debug, Copy, Clone)]
558pub struct BestPalette<C>(C);
559
560impl<C, P> ColorDepth for BestPalette<C>
561where
562    C: std::ops::Deref<Target = [P]>,
563    P: ColorDepth,
564{
565    fn convert_image_with_loss(
566        &self,
567        image: &RgbImage,
568        options: ColorOptions,
569    ) -> (Vec<Color>, u64) {
570        self.0
571            .iter()
572            .map(|cd| cd.convert_image_with_loss(image, options))
573            .min_by_key(|(_pixels, loss)| *loss)
574            .unwrap()
575    }
576}
577
578pub fn colors_to_image<I>(width: u32, height: u32, pixels: I) -> RgbImage
579where
580    I: IntoIterator<Item = Color>,
581{
582    let pixels = pixels
583        .into_iter()
584        .flat_map(|Color { r, g, b, .. }| [r, g, b])
585        .collect_vec();
586    ImageBuffer::from_raw(width, height, pixels).expect("there should be enough pixels")
587}
588
589/// 64 color palette established by the full-color EGA standard.
590pub static PALETTE_BW_1BIT: FixedPalette<&[[u8; 3]]> = FixedPalette(BW_1BIT);
591
592/// 64 color palette established by the full-color EGA standard.
593pub static BW_1BIT: &[[u8; 3]] = &[[0, 0, 0], [0xFF, 0xFF, 0xFF]];
retroimg/color.rs

retroimg/
color.rs
