nihav-codec-support/src/vq/generic_mediancut.rs

   1 use super::{VQElement, VQElementSum};
   2
   3 struct VQBox<'a, T: VQElement> {
   4     points: &'a mut [T],
   5     max:    T,
   6     min:    T,
   7 }
   8
   9 impl<'a, T: VQElement> VQBox<'a, T> {
  10     fn new(points: &'a mut [T]) -> Self {
  11         let mut max = T::min_cw();
  12         let mut min = T::max_cw();
  13         for point in points.iter() {
  14             max = max.max(*point);
  15             min = min.min(*point);
  16         }
  17         Self { points, max, min }
  18     }
  19     fn can_split(&self) -> bool {
  20         self.max != self.min
  21     }
  22     fn calc_min_and_max(points: &[T]) -> (T, T) {
  23         let mut max = T::min_cw();
  24         let mut min = T::max_cw();
  25         for point in points.iter() {
  26             max = max.max(*point);
  27             min = min.min(*point);
  28         }
  29         (min, max)
  30     }
  31     fn get_pivot(arr: &[T]) -> usize {
  32         if arr.len() < 2 {
  33             return 0;
  34         }
  35         let mut lastval = arr[0];
  36         let mut pivot = 0;
  37         let mut idx = 1;
  38         for el in arr.iter().skip(1) {
  39             if *el != lastval && (pivot == 0 || idx <= arr.len() / 2) {
  40                 pivot = idx;
  41                 lastval = *el;
  42             }
  43             idx += 1;
  44         }
  45         pivot
  46     }
  47     fn split(self) -> (VQBox<'a, T>, VQBox<'a, T>) {
  48         let sort_c = T::max_dist_component(&self.min, &self.max);
  49         T::sort_by_component(self.points, sort_c);
  50         let pivot = Self::get_pivot(self.points);
  51         let (part0, part1) = self.points.split_at_mut(pivot);
  52         let (min0, max0) = Self::calc_min_and_max(part0);
  53         let (min1, max1) = Self::calc_min_and_max(part1);
  54         let box0 = VQBox { points: part0, max: max0, min: min0 };
  55         let box1 = VQBox { points: part1, max: max1, min: min1 };
  56
  57         (box0, box1)
  58     }
  59 }
  60
  61 pub fn quantise_median_cut<T: VQElement, TS: VQElementSum<T>>(src: &[T], dst: &mut [T]) -> usize {
  62     let mut points = Vec::with_capacity(src.len());
  63     points.extend(src.into_iter());
  64     for comp in 0..T::num_components() {
  65         T::sort_by_component(points.as_mut_slice(), comp);
  66     }
  67     let box0 = VQBox::new(points.as_mut_slice());
  68     let mut boxes: Vec<VQBox<T>> = Vec::with_capacity(dst.len());
  69     boxes.push(box0);
  70     let mut changed = true;
  71     while changed && boxes.len() < dst.len() {
  72         let end = boxes.len();
  73         changed = false;
  74         let mut split_largest = false;
  75         for _ in 0..end {
  76             let curbox = boxes.remove(0);
  77             if curbox.can_split() {
  78                 let (box0, box1) = curbox.split();
  79                 boxes.push(box0);
  80                 boxes.push(box1);
  81                 changed = true;
  82             } else {
  83                 boxes.push(curbox);
  84                 split_largest = true;
  85                 break;
  86             }
  87             if boxes.len() == dst.len() {
  88                 break;
  89             }
  90         }
  91         if split_largest {
  92             let mut maxidx = 0;
  93             let mut lcount = 0;
  94             for (i, cbox) in boxes.iter().enumerate() {
  95                 if cbox.can_split() && cbox.points.len() > lcount {
  96                     lcount = cbox.points.len();
  97                     maxidx = i;
  98                 }
  99             }
 100             if lcount > 0 {
 101                 let curbox = boxes.remove(maxidx);
 102                 let (box0, box1) = curbox.split();
 103                 boxes.push(box0);
 104                 boxes.push(box1);
 105                 changed = true;
 106             }
 107         }
 108     }
 109     for (dst, curbox) in dst.iter_mut().zip(boxes.iter()) {
 110         let mut sum = TS::zero();
 111         sum.add(curbox.min, 1);
 112         sum.add(curbox.max, 1);
 113         *dst = sum.get_centroid();
 114     }
 115
 116     boxes.len()
 117 }