From 06ae1076d8ecc0d81a1ea738069360bf2168da97 Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Tue, 14 Apr 2026 20:26:24 +0200 Subject: [PATCH] bolt on global palette conversion functionality In the future it should probably implemented as more flexible filter pipeline and support calculating palette for the output, but for now even such hack should do. --- src/main.rs | 1 + src/palettise.rs | 402 ++++++++++++++++++++++++++++++++++++++++++++++ src/transcoder.rs | 48 +++++- 3 files changed, 445 insertions(+), 6 deletions(-) create mode 100644 src/palettise.rs diff --git a/src/main.rs b/src/main.rs index ac55eab..756f33a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -52,6 +52,7 @@ macro_rules! parse_and_apply_options { mod null; use null::*; mod acvt; +mod palettise; mod transcoder; use crate::transcoder::*; diff --git a/src/palettise.rs b/src/palettise.rs new file mode 100644 index 0000000..c1b75c3 --- /dev/null +++ b/src/palettise.rs @@ -0,0 +1,402 @@ +use nihav_core::frame::*; +use nihav_codec_support::codecs::qt_pal::*; + +use crate::transcoder::OptionArgs; + +fn find_nearest(pix: &[u8], pal: &[[u8; 3]]) -> usize { + let mut bestidx = 0; + let mut bestdist = 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 +} + +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) + } + 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 + } + fn search(&self, pix: [u8; 3]) -> usize { + let idx = Self::quant(pix); + let mut best_dist = 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, +} + +struct KDTree { + nodes: Vec, +} + +fn avg_u8(a: u8, b: u8) -> u8 { + (a & b) + ((a ^ b) >> 1) +} + +impl KDTree { + 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 clr in pal[start..end].iter() { + for ((mi, ma), &c) in min.iter_mut().zip(max.iter_mut()).zip(clr.iter()) { + *mi = (*mi).min(c); + *ma = (*ma).max(c); + } + } + 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 + } + 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; + } + } +} + +#[derive(Clone,Copy,Debug,PartialEq,Default)] +pub enum PaletteSearchMode { + Full, + #[default] + Local, + Tree, +} + +#[allow(clippy::large_enum_variant)] +enum PMode { + Full, + Local(LocalSearch), + Tree(KDTree), +} + +pub struct Palettiser { + pal: [[u8; 3]; 256], + pmode: PMode, +} + +#[allow(dead_code)] +impl Palettiser { + pub fn new(mode: PaletteSearchMode, pal: &[[u8; 3]; 256]) -> Self { + let pmode = match mode { + PaletteSearchMode::Full => PMode::Full, + PaletteSearchMode::Local => PMode::Local(LocalSearch::new(pal)), + PaletteSearchMode::Tree => PMode::Tree(KDTree::new(pal)), + }; + Self { pal: *pal, pmode } + } + pub fn search(&self, pix: [u8; 3]) -> usize { + match &self.pmode { + PMode::Full => find_nearest(&pix, &self.pal), + PMode::Local(ls) => ls.search(pix), + PMode::Tree(kdt) => kdt.search(pix), + } + } + pub fn set_pal(&mut self, pal: &[[u8; 3]; 256]) { + self.pal.copy_from_slice(pal); + match &mut self.pmode { + PMode::Full => {}, + PMode::Local(ref mut ls) => { *ls = LocalSearch::new(pal); }, + PMode::Tree(ref mut kdt) => { *kdt = KDTree::new(pal); }, + } + } + pub fn palettise_frame(&self, pic_in: &NABufferType, pic_out: &mut NABufferType) -> Result<(), &'static str> { +// todo remap already paletted format + 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(); + + 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]; + match &self.pmode { + PMode::Full => { + 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]]]; + *pix = find_nearest(&spixel, &self.pal) as u8; + } + } + }, + PMode::Local(ls) => { + 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]]]; + *pix = ls.search(spixel) as u8; + } + } + }, + PMode::Tree(kdt) => { + 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]]]; + *pix = kdt.search(spixel) as u8; + } + } + }, + } + } 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 dline in dst[doff..].chunks_exact_mut(dstride).take(h) { + for (x, pix) in dline[..w].iter_mut().enumerate() { + let spixel = [sdata[roff + x], sdata[goff + x], sdata[boff + x]]; + *pix = self.search(spixel) as u8; + } + roff += rstride; + goff += gstride; + boff += bstride; + } + } + + 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 dst[paloff..].chunks_mut(esize).zip(self.pal.iter()) { + dpal[coffs[0]] = spal[0]; + dpal[coffs[1]] = spal[1]; + dpal[coffs[2]] = spal[2]; + } + Ok(()) + } else { + Err("invalid buffer format") + } + } +} + +pub fn create_palettiser(enc_opts: &[OptionArgs]) -> Option { + let mut pmode = None; + let mut pal = std::array::from_fn(|i| [i as u8; 3]); + let mut pal_is_some = false; + for opt in enc_opts.iter() { + match opt.name.as_str() { + "pal.mode" => { + if let Some(pmode_val) = &opt.value { + pmode = match pmode_val.as_str() { + "full" => Some(PaletteSearchMode::Full), + "local" => Some(PaletteSearchMode::Local), + "tree" => Some(PaletteSearchMode::Tree), + _ => { + println!("invalid palettisation mode"); + None + } + }; + } else { + println!("option 'pal.mode' requires an argument"); + } + }, + "pal.set" => { + if let Some(pname) = &opt.value { + match pname.as_str() { + "grey" | "gray" => { + for (i, clr) in pal.iter_mut().enumerate() { + *clr = [i as u8; 3]; + } + }, + "bw" => { + pal[0] = [0x00; 3]; + pal[1] = [0xFF; 3]; + }, + "wb" => { + pal[0] = [0xFF; 3]; + pal[1] = [0x00; 3]; + }, + "systematic" => { + for (i, clr) in pal.iter_mut().enumerate() { + let idx = i as u8; + let r = idx >> 5; + let g = (idx >> 2) & 7; + let b = idx & 3; + *clr = [(r << 5) | (r << 2) | (r >> 1), + (g << 5) | (g << 2) | (g >> 1), + b * 0x55]; + } + }, + "qt4" => { + for (dclr, sclr) in pal.iter_mut().zip(MOV_DEFAULT_PAL_2BIT.chunks_exact(4)) { + dclr.copy_from_slice(&sclr[..3]); + } + }, + "qt16" => { + for (dclr, sclr) in pal.iter_mut().zip(MOV_DEFAULT_PAL_4BIT.chunks_exact(4)) { + dclr.copy_from_slice(&sclr[..3]); + } + }, + "qt256" => { + for (dclr, sclr) in pal.iter_mut().zip(MOV_DEFAULT_PAL_8BIT.chunks_exact(4)) { + dclr.copy_from_slice(&sclr[..3]); + } + }, + _ => { + println!("invalid or unknown palette mode"); + return None; + } + }; + pal_is_some = true; + } else { + println!("option 'pal.set' requires an argument"); + return None; + } + }, + _ => {}, + } + } + if pmode.is_some() || pal_is_some { + Some(Palettiser::new(pmode.unwrap_or_default(), &pal)) + } else { + None + } +} diff --git a/src/transcoder.rs b/src/transcoder.rs index dadf4bb..b37dfd8 100644 --- a/src/transcoder.rs +++ b/src/transcoder.rs @@ -10,6 +10,7 @@ use nihav_core::reorder::*; use nihav_core::scale::*; use crate::acvt::*; +use crate::palettise::*; use nihav_hlblocks::demux::*; use nihav_hlblocks::imgseqdec::*; @@ -233,18 +234,24 @@ pub struct VideoEncodeContext { pub tb_den: u32, pub cfr: bool, pub last_ts: Option, + pub plt: Option, + pub plt_buf: NABufferType, } impl EncoderInterface for VideoEncodeContext { fn encode_frame(&mut self, dst_id: u32, frm: NAFrameRef, scale_opts: &[(String, String)], queue: &mut OutputQueue, dbg: &mut Option) -> EncoderResult { let buf = frm.get_buffer(); if let Some(vinfo) = buf.get_video_info() { + let mut tgt_vinfo = self.vinfo; + if self.plt.is_some() { + tgt_vinfo.format = RGB24_FORMAT; + } if self.scaler.is_none() && vinfo != self.vinfo { let ifmt = get_scale_fmt_from_pic(&buf); let ofmt = ScaleInfo { - width: self.vinfo.width, - height: self.vinfo.height, - fmt: self.vinfo.format, + width: tgt_vinfo.width, + height: tgt_vinfo.height, + fmt: tgt_vinfo.format, }; if let Some(ref mut dlog) = dbg { dlog.log(DebugLog::ENCODE, &format!("Input for the output stream {dst_id} differs in format, inserting scaler")); @@ -263,7 +270,7 @@ impl EncoderInterface for VideoEncodeContext { } } } - let cbuf = if let NABufferType::None = buf { + let mut cbuf = if let NABufferType::None = buf { if (self.encoder.get_capabilities() & ENC_CAPS_SKIPFRAME) == 0 { if let NABufferType::None = self.scaler_buf { if let Some(ref mut dlog) = dbg { @@ -310,6 +317,23 @@ impl EncoderInterface for VideoEncodeContext { if self.scaler.is_none() && !matches!(cbuf, NABufferType::None) { self.scaler_buf = cbuf.clone(); } + if let Some(ref plt) = self.plt { + if matches!(self.plt_buf, NABufferType::None) { + let mut tgt_vinfo = self.vinfo; + tgt_vinfo.format = PAL8_FORMAT; + self.plt_buf = if let Ok(buf) = alloc_video_buffer(tgt_vinfo, 0) { + buf + } else { + println!("cannot create palettiser buffer"); + return Err(EncoderError::AllocError); + }; + } + if let Err(err) = plt.palettise_frame(&cbuf, &mut self.plt_buf) { + println!("palettisation error {err}"); + return Err(EncoderError::Bug); + } + cbuf = self.plt_buf.clone(); + } let ref_ts = frm.get_time_information(); let new_pts = if let Some(ts) = ref_ts.pts { @@ -943,6 +967,9 @@ impl Transcoder { } let ret_eparams = ret_eparams.unwrap(); + let plt = create_palettiser(&oopts.enc_opts); + oopts.enc_opts.retain(|opt| !opt.name.starts_with("pal.")); + let name = format!("output stream {}", out_id); parse_and_apply_options!(encoder, &oopts.enc_opts, name); @@ -978,16 +1005,23 @@ impl Transcoder { last_ts: None, tb_num: enc_stream.tb_num, tb_den: enc_stream.tb_den, + plt, + plt_buf: NABufferType::None, }) } else { - let ofmt = ScaleInfo { fmt: dvinfo.format, width: dvinfo.width, height: dvinfo.height }; + let tgt_fmt = if plt.is_some() { RGB24_FORMAT } else { dvinfo.format }; + let ofmt = ScaleInfo { fmt: tgt_fmt, width: dvinfo.width, height: dvinfo.height }; let ret = NAScale::new_with_options(ofmt, ofmt, scale_opts); if ret.is_err() { println!("cannot create scaler"); return None; } let scaler = ret.unwrap(); - let ret = alloc_video_buffer(*dvinfo, 4); + let mut sc_info = *dvinfo; + if plt.is_some() { + sc_info.format = RGB24_FORMAT; + } + let ret = alloc_video_buffer(sc_info, 4); if ret.is_err() { println!("cannot create scaler buffer"); return None; @@ -1002,6 +1036,8 @@ impl Transcoder { last_ts: None, tb_num: enc_stream.tb_num, tb_den: enc_stream.tb_den, + plt, + plt_buf: NABufferType::None, }) } }, -- 2.39.5