From a0506353ba651069a42293f32601737fe129a66f Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Fri, 17 Apr 2026 18:25:04 +0200 Subject: [PATCH] improve palettisation * add an alternative colour counter mode (faster but eats a lot of RAM) * add support for palette generation options (to that mode can be set) * change full palettisation mode into using 16MB LUT --- src/main.rs | 26 +++++-- src/palettise.rs | 177 +++++++++++++++++++++++++++++++++++----------- src/transcoder.rs | 2 +- 3 files changed, 158 insertions(+), 47 deletions(-) diff --git a/src/main.rs b/src/main.rs index 9bdde8c..2409467 100644 --- a/src/main.rs +++ b/src/main.rs @@ -110,6 +110,17 @@ macro_rules! parse_id { } } +fn split_options(opts: &mut Vec, in_opts: &str) { + for opt in in_opts.split(',') { + let oval: Vec<_> = opt.split('=').collect(); + match oval.len() { + 1 => { opts.push(OptionArgs{ name: oval[0].to_string(), value: None }); }, + 2 => { opts.push(OptionArgs{ name: oval[0].to_string(), value: Some(oval[1].to_string()) }); }, + _ => { println!(" unrecognised option '{opt}'"); }, + } + } +} + fn retrieve_packets(transcoder: &mut Transcoder, mux: &mut Muxer, vdata_size: &mut usize, adata_size: &mut usize, end: bool) -> bool { while let Some(pkt) = transcoder.queue.get_packet(&mut transcoder.debug) { transcoder.debug_log(DebugLog::MUX, &format!(" Got output packet for stream {} ts {:?}/{:?}", pkt.get_stream().get_id(), pkt.ts.pts, pkt.ts.dts)); @@ -205,6 +216,7 @@ fn main() { let mut custom_profile = false; let mut ignerr = false; let mut skip_unknown = false; + let mut palgen_options = Vec::new(); while arg_idx < args.len() { match args[arg_idx].as_str() { "--list-decoders" => { @@ -303,6 +315,10 @@ fn main() { "--generate-palette" => { transcoder.gen_pal = true; }, + "--palgen-options" => { + next_arg!(args, arg_idx); + split_options(&mut palgen_options, &args[arg_idx]); + }, "--output" | "-o" => { next_arg!(args, arg_idx); transcoder.output_name = args[arg_idx].clone(); @@ -601,7 +617,7 @@ fn main() { if nenc > 0 { let mut ccounters = Vec::with_capacity(nenc); for _ in 0..nenc { - ccounters.push(crate::palettise::ColourCounter::new()); + ccounters.push(crate::palettise::ColourCounter::new(&palgen_options)); } let mut cur_dmx = 0; let mut last_known_time = None; @@ -707,13 +723,15 @@ fn main() { } } } + if transcoder.verbose > 0 { + println!(); + println!(" generating palette(s)..."); + } for (&id, plt) in rev_map.iter().zip(ccounters.iter()) { let pal = plt.get_pal(); transcoder.qsupport.glbl_pal.push((id, pal)); } - if transcoder.verbose > 0 { - println!(); - } + ccounters.clear(); transcoder.debug_log(DebugLog::GENERAL, "Resetting state after palettisation"); // this is necessary since not all demuxers allow to seek even back to the start demuxers.clear(); diff --git a/src/palettise.rs b/src/palettise.rs index 5bffb70..7573954 100644 --- a/src/palettise.rs +++ b/src/palettise.rs @@ -28,6 +28,16 @@ fn find_nearest(pix: &[u8], pal: &[[u8; 3]]) -> usize { bestidx } +fn gen_full_lut(dst: &mut [u8], pal: &[[u8; 3]; 256]) { + for (r, chunkr) in dst.chunks_exact_mut(1 << 16).enumerate() { + for (g, chunkg) in chunkr.chunks_exact_mut(1 << 8).enumerate() { + for (b, el) in chunkg.iter_mut().enumerate() { + *el = find_nearest(&[r as u8, g as u8, b as u8], pal) as u8; + } + } + } +} + struct LocalSearch { pal: [[u8; 3]; 256], db: Vec>, @@ -196,7 +206,7 @@ impl KDTree { #[allow(clippy::large_enum_variant)] enum PMode { - Full, + Full(Vec), Local(LocalSearch), Tree(KDTree), } @@ -233,9 +243,14 @@ impl LookupCache { #[allow(dead_code)] impl Palettiser { + pub fn get_default_mode() -> PaletteSearchMode { PaletteSearchMode::Full } pub fn new(mode: PaletteSearchMode, pal: &[[u8; 3]; 256]) -> Self { let pmode = match mode { - PaletteSearchMode::Full => PMode::Full, + PaletteSearchMode::Full => { + let mut tab = vec![0; 1 << 24]; + gen_full_lut(&mut tab, pal); + PMode::Full(tab) + }, PaletteSearchMode::Local => PMode::Local(LocalSearch::new(pal)), PaletteSearchMode::KDTree => PMode::Tree(KDTree::new(pal)), }; @@ -243,7 +258,10 @@ impl Palettiser { } pub fn search(&self, pix: [u8; 3]) -> usize { match &self.pmode { - PMode::Full => find_nearest(&pix, &self.pal), + PMode::Full(tab) => { + let idx = (usize::from(pix[0]) << 16) | (usize::from(pix[1]) << 8) | usize::from(pix[2]); + usize::from(tab[idx]) + }, PMode::Local(ls) => ls.search(pix), PMode::Tree(kdt) => kdt.search(pix), } @@ -251,7 +269,7 @@ impl Palettiser { pub fn set_pal(&mut self, pal: &[[u8; 3]; 256]) { self.pal.copy_from_slice(pal); match &mut self.pmode { - PMode::Full => {}, + PMode::Full(ref mut tab) => { gen_full_lut(tab, pal); }, PMode::Local(ref mut ls) => { *ls = LocalSearch::new(pal); }, PMode::Tree(ref mut kdt) => { *kdt = KDTree::new(pal); }, } @@ -280,18 +298,14 @@ impl Palettiser { 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]; let mut cache = LookupCache::default(); match &self.pmode { - PMode::Full => { + PMode::Full(tab) => { for (src, dline) in sdata.chunks(istride) .zip(dst[doff..].chunks_exact_mut(dstride)).take(h) { for (pix, chunk) in dline.iter_mut() .zip(src.chunks_exact(esize)).take(w) { let spixel = [chunk[coffs[0]], chunk[coffs[1]], chunk[coffs[2]]]; - if let Some(idx) = cache.lookup(spixel) { - *pix = idx; - } else { - *pix = find_nearest(&spixel, &self.pal) as u8; - cache.add(spixel, *pix); - } + let idx = (usize::from(spixel[0]) << 16) | (usize::from(spixel[1]) << 8) | usize::from(spixel[2]); + *pix = tab[idx]; } } }, @@ -609,25 +623,19 @@ impl Bucket { } } -pub struct ColourCounter { +struct BucketCounter { buckets: Vec, nentries: usize, - scaler: Option, - sc_buf: NABufferType, - oinfo: NAVideoInfo, } -impl ColourCounter { - pub fn new() -> Self { +impl BucketCounter { + fn new() -> Self { Self { buckets: vec![Bucket::default(); NBUCKETS], nentries: 0, - scaler: None, - sc_buf: NABufferType::None, - oinfo: NAVideoInfo{ width: 0, height: 0, flipped: false, format: RGB24_FORMAT, bits: 24 }, } } - pub fn add(&mut self, clr: [u8; 3]) { + fn add(&mut self, clr: [u8; 3]) { let idx = Bucket::key(clr); if self.buckets[idx].add(clr) { self.nentries += 1; @@ -642,6 +650,91 @@ impl ColourCounter { } } } + fn get_pal(&self, debug: bool) -> [[u8; 3]; 256] { + let mut pal = [[0; 3]; 256]; + if debug { + println!(" {} entries in total", self.nentries); + } + if self.nentries <= 256 { + let mut dst = pal.iter_mut(); + for bucket in self.buckets.iter() { + for clr in bucket.clrs.iter() { + *dst.next().unwrap() = clr.clr; + } + } + } else { + let mut in_clrs = Vec::with_capacity(self.nentries); + for bucket in self.buckets.iter() { + for clr in bucket.clrs.iter() { + in_clrs.push(*clr); + } + } + let mut ppal = [Colour::default(); 256]; + let _prim_clrs = quantise_median_cut::(&in_clrs, &mut ppal); + let mut elbg: ELBG = ELBG::new(&ppal); + elbg.quantise(&in_clrs, &mut ppal); + for (dst, src) in pal.iter_mut().zip(ppal.iter()) { + *dst = src.clr; + } + } + pal + } +} + +enum CounterType { + Buckets(BucketCounter), + Brawn(Vec), +} + +pub struct ColourCounter { + scaler: Option, + sc_buf: NABufferType, + oinfo: NAVideoInfo, + ctype: CounterType, + debug: bool, +} + +impl ColourCounter { + pub fn new(options: &[OptionArgs]) -> Self { + let mut counter_type = "full"; + let mut debug = false; + for opt in options.iter() { + match opt.name.as_str() { + "counter" => { + match opt.value.as_deref() { + Some("full") => { counter_type = "full"; }, + Some("bucket") => { counter_type = "bucket"; }, + Some(_) => { println!("unknown counter method"); }, + None => { println!("counter option requires a value"); }, + } + }, + "debug" => { + debug = true; + }, + _ => {}, + } + } + let ctype = match counter_type { + "full" => CounterType::Brawn(vec![0; 1 << 24]), + "bucket" => CounterType::Buckets(BucketCounter::new()), + _ => unreachable!(), + }; + Self { + scaler: None, + sc_buf: NABufferType::None, + oinfo: NAVideoInfo{ width: 0, height: 0, flipped: false, format: RGB24_FORMAT, bits: 24 }, + ctype, debug, + } + } + pub fn add(&mut self, clr: [u8; 3]) { + match self.ctype { + CounterType::Buckets(ref mut bkt) => { bkt.add(clr); }, + CounterType::Brawn(ref mut hist) => { + let idx = usize::from(clr[0]) * (1 << 16) + usize::from(clr[1]) * (1 << 8) + usize::from(clr[2]); + hist[idx] += 1; + }, + } + } fn add_rgb24(&mut self, buf: &NAVideoBuffer, width: usize, height: usize) { let src = buf.get_data(); let stride = buf.get_stride(0); @@ -687,29 +780,29 @@ impl ColourCounter { } } pub fn get_pal(&self) -> [[u8; 3]; 256] { - let mut pal = [[0; 3]; 256]; - if self.nentries <= 256 { - let mut dst = pal.iter_mut(); - for bucket in self.buckets.iter() { - for clr in bucket.clrs.iter() { - *dst.next().unwrap() = clr.clr; + match self.ctype { + CounterType::Buckets(ref bkt) => bkt.get_pal(self.debug), + CounterType::Brawn(ref hist) => { + let mut pal = [[0; 3]; 256]; + let clrs: Vec = hist.iter().enumerate().filter(|(_i, &count)| count > 0).map(|(i, &count)| Colour { clr: [(i >> 16) as u8, (i >> 8) as u8, i as u8], count}).collect(); + if self.debug { + println!(" {} entries in total", clrs.len()); } - } - } else { - let mut in_clrs = Vec::with_capacity(self.nentries); - for bucket in self.buckets.iter() { - for clr in bucket.clrs.iter() { - in_clrs.push(*clr); + if clrs.len() <= 256 { + for (dclr, sclr) in pal.iter_mut().zip(clrs.iter()) { + *dclr = sclr.clr; + } + } else { + let mut ppal = [Colour::default(); 256]; + let _prim_clrs = quantise_median_cut::(&clrs, &mut ppal); + let mut elbg: ELBG = ELBG::new(&ppal); + elbg.quantise(&clrs, &mut ppal); + for (dst, src) in pal.iter_mut().zip(ppal.iter()) { + *dst = src.clr; + } } - } - let mut ppal = [Colour::default(); 256]; - let _prim_clrs = quantise_median_cut::(&in_clrs, &mut ppal); - let mut elbg: ELBG = ELBG::new(&ppal); - elbg.quantise(&in_clrs, &mut ppal); - for (dst, src) in pal.iter_mut().zip(ppal.iter()) { - *dst = src.clr; - } + pal + }, } - pal } } diff --git a/src/transcoder.rs b/src/transcoder.rs index 12601db..ab2afa7 100644 --- a/src/transcoder.rs +++ b/src/transcoder.rs @@ -977,7 +977,7 @@ impl Transcoder { if let Some(ref mut p) = plt { p.set_pal(gpal); } else { - plt = Some(Palettiser::new(PaletteSearchMode::default(), gpal)); + plt = Some(Palettiser::new(Palettiser::get_default_mode(), gpal)); } } } -- 2.39.5