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

   1 use super::Pixel;
   2
   3 struct RNG {
   4     seed: u16,
   5 }
   6
   7 impl RNG {
   8     fn new() -> Self { Self { seed: 0x1234 } }
   9     fn next(&mut self) -> u8 {
  10         if (self.seed & 0x8000) != 0 {
  11             self.seed = ((self.seed & 0x7FFF) * 2) ^ 0x1B2B;
  12         } else {
  13             self.seed <<= 1;
  14         }
  15         self.seed as u8
  16     }
  17 }
  18
  19 #[derive(Default,Clone,Copy,PartialEq,Debug)]
  20 struct Entry {
  21     pix:        Pixel,
  22     count:      u64,
  23 }
  24
  25 struct Cluster {
  26     centroid:   Pixel,
  27     dist:       u64,
  28     count:      u64,
  29     sum_r:      u64,
  30     sum_g:      u64,
  31     sum_b:      u64,
  32 }
  33
  34 impl Cluster {
  35     fn new(centroid: Pixel) -> Self {
  36         Self {
  37             centroid,
  38             dist:       0,
  39             count:      0,
  40             sum_r:      0,
  41             sum_g:      0,
  42             sum_b:      0,
  43         }
  44     }
  45     fn reset(&mut self) {
  46         self.count = 0;
  47         self.sum_r = 0;
  48         self.sum_g = 0;
  49         self.sum_b = 0;
  50         self.dist  = 0;
  51     }
  52     fn add_pixel(&mut self, entry: &Entry) {
  53         self.sum_r += u64::from(entry.pix.r) * entry.count;
  54         self.sum_g += u64::from(entry.pix.g) * entry.count;
  55         self.sum_b += u64::from(entry.pix.b) * entry.count;
  56         self.count += entry.count;
  57     }
  58     fn add_dist(&mut self, entry: &Entry) {
  59         self.dist += u64::from(self.centroid.dist(entry.pix)) * entry.count;
  60     }
  61     fn calc_centroid(&mut self) {
  62         if self.count != 0 {
  63             self.centroid.r = ((self.sum_r + self.count / 2) / self.count) as u8;
  64             self.centroid.g = ((self.sum_g + self.count / 2) / self.count) as u8;
  65             self.centroid.b = ((self.sum_b + self.count / 2) / self.count) as u8;
  66         }
  67     }
  68     fn calc_dist(&mut self) {
  69         if self.count != 0 {
  70             self.dist = (self.dist + self.count / 2) / self.count;
  71         }
  72     }
  73 }
  74
  75 pub struct ELBG {
  76     clusters:   Vec<Cluster>,
  77 }
  78
  79 impl ELBG {
  80     #[allow(dead_code)]
  81     pub fn new(initial_pal: &[[u8; 3]; 256]) -> Self {
  82         let mut clusters = Vec::with_capacity(256);
  83         for i in 0..256 {
  84             let pix = Pixel { r: initial_pal[i][0], g: initial_pal[i][1], b: initial_pal[i][2] };
  85             let cluster = Cluster::new(pix);
  86             clusters.push(cluster);
  87         }
  88         Self {
  89             clusters,
  90         }
  91     }
  92     #[allow(dead_code)]
  93     pub fn new_random() -> Self {
  94         let mut rng = RNG::new();
  95         let mut clusters = Vec::with_capacity(256);
  96         for _ in 0..256 {
  97             let pix = Pixel { r: rng.next(), g: rng.next(), b: rng.next() };
  98             let cluster = Cluster::new(pix);
  99             clusters.push(cluster);
 100         }
 101         Self {
 102             clusters,
 103         }
 104     }
 105     fn sort<F>(arr: &mut [Pixel], idx: F) where F: Fn(&Pixel) -> u8 {
 106         let mut dst = vec![Pixel::default(); arr.len()];
 107         let mut counts = [0; 256];
 108         for pix in arr.iter() {
 109             counts[idx(pix) as usize] += 1;
 110         }
 111         let mut last = counts[0];
 112         counts[0] = 0;
 113         for count in counts.iter_mut().skip(1) {
 114             let plast = last;
 115             last += *count;
 116             *count = plast;
 117         }
 118         for pix in arr.iter() {
 119             let bucket = idx(pix) as usize;
 120             dst[counts[bucket]] = *pix;
 121             counts[bucket] += 1;
 122         }
 123         arr.copy_from_slice(dst.as_slice());
 124     }
 125     fn new_split(old_index: usize, entries: &[Entry], indices: &[usize]) -> Option<(Pixel, Pixel)> {
 126         let mut max = Pixel { r:   0, g:   0, b:   0 };
 127         let mut min = Pixel { r: 255, g: 255, b: 255 };
 128         let mut found = false;
 129         for (entry, idx) in entries.iter().zip(indices) {
 130             if *idx == old_index {
 131                 max = max.max(entry.pix);
 132                 min = min.min(entry.pix);
 133                 found = true;
 134             }
 135         }
 136         if !found {
 137             return None;
 138         }
 139         let dr = max.r - min.r;
 140         let dg = max.g - min.g;
 141         let db = max.b - min.b;
 142         let cent0 = Pixel { r: min.r + dr / 3, g: min.g + dg / 3, b: min.b + db / 3 };
 143         let cent1 = Pixel { r: max.r - dr / 3, g: max.g - dg / 3, b: max.b - db / 3 };
 144         Some((cent0, cent1))
 145     }
 146     fn old_centre(&self, old_index1: usize, old_index2: usize, entries: &[Entry], indices: &[usize]) -> Pixel {
 147         let mut max = Pixel { r:   0, g:   0, b:   0 };
 148         let mut min = Pixel { r: 255, g: 255, b: 255 };
 149         let mut found = false;
 150         for (entry, idx) in entries.iter().zip(indices) {
 151             if *idx == old_index1 || *idx == old_index2 {
 152                 max = max.max(entry.pix);
 153                 min = min.min(entry.pix);
 154                 found = true;
 155             }
 156         }
 157         if !found {
 158             max = self.clusters[old_index1].centroid.max(self.clusters[old_index2].centroid);
 159             min = self.clusters[old_index1].centroid.min(self.clusters[old_index2].centroid);
 160         }
 161         let dr = max.r - min.r;
 162         let dg = max.g - min.g;
 163         let db = max.b - min.b;
 164         Pixel { r: min.r + dr / 2, g: min.g + dg / 2, b: min.b + db / 2 }
 165     }
 166     fn estimate_old(old_idx0: usize, old_idx1: usize, c: Pixel, entries: &[Entry], indices: &[usize]) -> u64 {
 167         let mut clu = Cluster::new(c);
 168         let mut count = 0;
 169         for (entry, idx) in entries.iter().zip(indices) {
 170             if *idx == old_idx0 || *idx == old_idx1 {
 171                 clu.add_dist(entry);
 172                 count += entry.count;
 173             }
 174         }
 175         clu.count = count;
 176         clu.calc_dist();
 177         clu.dist
 178     }
 179     fn estimate_new(c0: Pixel, c1: Pixel, old_idx: usize, entries: &[Entry], indices: &[usize]) -> u64 {
 180         let mut clu0 = Cluster::new(c0);
 181         let mut clu1 = Cluster::new(c1);
 182         let mut count0 = 0;
 183         let mut count1 = 0;
 184         for (entry, idx) in entries.iter().zip(indices) {
 185             if *idx == old_idx {
 186                 if c0.dist(entry.pix) < c1.dist(entry.pix) {
 187                     clu0.add_dist(entry);
 188                     count0 += entry.count;
 189                 } else {
 190                     clu1.add_dist(entry);
 191                     count1 += entry.count;
 192                 }
 193             }
 194         }
 195         clu0.count = count0;
 196         clu1.count = count1;
 197         clu0.calc_dist();
 198         clu1.calc_dist();
 199         clu0.dist + clu1.dist
 200     }
 201     #[allow(clippy::cyclomatic_complexity)]
 202     pub fn quantise(&mut self, src: &[Pixel], dst: &mut [[u8; 3]; 256]) {
 203         if src.len() < 3 {
 204             return;
 205         }
 206         let mut old_cb: [Pixel; 256] = [Pixel::default(); 256];
 207         let mut prev_dist = std::u64::MAX;
 208         let mut dist = std::u64::MAX / 2;
 209         let mut indices = Vec::with_capacity(src.len());
 210         let mut pixels = Vec::with_capacity(src.len());
 211         pixels.extend_from_slice(src);
 212         Self::sort(pixels.as_mut_slice(), |pix| pix.r);
 213         Self::sort(pixels.as_mut_slice(), |pix| pix.g);
 214         Self::sort(pixels.as_mut_slice(), |pix| pix.b);
 215         let mut entries = Vec::with_capacity(pixels.len() / 2);
 216         let mut lastval = pixels[0];
 217         let mut run = 1;
 218         for pix in pixels.iter().skip(1) {
 219             if &lastval == pix {
 220                 run += 1;
 221             } else {
 222                 entries.push(Entry { pix: lastval, count: run });
 223                 lastval = *pix;
 224                 run = 1;
 225             }
 226         }
 227         entries.push(Entry { pix: lastval, count: run });
 228         drop(pixels);
 229
 230         let mut low_u:  Vec<usize> = Vec::with_capacity(256);
 231         let mut high_u: Vec<usize> = Vec::with_capacity(256);
 232         let mut rng = RNG::new();
 233         let mut iterations = 0usize;
 234         let mut do_elbg_step = true;
 235         while (iterations < 20) && (dist < prev_dist - prev_dist / 1000) {
 236             prev_dist = dist;
 237             for i in 0..256 {
 238                 old_cb[i] = self.clusters[i].centroid;
 239                 self.clusters[i].reset();
 240             }
 241             // put pixels into the nearest clusters
 242             indices.truncate(0);
 243             for entry in entries.iter() {
 244                 let mut bestidx = 0;
 245                 let mut bestdist = std::u32::MAX;
 246                 for (i, cluster) in self.clusters.iter().enumerate() {
 247                     let dist = entry.pix.dist(cluster.centroid);
 248                     if bestdist > dist {
 249                         bestdist = dist;
 250                         bestidx = i;
 251                         if dist == 0 {
 252                             break;
 253                         }
 254                     }
 255                 }
 256                 indices.push(bestidx);
 257                 self.clusters[bestidx].add_pixel(entry);
 258             }
 259             // calculate params
 260             for cluster in self.clusters.iter_mut() {
 261                 cluster.calc_centroid();
 262             }
 263             dist = 0;
 264             for (idx, entry) in indices.iter().zip(entries.iter()) {
 265                 self.clusters[*idx].add_dist(entry);
 266             }
 267             for cluster in self.clusters.iter_mut() {
 268                 cluster.calc_dist();
 269                 dist += cluster.dist;
 270             }
 271
 272             let dmean = dist / 256;
 273             low_u.truncate(0);
 274             high_u.truncate(0);
 275             let mut used = [false; 256];
 276             for (i, cluster) in self.clusters.iter().enumerate() {
 277                 if cluster.dist < dmean {
 278                     low_u.push(i);
 279                 } else if cluster.dist > dmean * 2 {
 280                     high_u.push(i);
 281                     used[i] = true;
 282                 }
 283             }
 284
 285             if do_elbg_step {
 286                 do_elbg_step = false;
 287                 for low_idx in low_u.iter() {
 288                     if high_u.is_empty() {
 289                         break;
 290                     }
 291                     let high_idx_idx = (rng.next() as usize) % high_u.len();
 292                     let high_idx = high_u[high_idx_idx];
 293                     let mut closest_idx = *low_idx;
 294                     let mut closest_dist = std::u32::MAX;
 295                     let low_centr = self.clusters[*low_idx].centroid;
 296                     for i in 0..256 {//low_u.iter() {
 297                         if i == *low_idx || used[i] {
 298                             continue;
 299                         }
 300                         let dist = self.clusters[i].centroid.dist(low_centr);
 301                         if closest_dist > dist {
 302                             closest_dist = dist;
 303                             closest_idx  = i;
 304                         }
 305                     }
 306                     if closest_idx == *low_idx {
 307                         continue;
 308                     }
 309                     let old_dist = self.clusters[*low_idx].dist + self.clusters[closest_idx].dist + self.clusters[high_idx].dist;
 310                     let old_centr = self.old_centre(*low_idx, closest_idx, entries.as_slice(), indices.as_slice());
 311                     let ret = Self::new_split(high_idx, entries.as_slice(), indices.as_slice());
 312                     if ret.is_none() {
 313                         continue;
 314                     }
 315                     let (centr0, centr1) = ret.unwrap();
 316                     let dist_o = if old_dist > self.clusters[high_idx].dist {
 317                             Self::estimate_old(*low_idx, closest_idx, old_centr, entries.as_slice(), indices.as_slice())
 318                         } else { 0 };
 319                     let dist_n = Self::estimate_new(centr0, centr1, high_idx, entries.as_slice(), indices.as_slice());
 320                     if dist_o + dist_n < old_dist {
 321                         self.clusters[*low_idx   ].centroid = old_centr;
 322                         self.clusters[closest_idx].centroid = centr0;
 323                         self.clusters[high_idx   ].centroid = centr1;
 324                         used[*low_idx]    = true;
 325                         used[closest_idx] = true;
 326                         used[high_idx]    = true;
 327                         high_u.remove(high_idx_idx);
 328                         do_elbg_step = true;
 329                     }
 330                 }
 331             }
 332             iterations += 1;
 333         }
 334         if dist < prev_dist {
 335             for i in 0..256 {
 336                 old_cb[i] = self.clusters[i].centroid;
 337             }
 338         }
 339         for i in 0..256 {
 340             dst[i][0] = old_cb[i].r;
 341             dst[i][1] = old_cb[i].g;
 342             dst[i][2] = old_cb[i].b;
 343         }
 344     }
 345 }