--- /dev/null
+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<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)
+ }
+ 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<KDNode>,
+}
+
+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<Palettiser> {
+ 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
+ }
+}
use nihav_core::scale::*;
use crate::acvt::*;
+use crate::palettise::*;
use nihav_hlblocks::demux::*;
use nihav_hlblocks::imgseqdec::*;
pub tb_den: u32,
pub cfr: bool,
pub last_ts: Option<u64>,
+ pub plt: Option<Palettiser>,
+ 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<DebugLog>) -> EncoderResult<bool> {
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"));
}
}
}
- 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 {
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 {
}
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);
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;
last_ts: None,
tb_num: enc_stream.tb_num,
tb_den: enc_stream.tb_den,
+ plt,
+ plt_buf: NABufferType::None,
})
}
},
projects / nihav-encoder.git / commitdiff