impl Kernel for YuvToRgb {
fn init(&mut self, in_fmt: &ScaleInfo, dest_fmt: &ScaleInfo) -> ScaleResult<NABufferType> {
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() {
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 {
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 },
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<Stage> = None;
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();
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);
+ }
}
--- /dev/null
+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<Cluster>,
+}
+
+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<F>(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<usize> = Vec::with_capacity(256);
+ let mut high_u: Vec<usize> = 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;
+ }
+ }
+}
--- /dev/null
+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<F>(arr: &mut [Pixel], idx: F) -> usize where F: Fn(&Pixel) -> usize {
+ let mut buckets: Vec<Vec<Pixel>> = 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<VQBox> = 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()
+}
--- /dev/null
+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<Pixel>) -> 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<Pixel>,
+ 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<NABufferType> {
+ 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<dyn Kernel> {
+ Box::new(PalettiseKernel::new())
+}
--- /dev/null
+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;
+ }
+ }
+}
--- /dev/null
+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<Vec<[u8; 4]>>,
+}
+
+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<KDNode>,
+}
+
+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;
+ }
+ }
+}
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); }
projects / nihav.git / commitdiff