From 4b459d0b9b76cb51c1029e6d1ffb17cf5f2d44d9 Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Wed, 13 May 2020 14:27:39 +0200 Subject: [PATCH] core/scale: add conversion into paletted format --- nihav-core/src/scale/colorcvt.rs | 1 + nihav-core/src/scale/mod.rs | 33 ++- nihav-core/src/scale/palette/elbg.rs | 344 ++++++++++++++++++++++ nihav-core/src/scale/palette/mediancut.rs | 162 ++++++++++ nihav-core/src/scale/palette/mod.rs | 236 +++++++++++++++ nihav-core/src/scale/palette/neuquant.rs | 133 +++++++++ nihav-core/src/scale/palette/palettise.rs | 185 ++++++++++++ nihav-core/src/scale/repack.rs | 1 + 8 files changed, 1094 insertions(+), 1 deletion(-) create mode 100644 nihav-core/src/scale/palette/elbg.rs create mode 100644 nihav-core/src/scale/palette/mediancut.rs create mode 100644 nihav-core/src/scale/palette/mod.rs create mode 100644 nihav-core/src/scale/palette/neuquant.rs create mode 100644 nihav-core/src/scale/palette/palettise.rs diff --git a/nihav-core/src/scale/colorcvt.rs b/nihav-core/src/scale/colorcvt.rs index d7ea7c8..351d79c 100644 --- a/nihav-core/src/scale/colorcvt.rs +++ b/nihav-core/src/scale/colorcvt.rs @@ -253,6 +253,7 @@ impl YuvToRgb { impl Kernel for YuvToRgb { fn init(&mut self, in_fmt: &ScaleInfo, dest_fmt: &ScaleInfo) -> ScaleResult { let mut df = dest_fmt.fmt; + df.palette = false; //todo coeff selection make_yuv2rgb(YUV_PARAMS[2][0], YUV_PARAMS[2][1], &mut self.matrix); if let ColorModel::YUV(yuvsm) = in_fmt.fmt.get_model() { diff --git a/nihav-core/src/scale/mod.rs b/nihav-core/src/scale/mod.rs index 0ba1e21..e66d393 100644 --- a/nihav-core/src/scale/mod.rs +++ b/nihav-core/src/scale/mod.rs @@ -23,6 +23,10 @@ mod colorcvt; mod repack; mod scale; +mod palette; + +pub use crate::scale::palette::{palettise_frame, QuantisationMode, PaletteSearchMode}; + /// Image format information used by the converter. #[derive(Clone,Copy,PartialEq)] pub struct ScaleInfo { @@ -84,6 +88,7 @@ const KERNELS: &[KernelDesc] = &[ KernelDesc { name: "pack", create: repack::create_pack }, KernelDesc { name: "unpack", create: repack::create_unpack }, KernelDesc { name: "depal", create: repack::create_depal }, + KernelDesc { name: "palette", create: palette::create_palettise }, KernelDesc { name: "scale", create: scale::create_scale }, KernelDesc { name: "rgb_to_yuv", create: colorcvt::create_rgb2yuv }, KernelDesc { name: "yuv_to_rgb", create: colorcvt::create_yuv2rgb }, @@ -227,6 +232,7 @@ println!("convert {} -> {}", ifmt, ofmt); let needs_convert = inname != outname; let scale_before_cvt = is_better_fmt(&ifmt, &ofmt) && needs_convert && (ofmt.fmt.get_max_subsampling() == 0); + let needs_palettise = ofmt.fmt.palette; //todo stages for model and gamma conversion let mut stages: Option = None; @@ -264,12 +270,18 @@ println!("[adding scale]"); cur_fmt = new_stage.fmt_out; add_stage!(stages, new_stage); } - if needs_pack { + if needs_pack && !needs_palettise { println!("[adding pack]"); let new_stage = Stage::new("pack", &cur_fmt, &ofmt)?; //cur_fmt = new_stage.fmt_out; add_stage!(stages, new_stage); } + if needs_palettise { +println!("[adding palettise]"); + let new_stage = Stage::new("palette", &cur_fmt, &ofmt)?; + //cur_fmt = new_stage.fmt_out; + add_stage!(stages, new_stage); + } if let Some(ref mut head) = stages { head.drop_last_tmp(); @@ -479,4 +491,23 @@ mod test { assert_eq!(odata[uoff], 154); assert_eq!(odata[voff], 103); } + #[test] + fn test_scale_and_convert_to_pal() { + let mut in_pic = alloc_video_buffer(NAVideoInfo::new(7, 3, false, YUV420_FORMAT), 3).unwrap(); + fill_pic(&mut in_pic, 142); + let mut out_pic = alloc_video_buffer(NAVideoInfo::new(4, 4, false, PAL8_FORMAT), 0).unwrap(); + fill_pic(&mut out_pic, 0); + let ifmt = get_scale_fmt_from_pic(&in_pic); + let ofmt = get_scale_fmt_from_pic(&out_pic); + let mut scaler = NAScale::new(ifmt, ofmt).unwrap(); + scaler.convert(&in_pic, &mut out_pic).unwrap(); + let obuf = out_pic.get_vbuf().unwrap(); + let dataoff = obuf.get_offset(0); + let paloff = obuf.get_offset(1); + let odata = obuf.get_data(); + assert_eq!(odata[dataoff], 0); + assert_eq!(odata[paloff], 157); + assert_eq!(odata[paloff + 1], 99); + assert_eq!(odata[paloff + 2], 170); + } } diff --git a/nihav-core/src/scale/palette/elbg.rs b/nihav-core/src/scale/palette/elbg.rs new file mode 100644 index 0000000..0926901 --- /dev/null +++ b/nihav-core/src/scale/palette/elbg.rs @@ -0,0 +1,344 @@ +use super::Pixel; + +struct RNG { + seed: u16, +} + +impl RNG { + fn new() -> Self { Self { seed: 0x1234 } } + fn next(&mut self) -> u8 { + if (self.seed & 0x8000) != 0 { + self.seed = (self.seed & 0x7FFF) * 2 ^ 0x1B2B; + } else { + self.seed <<= 1; + } + self.seed as u8 + } +} + +#[derive(Default,Clone,Copy,PartialEq,Debug)] +struct Entry { + pix: Pixel, + count: u64, +} + +struct Cluster { + centroid: Pixel, + dist: u64, + count: u64, + sum_r: u64, + sum_g: u64, + sum_b: u64, +} + +impl Cluster { + fn new(centroid: Pixel) -> Self { + Self { + centroid, + dist: 0, + count: 0, + sum_r: 0, + sum_g: 0, + sum_b: 0, + } + } + fn reset(&mut self) { + self.count = 0; + self.sum_r = 0; + self.sum_g = 0; + self.sum_b = 0; + self.dist = 0; + } + fn add_pixel(&mut self, entry: &Entry) { + self.sum_r += u64::from(entry.pix.r) * entry.count; + self.sum_g += u64::from(entry.pix.g) * entry.count; + self.sum_b += u64::from(entry.pix.b) * entry.count; + self.count += entry.count; + } + fn add_dist(&mut self, entry: &Entry) { + self.dist += u64::from(self.centroid.dist(entry.pix)) * entry.count; + } + fn calc_centroid(&mut self) { + if self.count != 0 { + self.centroid.r = ((self.sum_r + self.count / 2) / self.count) as u8; + self.centroid.g = ((self.sum_g + self.count / 2) / self.count) as u8; + self.centroid.b = ((self.sum_b + self.count / 2) / self.count) as u8; + } + } + fn calc_dist(&mut self) { + if self.count != 0 { + self.dist = (self.dist + self.count / 2) / self.count; + } + } +} + +pub struct ELBG { + clusters: Vec, +} + +impl ELBG { + #[allow(dead_code)] + pub fn new(initial_pal: &[[u8; 3]; 256]) -> Self { + let mut clusters = Vec::with_capacity(256); + for i in 0..256 { + let pix = Pixel { r: initial_pal[i][0], g: initial_pal[i][1], b: initial_pal[i][2] }; + let cluster = Cluster::new(pix); + clusters.push(cluster); + } + Self { + clusters, + } + } + #[allow(dead_code)] + pub fn new_random() -> Self { + let mut rng = RNG::new(); + let mut clusters = Vec::with_capacity(256); + for _ in 0..256 { + let pix = Pixel { r: rng.next(), g: rng.next(), b: rng.next() }; + let cluster = Cluster::new(pix); + clusters.push(cluster); + } + Self { + clusters, + } + } + fn sort(arr: &mut [Pixel], idx: F) where F: Fn(&Pixel) -> u8 { + let mut dst = vec![Pixel::default(); arr.len()]; + let mut counts = [0; 256]; + for pix in arr.iter() { + counts[idx(pix) as usize] += 1; + } + let mut last = counts[0]; + counts[0] = 0; + for count in counts.iter_mut().skip(1) { + let plast = last; + last += *count; + *count = plast; + } + for pix in arr.iter() { + let bucket = idx(pix) as usize; + dst[counts[bucket]] = *pix; + counts[bucket] += 1; + } + arr.copy_from_slice(dst.as_slice()); + } + fn new_split(old_index: usize, entries: &[Entry], indices: &[usize]) -> Option<(Pixel, Pixel)> { + let mut max = Pixel { r: 0, g: 0, b: 0 }; + let mut min = Pixel { r: 255, g: 255, b: 255 }; + let mut found = false; + for (entry, idx) in entries.iter().zip(indices) { + if *idx == old_index { + max = max.max(entry.pix); + min = min.min(entry.pix); + found = true; + } + } + if !found { + return None; + } + let dr = max.r - min.r; + let dg = max.g - min.g; + let db = max.b - min.b; + let cent0 = Pixel { r: min.r + dr / 3, g: min.g + dg / 3, b: min.b + db / 3 }; + let cent1 = Pixel { r: max.r - dr / 3, g: max.g - dg / 3, b: max.b - db / 3 }; + Some((cent0, cent1)) + } + fn old_centre(&self, old_index1: usize, old_index2: usize, entries: &[Entry], indices: &[usize]) -> Pixel { + let mut max = Pixel { r: 0, g: 0, b: 0 }; + let mut min = Pixel { r: 255, g: 255, b: 255 }; + let mut found = false; + for (entry, idx) in entries.iter().zip(indices) { + if *idx == old_index1 || *idx == old_index2 { + max = max.max(entry.pix); + min = min.min(entry.pix); + found = true; + } + } + if !found { + max = self.clusters[old_index1].centroid.max(self.clusters[old_index2].centroid); + min = self.clusters[old_index1].centroid.min(self.clusters[old_index2].centroid); + } + let dr = max.r - min.r; + let dg = max.g - min.g; + let db = max.b - min.b; + Pixel { r: min.r + dr / 2, g: min.g + dg / 2, b: min.b + db / 2 } + } + fn estimate_old(old_idx0: usize, old_idx1: usize, c: Pixel, entries: &[Entry], indices: &[usize]) -> u64 { + let mut clu = Cluster::new(c); + let mut count = 0; + for (entry, idx) in entries.iter().zip(indices) { + if *idx == old_idx0 || *idx == old_idx1 { + clu.add_dist(entry); + count += entry.count; + } + } + clu.count = count; + clu.calc_dist(); + clu.dist + } + fn estimate_new(c0: Pixel, c1: Pixel, old_idx: usize, entries: &[Entry], indices: &[usize]) -> u64 { + let mut clu0 = Cluster::new(c0); + let mut clu1 = Cluster::new(c1); + let mut count0 = 0; + let mut count1 = 0; + for (entry, idx) in entries.iter().zip(indices) { + if *idx == old_idx { + if c0.dist(entry.pix) < c1.dist(entry.pix) { + clu0.add_dist(entry); + count0 += entry.count; + } else { + clu1.add_dist(entry); + count1 += entry.count; + } + } + } + clu0.count = count0; + clu1.count = count1; + clu0.calc_dist(); + clu1.calc_dist(); + clu0.dist + clu1.dist + } + pub fn quantise(&mut self, src: &[Pixel], dst: &mut [[u8; 3]; 256]) { + if src.len() < 3 { + return; + } + let mut old_cb: [Pixel; 256] = [Pixel::default(); 256]; + let mut prev_dist = std::u64::MAX; + let mut dist = std::u64::MAX / 2; + let mut indices = Vec::with_capacity(src.len()); + let mut pixels = Vec::with_capacity(src.len()); + pixels.extend_from_slice(src); + Self::sort(pixels.as_mut_slice(), |pix| pix.r); + Self::sort(pixels.as_mut_slice(), |pix| pix.g); + Self::sort(pixels.as_mut_slice(), |pix| pix.b); + let mut entries = Vec::with_capacity(pixels.len() / 2); + let mut lastval = pixels[0]; + let mut run = 1; + for pix in pixels.iter().skip(1) { + if &lastval == pix { + run += 1; + } else { + entries.push(Entry { pix: lastval, count: run }); + lastval = *pix; + run = 1; + } + } + entries.push(Entry { pix: lastval, count: run }); + drop(pixels); + + let mut low_u: Vec = Vec::with_capacity(256); + let mut high_u: Vec = Vec::with_capacity(256); + let mut rng = RNG::new(); + let mut iterations = 0usize; + let mut do_elbg_step = true; + while (iterations < 20) && (dist < prev_dist - prev_dist / 1000) { + prev_dist = dist; + for i in 0..256 { + old_cb[i] = self.clusters[i].centroid; + self.clusters[i].reset(); + } + // put pixels into the nearest clusters + indices.truncate(0); + for entry in entries.iter() { + let mut bestidx = 0; + let mut bestdist = std::u32::MAX; + for (i, cluster) in self.clusters.iter().enumerate() { + let dist = entry.pix.dist(cluster.centroid); + if bestdist > dist { + bestdist = dist; + bestidx = i; + if dist == 0 { + break; + } + } + } + indices.push(bestidx); + self.clusters[bestidx].add_pixel(entry); + } + // calculate params + for cluster in self.clusters.iter_mut() { + cluster.calc_centroid(); + } + dist = 0; + for (idx, entry) in indices.iter().zip(entries.iter()) { + self.clusters[*idx].add_dist(entry); + } + for cluster in self.clusters.iter_mut() { + cluster.calc_dist(); + dist += cluster.dist; + } + + let dmean = dist / 256; + low_u.truncate(0); + high_u.truncate(0); + let mut used = [false; 256]; + for (i, cluster) in self.clusters.iter().enumerate() { + if cluster.dist < dmean { + low_u.push(i); + } else if cluster.dist > dmean * 2 { + high_u.push(i); + used[i] = true; + } + } + + if do_elbg_step { + do_elbg_step = false; + for low_idx in low_u.iter() { + if high_u.len() == 0 { + break; + } + let high_idx_idx = (rng.next() as usize) % high_u.len(); + let high_idx = high_u[high_idx_idx]; + let mut closest_idx = *low_idx; + let mut closest_dist = std::u32::MAX; + let low_centr = self.clusters[*low_idx].centroid; + for i in 0..256 {//low_u.iter() { + if i == *low_idx || used[i] { + continue; + } + let dist = self.clusters[i].centroid.dist(low_centr); + if closest_dist > dist { + closest_dist = dist; + closest_idx = i; + } + } + if closest_idx == *low_idx { + continue; + } + let old_dist = self.clusters[*low_idx].dist + self.clusters[closest_idx].dist + self.clusters[high_idx].dist; + let old_centr = self.old_centre(*low_idx, closest_idx, entries.as_slice(), indices.as_slice()); + let ret = Self::new_split(high_idx, entries.as_slice(), indices.as_slice()); + if ret.is_none() { + continue; + } + let (centr0, centr1) = ret.unwrap(); + let dist_o = if old_dist > self.clusters[high_idx].dist { + Self::estimate_old(*low_idx, closest_idx, old_centr, entries.as_slice(), indices.as_slice()) + } else { 0 }; + let dist_n = Self::estimate_new(centr0, centr1, high_idx, entries.as_slice(), indices.as_slice()); + if dist_o + dist_n < old_dist { + self.clusters[*low_idx ].centroid = old_centr; + self.clusters[closest_idx].centroid = centr0; + self.clusters[high_idx ].centroid = centr1; + used[*low_idx] = true; + used[closest_idx] = true; + used[high_idx] = true; + high_u.remove(high_idx_idx); + do_elbg_step = true; + } + } + } + iterations += 1; + } + if dist < prev_dist { + for i in 0..256 { + old_cb[i] = self.clusters[i].centroid; + } + } + for i in 0..256 { + dst[i][0] = old_cb[i].r; + dst[i][1] = old_cb[i].g; + dst[i][2] = old_cb[i].b; + } + } +} diff --git a/nihav-core/src/scale/palette/mediancut.rs b/nihav-core/src/scale/palette/mediancut.rs new file mode 100644 index 0000000..497cc4a --- /dev/null +++ b/nihav-core/src/scale/palette/mediancut.rs @@ -0,0 +1,162 @@ +use super::Pixel; + +struct VQBox<'a> { + pixels: &'a mut [Pixel], + max: Pixel, + min: Pixel, +} + +fn avg_u8(a: u8, b: u8) -> u8 { + (a & b) + ((a ^ b) >> 1) +} + +impl<'a> VQBox<'a> { + fn new(pixels: &'a mut [Pixel]) -> Self { + let mut max = Pixel{ r: 0, g: 0, b: 0 }; + let mut min = Pixel{ r: 255, g: 255, b: 255 }; + for pix in pixels.iter() { + max = max.max(*pix); + min = min.min(*pix); + } + Self { pixels, max, min } + } + fn can_split(&self) -> bool { + let dr = self.max.r - self.min.r; + let dg = self.max.g - self.min.g; + let db = self.max.b - self.min.b; + (dr | dg | db) != 0 + } + fn sort(arr: &mut [Pixel], idx: F) -> usize where F: Fn(&Pixel) -> usize { + let mut buckets: Vec> = Vec::with_capacity(256); + let mut counts = [0; 256]; + for pix in arr.iter() { + counts[idx(pix)] += 1; + } + for i in 0..256 { + buckets.push(Vec::with_capacity(counts[i])); + } + for pix in arr.iter() { + buckets[idx(pix)].push(*pix); + } + let mut start = 0; + let mut pivot = 0; + let mut cand1 = 0; + for bucket in buckets.iter() { + let end = bucket.len(); + if (start > 0) && (start <= arr.len() / 2) { + pivot = start + end; + cand1 = start; + } + (&mut arr[start..][..end]).copy_from_slice(bucket.as_slice()); + start += end; + } + for bucket in buckets.iter() { + if pivot != 0 { + break; + } + pivot += bucket.len(); + } + if pivot < arr.len() || cand1 == 0 { + pivot + } else { + cand1 + } + } + fn calc_min_and_max(pixels: &[Pixel]) -> (Pixel, Pixel) { + let mut max = Pixel{ r: 0, g: 0, b: 0 }; + let mut min = Pixel{ r: 255, g: 255, b: 255 }; + for pix in pixels.iter() { + max = max.max(*pix); + min = min.min(*pix); + } + (min, max) + } + fn split(self) -> (VQBox<'a>, VQBox<'a>) { + let dr = self.max.r - self.min.r; + let dg = self.max.g - self.min.g; + let db = self.max.b - self.min.b; + + let box0; + let box1; + if (dr > dg) && (dr > db) { + let pivot = Self::sort(self.pixels, |pix| pix.r as usize); + let (part0, part1) = self.pixels.split_at_mut(pivot); + let (min0, max0) = Self::calc_min_and_max(part0); + let (min1, max1) = Self::calc_min_and_max(part1); + box0 = VQBox { pixels: part0, max: max0, min: min0 }; + box1 = VQBox { pixels: part1, max: max1, min: min1 }; + } else if (db > dr) && (db > dg) { + let pivot = Self::sort(self.pixels, |pix| pix.b as usize); + let (part0, part1) = self.pixels.split_at_mut(pivot); + let (min0, max0) = Self::calc_min_and_max(part0); + let (min1, max1) = Self::calc_min_and_max(part1); + box0 = VQBox { pixels: part0, max: max0, min: min0 }; + box1 = VQBox { pixels: part1, max: max1, min: min1 }; + } else { + let pivot = Self::sort(self.pixels, |pix| pix.g as usize); + let (part0, part1) = self.pixels.split_at_mut(pivot); + let (min0, max0) = Self::calc_min_and_max(part0); + let (min1, max1) = Self::calc_min_and_max(part1); + box0 = VQBox { pixels: part0, max: max0, min: min0 }; + box1 = VQBox { pixels: part1, max: max1, min: min1 }; + } + (box0, box1) + } +} + +pub fn quantise_median_cut(src: &[Pixel], pal: &mut [[u8; 3]; 256]) -> usize { + let mut pixels = Vec::with_capacity(src.len()); + pixels.extend_from_slice(src); + VQBox::sort(pixels.as_mut_slice(), |pix| pix.r as usize); + VQBox::sort(pixels.as_mut_slice(), |pix| pix.g as usize); + VQBox::sort(pixels.as_mut_slice(), |pix| pix.b as usize); + let box0 = VQBox::new(pixels.as_mut_slice()); + let mut boxes: Vec = Vec::with_capacity(256); + boxes.push(box0); + let mut changed = true; + while changed && boxes.len() < 256 { + let end = boxes.len(); + changed = false; + let mut split_largest = false; + for _ in 0..end { + let curbox = boxes.remove(0); + if curbox.can_split() { + let (box0, box1) = curbox.split(); + boxes.push(box0); + boxes.push(box1); + changed = true; + } else { + boxes.push(curbox); + split_largest = true; + break; + } + if boxes.len() == 256 { + break; + } + } + if split_largest { + let mut maxidx = 0; + let mut lcount = 0; + for (i, cbox) in boxes.iter().enumerate() { + if cbox.can_split() && cbox.pixels.len() > lcount { + lcount = cbox.pixels.len(); + maxidx = i; + } + } + if lcount > 0 { + let curbox = boxes.remove(maxidx); + let (box0, box1) = curbox.split(); + boxes.push(box0); + boxes.push(box1); + changed = true; + } + } + } + for (curbox, palentry) in boxes.iter().zip(pal.iter_mut()) { + palentry[0] = avg_u8(curbox.min.r, curbox.max.r); + palentry[1] = avg_u8(curbox.min.g, curbox.max.g); + palentry[2] = avg_u8(curbox.min.b, curbox.max.b); + } + + boxes.len() +} diff --git a/nihav-core/src/scale/palette/mod.rs b/nihav-core/src/scale/palette/mod.rs new file mode 100644 index 0000000..2005cd8 --- /dev/null +++ b/nihav-core/src/scale/palette/mod.rs @@ -0,0 +1,236 @@ +use crate::formats::*; +use super::*; +use super::kernel::Kernel; + +#[derive(Default,Clone,Copy,PartialEq,Debug)] +pub struct Pixel { + pub r: u8, + pub g: u8, + pub b: u8, +} + +impl Pixel { + #[allow(dead_code)] + fn new(src: &[u8]) -> Self { + Self { r: src[0], g: src[1], b: src[2] } + } + fn to_rgb(&self) -> [u8; 3] { + [self.r, self.g, self.b] + } + fn dist(&self, pix: Pixel) -> u32 { + let dr = i32::from(self.r) - i32::from(pix.r); + let dg = i32::from(self.g) - i32::from(pix.g); + let db = i32::from(self.b) - i32::from(pix.b); + (dr * dr + dg * dg + db * db) as u32 + } + fn min(&self, pix: Pixel) -> Pixel { + Pixel { r: self.r.min(pix.r), g: self.g.min(pix.g), b: self.b.min(pix.b) } + } + fn max(&self, pix: Pixel) -> Pixel { + Pixel { r: self.r.max(pix.r), g: self.g.max(pix.g), b: self.b.max(pix.b) } + } +} + +#[allow(dead_code)] +fn avg_u8(a: u8, b: u8) -> u8 { + (a & b) + ((a ^ b) >> 1) +} + +mod elbg; +mod mediancut; +mod neuquant; +mod palettise; + +//use elbg::ELBG; +//use mediancut::quantise_median_cut; +//use neuquant::NeuQuantQuantiser; + +#[derive(Clone,Copy,Debug,PartialEq)] +/// Palette quantisation algorithms. +pub enum QuantisationMode { + /// Median cut approach proposed by Paul Heckbert. + /// + /// This is moderately fast and moderately good. + MedianCut, + /// Enhanced LBG algorithm proposed by Giuseppe Patane and Marco Russo. + /// + /// This is slow but good method. + ELBG, + /// NeuQuant algorithm proposed by Anthony Dekker. + /// + /// It may operate on randomly subsampled image with subsampling factors 1-30. + /// This algorithm is fast especially with high subsampling factor but output palette may be far from optimal one. + NeuQuant(u8), +} + +impl Default for QuantisationMode { + fn default() -> Self { QuantisationMode::MedianCut } +} + +#[derive(Clone,Copy,Debug,PartialEq)] +/// Algorithms for seaching an appropriate palette entry for a given pixel. +pub enum PaletteSearchMode { + /// Full search (slowest). + Full, + /// Local search (faster but may be not so good). + Local, + /// k-d tree based one (the fastest but not so accurate). + KDTree, +} + +impl Default for PaletteSearchMode { + fn default() -> Self { PaletteSearchMode::Local } +} + +use crate::scale::palette::elbg::ELBG; +use crate::scale::palette::mediancut::quantise_median_cut; +use crate::scale::palette::neuquant::NeuQuantQuantiser; +use crate::scale::palette::palettise::*; + +fn palettise_frame_internal(pic_in: &NABufferType, pic_out: &mut NABufferType, qmode: QuantisationMode, palmode: PaletteSearchMode, pixels: &mut Vec) -> ScaleResult<()> { + if let (Some(ref sbuf), Some(ref mut dbuf)) = (pic_in.get_vbuf(), pic_out.get_vbuf()) { + let ioff = sbuf.get_offset(0); + let (w, h) = sbuf.get_dimensions(0); + let istride = sbuf.get_stride(0); + let ifmt = sbuf.get_info().get_format(); + let sdata1 = sbuf.get_data(); + let sdata = &sdata1[ioff..]; + + let doff = dbuf.get_offset(0); + let paloff = dbuf.get_offset(1); + let dstride = dbuf.get_stride(0); + let ofmt = dbuf.get_info().get_format(); + let dst = dbuf.get_data_mut().unwrap(); + + pixels.truncate(0); + if !ifmt.is_unpacked() { + let esize = ifmt.elem_size as usize; + let coffs = [ifmt.comp_info[0].unwrap().comp_offs as usize, ifmt.comp_info[1].unwrap().comp_offs as usize, ifmt.comp_info[2].unwrap().comp_offs as usize]; + for src in sdata.chunks(istride).take(h) { + for chunk in src.chunks_exact(esize).take(w) { + let pixel = Pixel{ r: chunk[coffs[0]], g: chunk[coffs[1]], b: chunk[coffs[2]] }; + pixels.push(pixel); + } + } + } else { + let mut roff = ioff; + let mut goff = sbuf.get_offset(1); + let mut boff = sbuf.get_offset(2); + let rstride = istride; + let gstride = sbuf.get_stride(1); + let bstride = sbuf.get_stride(2); + for _ in 0..h { + for x in 0..w { + let pixel = Pixel{ r: sdata[roff + x], g: sdata[goff + x], b: sdata[boff + x] }; + pixels.push(pixel); + } + roff += rstride; + goff += gstride; + boff += bstride; + } + } + let mut pal = [[0u8; 3]; 256]; + match qmode { + QuantisationMode::ELBG => { + let mut elbg = ELBG::new_random(); + elbg.quantise(pixels.as_slice(), &mut pal); + }, + QuantisationMode::MedianCut => { + quantise_median_cut(pixels.as_slice(), &mut pal); + }, + QuantisationMode::NeuQuant(fact) => { + let mut nq = NeuQuantQuantiser::new(fact as usize); + nq.learn(pixels.as_slice()); + nq.make_pal(&mut pal); + }, + }; + let esize = ofmt.elem_size as usize; + let coffs = [ofmt.comp_info[0].unwrap().comp_offs as usize, ofmt.comp_info[1].unwrap().comp_offs as usize, ofmt.comp_info[2].unwrap().comp_offs as usize]; + for (dpal, spal) in (&mut dst[paloff..]).chunks_mut(esize).zip(pal.iter()) { + dpal[coffs[0]] = spal[0]; + dpal[coffs[1]] = spal[1]; + dpal[coffs[2]] = spal[2]; + } + + let dst = &mut dst[doff..]; + match palmode { + PaletteSearchMode::Full => { + for (dline, sline) in dst.chunks_mut(dstride).take(h).zip(pixels.chunks(w)) { + for (didx, pix) in dline.iter_mut().take(w).zip(sline.iter()) { + let rgb = pix.to_rgb(); + *didx = find_nearest(&rgb, &pal) as u8; + } + } + }, + PaletteSearchMode::Local => { + let ls = LocalSearch::new(&pal); + for (dline, sline) in dst.chunks_mut(dstride).take(h).zip(pixels.chunks(w)) { + for (didx, pix) in dline.iter_mut().take(w).zip(sline.iter()) { + *didx = ls.search(pix.to_rgb()) as u8; + } + } + }, + PaletteSearchMode::KDTree => { + let kdtree = KDTree::new(&pal); + for (dline, sline) in dst.chunks_mut(dstride).take(h).zip(pixels.chunks(w)) { + for (didx, pix) in dline.iter_mut().take(w).zip(sline.iter()) { + *didx = kdtree.search(pix.to_rgb()) as u8; + } + } + }, + }; + Ok(()) + } else { + Err(ScaleError::InvalidArgument) + } +} + +/// Converts packed RGB frame into palettised one. +/// +/// This function can operate in several modes of both palette generation and colour substitution with palette indices. +/// Some may work fast but produce worse output image. +/// See [`QuantisationMode`] and [`PaletteSearchMode`] for details. +/// If you are not sure what to use there are `QuantisationMode::default()` and `PaletteSearchMode::default()`. +/// +/// [`QuantisationMode`]: ./enum.QuantisationMode.html +/// [`PaletteSearchMode`]: ./enum.PaletteSearchMode.html +pub fn palettise_frame(pic_in: &NABufferType, pic_out: &mut NABufferType, qmode: QuantisationMode, palmode: PaletteSearchMode) -> ScaleResult<()> { + let size; + if let Some(ref vbuf) = pic_in.get_vbuf() { +//todo check format for being packed RGB in and pal out + let (w, h) = vbuf.get_dimensions(0); + size = w * h; + } else { + return Err(ScaleError::InvalidArgument); + } + let mut pixels = Vec::with_capacity(size); + palettise_frame_internal(pic_in, pic_out, qmode, palmode, &mut pixels) +} + + +#[derive(Default)] +struct PalettiseKernel { + pixels: Vec, + qmode: QuantisationMode, + palmode: PaletteSearchMode, +} + +impl PalettiseKernel { + fn new() -> Self { Self::default() } +} + +impl Kernel for PalettiseKernel { + fn init(&mut self, in_fmt: &ScaleInfo, _dest_fmt: &ScaleInfo) -> ScaleResult { + self.pixels = Vec::with_capacity(in_fmt.width * in_fmt.height); + let res = alloc_video_buffer(NAVideoInfo::new(in_fmt.width, in_fmt.height, false, PAL8_FORMAT), 0); + if res.is_err() { return Err(ScaleError::AllocError); } + Ok(res.unwrap()) + } + fn process(&mut self, pic_in: &NABufferType, pic_out: &mut NABufferType) { + palettise_frame_internal(pic_in, pic_out, self.qmode, self.palmode, &mut self.pixels).unwrap(); + } +} + +pub fn create_palettise() -> Box { + Box::new(PalettiseKernel::new()) +} diff --git a/nihav-core/src/scale/palette/neuquant.rs b/nihav-core/src/scale/palette/neuquant.rs new file mode 100644 index 0000000..7b19eb5 --- /dev/null +++ b/nihav-core/src/scale/palette/neuquant.rs @@ -0,0 +1,133 @@ +use super::Pixel; + +pub struct NeuQuantQuantiser { + weights: [[f64; 3]; 256], + freq: [f64; 256], + bias: [f64; 256], + factor: usize, +} + +const SPECIAL_NODES: usize = 2; +impl NeuQuantQuantiser { + pub fn new(factor: usize) -> Self { + let mut weights = [[0.0; 3]; 256]; + if SPECIAL_NODES > 1 { + weights[1] = [255.0; 3]; // for white + } + for i in SPECIAL_NODES..256 { + let w = 255.0 * ((i - SPECIAL_NODES) as f64) / ((256 - SPECIAL_NODES) as f64); + weights[i] = [w, w, w]; + } + Self { + weights, + freq: [1.0 / 256.0; 256], + bias: [0.0; 256], + factor, + } + } + fn update_node(&mut self, idx: usize, clr: &[f64; 3], alpha: f64) { + self.weights[idx][0] -= alpha * (self.weights[idx][0] - clr[0]); + self.weights[idx][1] -= alpha * (self.weights[idx][1] - clr[1]); + self.weights[idx][2] -= alpha * (self.weights[idx][2] - clr[2]); + } + fn update_neighbours(&mut self, idx: usize, clr: &[f64; 3], alpha: f64, radius: usize) { + let low = idx.saturating_sub(radius).max(SPECIAL_NODES - 1); + let high = (idx + radius).min(self.weights.len() - 1); + + let mut idx0 = idx + 1; + let mut idx1 = idx - 1; + let mut range = 0; + let sqradius = (radius * radius) as f64; + while (idx0 < high) || (idx1 > low) { + let sqrng = (range * range) as f64; + let a = alpha * (sqradius - sqrng) / sqradius; + range += 1; + if idx0 < high { + self.update_node(idx0, clr, a); + idx0 += 1; + } + if idx1 > low { + self.update_node(idx1, clr, a); + idx1 -= 1; + } + } + } + fn find_node(&mut self, clr: &[f64; 3]) -> usize { + for i in 0..SPECIAL_NODES { + if &self.weights[i] == clr { + return i; + } + } + let mut bestdist = std::f64::MAX; + let mut distidx = 0; + let mut bestbias = std::f64::MAX; + let mut biasidx = 0; + for i in SPECIAL_NODES..256 { + let dist = (self.weights[i][0] - clr[0]) * (self.weights[i][0] - clr[0]) + + (self.weights[i][1] - clr[1]) * (self.weights[i][1] - clr[1]) + + (self.weights[i][2] - clr[2]) * (self.weights[i][2] - clr[2]); + if bestdist > dist { + bestdist = dist; + distidx = i; + } + let biasdiff = dist - self.bias[i]; + if bestbias > biasdiff { + bestbias = biasdiff; + biasidx = i; + } + self.freq[i] -= self.freq[i] / 1024.0; + self.bias[i] += self.freq[i]; + } + self.freq[distidx] += 1.0 / 1024.0; + self.bias[distidx] -= 1.0; + biasidx + } + pub fn learn(&mut self, src: &[Pixel]) { + let mut bias_radius = (256 / 8) << 6; + let alphadec = (30 + (self.factor - 1) / 3) as f64; + let initial_alpha = (1 << 10) as f64; + + let npixels = src.len(); + + let mut radius = bias_radius >> 6; + if radius == 1 { radius = 0 }; + let samples = npixels / self.factor; + let delta = samples / 100; + let mut alpha = initial_alpha; + + let mut pos = 0; + const PRIMES: [usize; 4] = [ 499, 491, 487, 503 ]; + let mut step = PRIMES[3]; + for prime in PRIMES.iter().rev() { + if npixels % *prime != 0 { + step = *prime; + } + } + + for i in 0..samples { + let clr = [src[pos].r as f64, src[pos].g as f64, src[pos].b as f64]; + let idx = self.find_node(&clr); + if idx >= SPECIAL_NODES { + let new_alpha = alphadec / initial_alpha; + self.update_node(idx, &clr, new_alpha); + if radius > 0 { + self.update_neighbours(idx, &clr, new_alpha, radius); + } + } + pos = (pos + step) % npixels; + if (i + 1) % delta == 0 { + alpha -= alpha / alphadec; + bias_radius -= bias_radius / 30; + radius = bias_radius >> 6; + if radius == 1 { radius = 0 }; + } + } + } + pub fn make_pal(&self, pal: &mut [[u8; 3]; 256]) { + for (pal, node) in pal.iter_mut().zip(self.weights.iter()) { + pal[0] = (node[0] + 0.5).max(0.0).min(255.0) as u8; + pal[1] = (node[1] + 0.5).max(0.0).min(255.0) as u8; + pal[2] = (node[2] + 0.5).max(0.0).min(255.0) as u8; + } + } +} diff --git a/nihav-core/src/scale/palette/palettise.rs b/nihav-core/src/scale/palette/palettise.rs new file mode 100644 index 0000000..54467d3 --- /dev/null +++ b/nihav-core/src/scale/palette/palettise.rs @@ -0,0 +1,185 @@ +pub fn find_nearest(pix: &[u8], pal: &[[u8; 3]; 256]) -> usize { + let mut bestidx = 0; + let mut bestdist = std::i32::MAX; + + for (idx, entry) in pal.iter().enumerate() { + let dist0 = i32::from(pix[0]) - i32::from(entry[0]); + let dist1 = i32::from(pix[1]) - i32::from(entry[1]); + let dist2 = i32::from(pix[2]) - i32::from(entry[2]); + if (dist0 | dist1 | dist2) == 0 { + return idx; + } + let dist = dist0 * dist0 + dist1 * dist1 + dist2 * dist2; + if bestdist > dist { + bestdist = dist; + bestidx = idx; + } + } + bestidx +} + +pub struct LocalSearch { + pal: [[u8; 3]; 256], + db: Vec>, +} + +impl LocalSearch { + fn quant(key: [u8; 3]) -> usize { + (((key[0] >> 3) as usize) << 10) | + (((key[1] >> 3) as usize) << 5) | + ((key[2] >> 3) as usize) + } + pub fn new(in_pal: &[[u8; 3]; 256]) -> Self { + let mut db = Vec::with_capacity(1 << 15); + let pal = *in_pal; + for _ in 0..(1 << 15) { + db.push(Vec::new()); + } + for (i, palentry) in pal.iter().enumerate() { + let r0 = (palentry[0] >> 3) as usize; + let g0 = (palentry[1] >> 3) as usize; + let b0 = (palentry[2] >> 3) as usize; + for r in r0.saturating_sub(1)..=(r0 + 1).min(31) { + for g in g0.saturating_sub(1)..=(g0 + 1).min(31) { + for b in b0.saturating_sub(1)..=(b0 + 1).min(31) { + let idx = (r << 10) | (g << 5) | b; + db[idx].push([palentry[0], palentry[1], palentry[2], i as u8]); + } + } + } + } + Self { pal, db } + } + fn dist(a: &[u8; 4], b: [u8; 3]) -> u32 { + let d0 = i32::from(a[0]) - i32::from(b[0]); + let d1 = i32::from(a[1]) - i32::from(b[1]); + let d2 = i32::from(a[2]) - i32::from(b[2]); + (d0 * d0 + d1 * d1 + d2 * d2) as u32 + } + pub fn search(&self, pix: [u8; 3]) -> usize { + let idx = Self::quant(pix); + let mut best_dist = std::u32::MAX; + let mut best_idx = 0; + let mut count = 0; + for clr in self.db[idx].iter() { + let dist = Self::dist(clr, pix); + count += 1; + if best_dist > dist { + best_dist = dist; + best_idx = clr[3] as usize; + if dist == 0 { break; } + } + } + if count > 0 { + best_idx + } else { + find_nearest(&pix, &self.pal) + } + } +} + +struct KDNode { + key: [u8; 3], + comp: u8, + idx: u8, + child0: usize, + child1: usize, +} + +pub struct KDTree { + nodes: Vec, +} + +fn avg_u8(a: u8, b: u8) -> u8 { + (a & b) + ((a ^ b) >> 1) +} + +impl KDTree { + pub fn new(pal: &[[u8; 3]; 256]) -> Self { + let mut npal = [[0; 4]; 256]; + for i in 0..256 { + npal[i][0] = pal[i][0]; + npal[i][1] = pal[i][1]; + npal[i][2] = pal[i][2]; + npal[i][3] = i as u8; + } + let mut tree = Self { nodes: Vec::with_capacity(512) }; + tree.build(&mut npal, 0, 256, 1024, false); + tree + } + fn build(&mut self, pal: &mut [[u8; 4]; 256], start: usize, end: usize, root: usize, child0: bool) { + if start + 1 == end { + let key = [pal[start][0], pal[start][1], pal[start][2]]; + let newnode = KDNode { key, comp: 0, idx: pal[start][3], child0: 0, child1: 0 }; + let cur_node = self.nodes.len(); + self.nodes.push(newnode); + if child0 { + self.nodes[root].child0 = cur_node; + } else { + self.nodes[root].child1 = cur_node; + } + return; + } + let mut min = [255u8; 3]; + let mut max = [0u8; 3]; + for i in start..end { + for comp in 0..3 { + min[comp] = min[comp].min(pal[i][comp]); + max[comp] = max[comp].max(pal[i][comp]); + } + } + let dr = max[0] - min[0]; + let dg = max[1] - min[1]; + let db = max[2] - min[2]; + let med = [avg_u8(min[0], max[0]), avg_u8(min[1], max[1]), avg_u8(min[2], max[2])]; + let comp = if dr > dg && dr > db { + 0 + } else if db > dr && db > dg { + 2 + } else { + 1 + }; + let pivot = Self::reorder(&mut pal[start..end], comp, med[comp]) + start; + let newnode = KDNode { key: med, comp: comp as u8, idx: 0, child0: 0, child1: 0 }; + let cur_node = self.nodes.len(); + self.nodes.push(newnode); + if root != 1024 { + if child0 { + self.nodes[root].child0 = cur_node; + } else { + self.nodes[root].child1 = cur_node; + } + } + self.build(pal, start, pivot, cur_node, true); + self.build(pal, pivot, end, cur_node, false); + } + fn reorder(pal: &mut[[u8; 4]], comp: usize, med: u8) -> usize { + let mut start = 0; + let mut end = pal.len() - 1; + while start < end { + while start < end && pal[start][comp] <= med { + start += 1; + } + while start < end && pal[end][comp] > med { + end -= 1; + } + if start < end { + pal.swap(start, end); + start += 1; + end -= 1; + } + } + start + } + pub fn search(&self, pix: [u8; 3]) -> usize { + let mut idx = 0; + loop { + let cnode = &self.nodes[idx]; + if cnode.child0 == 0 { + return cnode.idx as usize; + } + let nidx = if cnode.key[cnode.comp as usize] >= pix[cnode.comp as usize] { cnode.child0 } else { cnode.child1 }; + idx = nidx; + } + } +} diff --git a/nihav-core/src/scale/repack.rs b/nihav-core/src/scale/repack.rs index fd59e78..faa5fe7 100644 --- a/nihav-core/src/scale/repack.rs +++ b/nihav-core/src/scale/repack.rs @@ -146,6 +146,7 @@ impl Kernel for UnpackKernel { for i in 0..self.ncomps { df.comp_info[i] = chr[i]; } + df.palette = false; println!(" [intermediate format {}]", df); let res = alloc_video_buffer(NAVideoInfo::new(in_fmt.width, in_fmt.height, false, df), 3); if res.is_err() { return Err(ScaleError::AllocError); } -- 2.30.2