summaryrefslogtreecommitdiffstats
path: root/src/frontier/mean.rs
diff options
context:
space:
mode:
authorTavian Barnes <tavianator@tavianator.com>2020-05-03 10:55:16 -0400
committerGitHub <noreply@github.com>2020-05-03 10:55:16 -0400
commitce2904b4840611f769b92b55bf6d9b5afe84d3d7 (patch)
treea133319a302f95edf7a7a261262a8f24473bd21c /src/frontier/mean.rs
parentd95e93bf70f3351e6fd489284794ef7909fd94ce (diff)
parent2984e8f93fe88d0ee7eb3c0561dcd2da44807429 (diff)
downloadkd-forest-ce2904b4840611f769b92b55bf6d9b5afe84d3d7.tar.xz
Merge pull request #1 from tavianator/rust
Rewrite in rust
Diffstat (limited to 'src/frontier/mean.rs')
-rw-r--r--src/frontier/mean.rs140
1 files changed, 140 insertions, 0 deletions
diff --git a/src/frontier/mean.rs b/src/frontier/mean.rs
new file mode 100644
index 0000000..889c5ba
--- /dev/null
+++ b/src/frontier/mean.rs
@@ -0,0 +1,140 @@
+//! Mean selection frontier.
+
+use super::{neighbors, Frontier, Pixel};
+
+use crate::color::{ColorSpace, Rgb8};
+use crate::metric::soft::SoftKdForest;
+use crate::metric::NearestNeighbors;
+
+use std::iter;
+use std::rc::Rc;
+
+/// A pixel on a mean frontier.
+#[derive(Debug)]
+enum MeanPixel<C> {
+ Empty,
+ Fillable(Rc<Pixel<C>>),
+ Filled(C),
+}
+
+impl<C: ColorSpace> MeanPixel<C> {
+ fn filled_color(&self) -> Option<C> {
+ match self {
+ Self::Filled(color) => Some(*color),
+ _ => None,
+ }
+ }
+}
+
+/// A [Frontier] that looks at the average color of each pixel's neighbors.
+pub struct MeanFrontier<C> {
+ pixels: Vec<MeanPixel<C>>,
+ forest: SoftKdForest<Rc<Pixel<C>>>,
+ width: u32,
+ height: u32,
+ len: usize,
+ deleted: usize,
+}
+
+impl<C: ColorSpace> MeanFrontier<C> {
+ /// Create a MeanFrontier with the given dimensions and initial pixel location.
+ pub fn new(width: u32, height: u32, x0: u32, y0: u32) -> Self {
+ let size = (width as usize) * (height as usize);
+ let mut pixels = Vec::with_capacity(size);
+ for _ in 0..size {
+ pixels.push(MeanPixel::Empty);
+ }
+
+ let pixel0 = Rc::new(Pixel::new(x0, y0, C::from(Rgb8::from([0, 0, 0]))));
+ let i = (x0 + y0 * width) as usize;
+ pixels[i] = MeanPixel::Fillable(pixel0.clone());
+
+ Self {
+ pixels,
+ forest: iter::once(pixel0).collect(),
+ width,
+ height,
+ len: 1,
+ deleted: 0,
+ }
+ }
+
+ fn pixel_index(&self, x: u32, y: u32) -> usize {
+ debug_assert!(x < self.width);
+ debug_assert!(y < self.height);
+
+ (x + y * self.width) as usize
+ }
+
+ fn fill(&mut self, x: u32, y: u32, color: C) {
+ let i = self.pixel_index(x, y);
+ match &self.pixels[i] {
+ MeanPixel::Empty => {}
+ MeanPixel::Fillable(pixel) => {
+ pixel.delete();
+ self.deleted += 1;
+ }
+ _ => unreachable!(),
+ }
+ self.pixels[i] = MeanPixel::Filled(color);
+
+ let mut pixels = Vec::new();
+ for &(x, y) in &neighbors(x, y) {
+ if x < self.width && y < self.height {
+ let i = self.pixel_index(x, y);
+ match &self.pixels[i] {
+ MeanPixel::Empty => {}
+ MeanPixel::Fillable(pixel) => {
+ pixel.delete();
+ self.deleted += 1;
+ }
+ MeanPixel::Filled(_) => continue,
+ }
+ let color = C::average(
+ neighbors(x, y)
+ .iter()
+ .filter(|(x, y)| *x < self.width && *y < self.height)
+ .map(|(x, y)| self.pixel_index(*x, *y))
+ .map(|i| &self.pixels[i])
+ .map(MeanPixel::filled_color)
+ .flatten(),
+ );
+ let pixel = Rc::new(Pixel::new(x, y, color));
+ self.pixels[i] = MeanPixel::Fillable(pixel.clone());
+ pixels.push(pixel);
+ }
+ }
+
+ self.len += pixels.len();
+ self.forest.extend(pixels);
+
+ if 2 * self.deleted >= self.len {
+ self.forest.rebuild();
+ self.len -= self.deleted;
+ self.deleted = 0;
+ }
+ }
+}
+
+impl<C: ColorSpace> Frontier for MeanFrontier<C> {
+ fn width(&self) -> u32 {
+ self.width
+ }
+
+ fn height(&self) -> u32 {
+ self.height
+ }
+
+ fn len(&self) -> usize {
+ self.len - self.deleted
+ }
+
+ fn place(&mut self, rgb8: Rgb8) -> Option<(u32, u32)> {
+ let color = C::from(rgb8);
+ let (x, y) = self.forest.nearest(&color).map(|n| n.item.pos)?;
+
+ self.fill(x, y, color);
+
+ Some((x, y))
+ }
+}