nihav-core/src/scale/palette/mediancut.rs

   1 use super::Pixel;
   2
   3 struct VQBox<'a> {
   4     pixels: &'a mut [Pixel],
   5     max:    Pixel,
   6     min:    Pixel,
   7 }
   8
   9 fn avg_u8(a: u8, b: u8) -> u8 {
  10     (a & b) + ((a ^ b) >> 1)
  11 }
  12
  13 impl<'a> VQBox<'a> {
  14     fn new(pixels: &'a mut [Pixel]) -> Self {
  15         let mut max = Pixel{ r:   0, g:   0, b:   0 };
  16         let mut min = Pixel{ r: 255, g: 255, b: 255 };
  17         for pix in pixels.iter() {
  18             max = max.max(*pix);
  19             min = min.min(*pix);
  20         }
  21         Self { pixels, max, min }
  22     }
  23     fn can_split(&self) -> bool {
  24         let dr = self.max.r - self.min.r;
  25         let dg = self.max.g - self.min.g;
  26         let db = self.max.b - self.min.b;
  27         (dr | dg | db) != 0
  28     }
  29     fn sort<F>(arr: &mut [Pixel], idx: F) -> usize where F: Fn(&Pixel) -> usize {
  30         let mut buckets: Vec<Vec<Pixel>> = Vec::with_capacity(256);
  31         let mut counts = [0; 256];
  32         for pix in arr.iter() {
  33             counts[idx(pix)] += 1;
  34         }
  35         for i in 0..256 {
  36             buckets.push(Vec::with_capacity(counts[i]));
  37         }
  38         for pix in arr.iter() {
  39             buckets[idx(pix)].push(*pix);
  40         }
  41         let mut start = 0;
  42         let mut pivot = 0;
  43         let mut cand1 = 0;
  44         for bucket in buckets.iter() {
  45             let end = bucket.len();
  46             if (start > 0) && (start <= arr.len() / 2) {
  47                 pivot = start + end;
  48                 cand1 = start;
  49             }
  50             (&mut arr[start..][..end]).copy_from_slice(bucket.as_slice());
  51             start += end;
  52         }
  53         for bucket in buckets.iter() {
  54             if pivot != 0 {
  55                 break;
  56             }
  57             pivot += bucket.len();
  58         }
  59         if pivot < arr.len() || cand1 == 0 {
  60             pivot
  61         } else {
  62             cand1
  63         }
  64     }
  65     fn calc_min_and_max(pixels: &[Pixel]) -> (Pixel, Pixel) {
  66         let mut max = Pixel{ r:   0, g:   0, b:   0 };
  67         let mut min = Pixel{ r: 255, g: 255, b: 255 };
  68         for pix in pixels.iter() {
  69             max = max.max(*pix);
  70             min = min.min(*pix);
  71         }
  72         (min, max)
  73     }
  74     fn split(self) -> (VQBox<'a>, VQBox<'a>) {
  75         let dr = self.max.r - self.min.r;
  76         let dg = self.max.g - self.min.g;
  77         let db = self.max.b - self.min.b;
  78
  79         let box0;
  80         let box1;
  81         if (dr > dg) && (dr >= db) {
  82             let pivot = Self::sort(self.pixels, |pix| pix.r as usize);
  83             let (part0, part1) = self.pixels.split_at_mut(pivot);
  84             let (min0, max0) = Self::calc_min_and_max(part0);
  85             let (min1, max1) = Self::calc_min_and_max(part1);
  86             box0 = VQBox { pixels: part0, max: max0, min: min0 };
  87             box1 = VQBox { pixels: part1, max: max1, min: min1 };
  88         } else if (db > dr) && (db > dg) {
  89             let pivot = Self::sort(self.pixels, |pix| pix.b as usize);
  90             let (part0, part1) = self.pixels.split_at_mut(pivot);
  91             let (min0, max0) = Self::calc_min_and_max(part0);
  92             let (min1, max1) = Self::calc_min_and_max(part1);
  93             box0 = VQBox { pixels: part0, max: max0, min: min0 };
  94             box1 = VQBox { pixels: part1, max: max1, min: min1 };
  95         } else {
  96             let pivot = Self::sort(self.pixels, |pix| pix.g as usize);
  97             let (part0, part1) = self.pixels.split_at_mut(pivot);
  98             let (min0, max0) = Self::calc_min_and_max(part0);
  99             let (min1, max1) = Self::calc_min_and_max(part1);
 100             box0 = VQBox { pixels: part0, max: max0, min: min0 };
 101             box1 = VQBox { pixels: part1, max: max1, min: min1 };
 102         }
 103         (box0, box1)
 104     }
 105 }
 106
 107 pub fn quantise_median_cut(src: &[Pixel], pal: &mut [[u8; 3]; 256]) -> usize {
 108     let mut pixels = Vec::with_capacity(src.len());
 109     pixels.extend_from_slice(src);
 110     VQBox::sort(pixels.as_mut_slice(), |pix| pix.r as usize);
 111     VQBox::sort(pixels.as_mut_slice(), |pix| pix.g as usize);
 112     VQBox::sort(pixels.as_mut_slice(), |pix| pix.b as usize);
 113     let box0 = VQBox::new(pixels.as_mut_slice());
 114     let mut boxes: Vec<VQBox> = Vec::with_capacity(256);
 115     boxes.push(box0);
 116     let mut changed = true;
 117     while changed && boxes.len() < 256 {
 118         let end = boxes.len();
 119         changed = false;
 120         let mut split_largest = false;
 121         for _ in 0..end {
 122             let curbox = boxes.remove(0);
 123             if curbox.can_split() {
 124                 let (box0, box1) = curbox.split();
 125                 boxes.push(box0);
 126                 boxes.push(box1);
 127                 changed = true;
 128             } else {
 129                 boxes.push(curbox);
 130                 split_largest = true;
 131                 break;
 132             }
 133             if boxes.len() == 256 {
 134                 break;
 135             }
 136         }
 137         if split_largest {
 138             let mut maxidx = 0;
 139             let mut lcount = 0;
 140             for (i, cbox) in boxes.iter().enumerate() {
 141                 if cbox.can_split() && cbox.pixels.len() > lcount {
 142                     lcount = cbox.pixels.len();
 143                     maxidx = i;
 144                 }
 145             }
 146             if lcount > 0 {
 147                 let curbox = boxes.remove(maxidx);
 148                 let (box0, box1) = curbox.split();
 149                 boxes.push(box0);
 150                 boxes.push(box1);
 151                 changed = true;
 152             }
 153         }
 154     }
 155     for (curbox, palentry) in boxes.iter().zip(pal.iter_mut()) {
 156         palentry[0] = avg_u8(curbox.min.r, curbox.max.r);
 157         palentry[1] = avg_u8(curbox.min.g, curbox.max.g);
 158         palentry[2] = avg_u8(curbox.min.b, curbox.max.b);
 159     }
 160
 161     boxes.len()
 162 }