main: Implement the main binaryrust

3 files changed, 591 insertions, 1 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 7f81f15..c066dd8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -7,18 +7,67 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5d2e7343e7fc9de883d1b0341e0b13970f764c14101234857d2ddafa1cb1cac2"
 
 [[package]]
+name = "ansi_term"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ee49baf6cb617b853aa8d93bf420db2383fab46d314482ca2803b40d5fde979b"
+dependencies = [
+ "winapi",
+]
+
+[[package]]
+name = "arrayref"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a4c527152e37cf757a3f78aae5a06fbeefdb07ccc535c980a3208ee3060dd544"
+
+[[package]]
+name = "arrayvec"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cff77d8686867eceff3105329d4698d96c2391c176d5d03adc90c7389162b5b8"
+
+[[package]]
+name = "atty"
+version = "0.2.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8"
+dependencies = [
+ "hermit-abi",
+ "libc",
+ "winapi",
+]
+
+[[package]]
 name = "autocfg"
 version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f8aac770f1885fd7e387acedd76065302551364496e46b3dd00860b2f8359b9d"
 
 [[package]]
+name = "base64"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b41b7ea54a0c9d92199de89e20e58d49f02f8e699814ef3fdf266f6f748d15c7"
+
+[[package]]
 name = "bitflags"
 version = "1.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cf1de2fe8c75bc145a2f577add951f8134889b4795d47466a54a5c846d691693"
 
 [[package]]
+name = "blake2b_simd"
+version = "0.5.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d8fb2d74254a3a0b5cac33ac9f8ed0e44aa50378d9dbb2e5d83bd21ed1dc2c8a"
+dependencies = [
+ "arrayref",
+ "arrayvec",
+ "constant_time_eq",
+]
+
+[[package]]
 name = "bytemuck"
 version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -37,12 +86,33 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4785bdd1c96b2a846b2bd7cc02e86b6b3dbf14e7e53446c4f54c92a361040822"
 
 [[package]]
+name = "clap"
+version = "2.33.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5067f5bb2d80ef5d68b4c87db81601f0b75bca627bc2ef76b141d7b846a3c6d9"
+dependencies = [
+ "ansi_term",
+ "atty",
+ "bitflags",
+ "strsim",
+ "textwrap",
+ "unicode-width",
+ "vec_map",
+]
+
+[[package]]
 name = "color_quant"
 version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0dbbb57365263e881e805dc77d94697c9118fd94d8da011240555aa7b23445bd"
 
 [[package]]
+name = "constant_time_eq"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc"
+
+[[package]]
 name = "crc32fast"
 version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -109,6 +179,28 @@ dependencies = [
 ]
 
 [[package]]
+name = "dirs"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13aea89a5c93364a98e9b37b2fa237effbb694d5cfe01c5b70941f7eb087d5e3"
+dependencies = [
+ "cfg-if",
+ "dirs-sys",
+]
+
+[[package]]
+name = "dirs-sys"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "afa0b23de8fd801745c471deffa6e12d248f962c9fd4b4c33787b055599bde7b"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "redox_users",
+ "winapi",
+]
+
+[[package]]
 name = "either"
 version = "1.5.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -185,9 +277,12 @@ dependencies = [
 name = "kd-forest"
 version = "2.0.0"
 dependencies = [
+ "clap",
  "image",
  "ordered-float",
  "rand",
+ "rand_pcg",
+ "term",
 ]
 
 [[package]]
@@ -352,6 +447,15 @@ dependencies = [
 ]
 
 [[package]]
+name = "rand_pcg"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "16abd0c1b639e9eb4d7c50c0b8100b0d0f849be2349829c740fe8e6eb4816429"
+dependencies = [
+ "rand_core",
+]
+
+[[package]]
 name = "rayon"
 version = "1.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -376,6 +480,35 @@ dependencies = [
 ]
 
 [[package]]
+name = "redox_syscall"
+version = "0.1.56"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2439c63f3f6139d1b57529d16bc3b8bb855230c8efcc5d3a896c8bea7c3b1e84"
+
+[[package]]
+name = "redox_users"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09b23093265f8d200fa7b4c2c76297f47e681c655f6f1285a8780d6a022f7431"
+dependencies = [
+ "getrandom",
+ "redox_syscall",
+ "rust-argon2",
+]
+
+[[package]]
+name = "rust-argon2"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2bc8af4bda8e1ff4932523b94d3dd20ee30a87232323eda55903ffd71d2fb017"
+dependencies = [
+ "base64",
+ "blake2b_simd",
+ "constant_time_eq",
+ "crossbeam-utils",
+]
+
+[[package]]
 name = "scoped_threadpool"
 version = "0.1.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -388,6 +521,31 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
 
 [[package]]
+name = "strsim"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a"
+
+[[package]]
+name = "term"
+version = "0.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c0863a3345e70f61d613eab32ee046ccd1bcc5f9105fe402c61fcd0c13eeb8b5"
+dependencies = [
+ "dirs",
+ "winapi",
+]
+
+[[package]]
+name = "textwrap"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d326610f408c7a4eb6f51c37c330e496b08506c9457c9d34287ecc38809fb060"
+dependencies = [
+ "unicode-width",
+]
+
+[[package]]
 name = "tiff"
 version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -399,7 +557,41 @@ dependencies = [
 ]
 
 [[package]]
+name = "unicode-width"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "caaa9d531767d1ff2150b9332433f32a24622147e5ebb1f26409d5da67afd479"
+
+[[package]]
+name = "vec_map"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05c78687fb1a80548ae3250346c3db86a80a7cdd77bda190189f2d0a0987c81a"
+
+[[package]]
 name = "wasi"
 version = "0.9.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519"
+
+[[package]]
+name = "winapi"
+version = "0.3.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8093091eeb260906a183e6ae1abdba2ef5ef2257a21801128899c3fc699229c6"
+dependencies = [
+ "winapi-i686-pc-windows-gnu",
+ "winapi-x86_64-pc-windows-gnu",
+]
+
+[[package]]
+name = "winapi-i686-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
+
+[[package]]
+name = "winapi-x86_64-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
diff --git a/Cargo.toml b/Cargo.toml
index 902c3a3..de65bb7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,6 +5,9 @@ authors = ["Tavian Barnes <tavianator@tavianator.com>"]
 edition = "2018"
 
 [dependencies]
+clap = "2.33.0"
 image = "0.23.4"
 ordered-float = "1.0.2"
 rand = "0.7.3"
+rand_pcg = "0.2.1"
+term = "0.6.1"
diff --git a/src/main.rs b/src/main.rs
index 07e138a..f016b4c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,4 +3,399 @@ pub mod frontier;
 pub mod hilbert;
 pub mod metric;
 
-fn main() {}
+use crate::color::source::{AllColors, ColorSource, ImageColors};
+use crate::color::{order, ColorSpace, LabSpace, LuvSpace, Rgb8, RgbSpace};
+use crate::frontier::image::ImageFrontier;
+use crate::frontier::mean::MeanFrontier;
+use crate::frontier::min::MinFrontier;
+use crate::frontier::Frontier;
+
+use clap::{self, clap_app, crate_authors, crate_name, crate_version, AppSettings};
+
+use image::{self, Rgba, RgbaImage};
+
+use rand::SeedableRng;
+use rand_pcg::Pcg64;
+
+use term;
+
+use std::cmp;
+use std::error::Error;
+use std::fs;
+use std::io::{self, Write};
+use std::path::PathBuf;
+use std::str::FromStr;
+use std::time::Instant;
+
+/// The color source specified on the command line.
+#[derive(Debug, Eq, PartialEq)]
+enum SourceArg {
+    /// All RGB colors of the given bit depth.
+    AllRgb(u32),
+    /// Take the colors from an image.
+    Image(PathBuf),
+}
+
+/// The order to process colors in.
+#[derive(Debug, Eq, PartialEq)]
+enum OrderArg {
+    /// Sorted by hue.
+    HueSort,
+    /// Shuffled randomly.
+    Random,
+    /// Morton/Z-order.
+    Morton,
+    /// Hilbert curve order.
+    Hilbert,
+}
+
+/// The frontier implementation.
+#[derive(Debug, Eq, PartialEq)]
+enum FrontierArg {
+    /// Pick a neighbor of the closest pixel so far.
+    Min,
+    /// Pick the pixel with the closest mean color of all its neighbors.
+    Mean,
+    /// Target the closest pixel on an image.
+    Image(PathBuf),
+}
+
+/// The color space to operate in.
+#[derive(Debug, Eq, PartialEq)]
+enum ColorSpaceArg {
+    /// sRGB space.
+    Rgb,
+    /// CIE L*a*b* space.
+    Lab,
+    /// CIE L*u*v* space.
+    Luv,
+}
+
+/// Parse an argument into the appropriate type.
+fn parse_arg<F>(arg: Option<&str>) -> clap::Result<Option<F>>
+where
+    F: FromStr,
+    F::Err: Error,
+{
+    match arg.map(|s| s.parse()) {
+        Some(Ok(f)) => Ok(Some(f)),
+        Some(Err(e)) => Err(clap::Error::with_description(
+            &e.to_string(),
+            clap::ErrorKind::InvalidValue,
+        )),
+        None => Ok(None),
+    }
+}
+
+/// The parsed command line arguments.
+#[derive(Debug)]
+struct Args {
+    source: SourceArg,
+    order: OrderArg,
+    stripe: bool,
+    frontier: FrontierArg,
+    space: ColorSpaceArg,
+    width: Option<u32>,
+    height: Option<u32>,
+    x0: Option<u32>,
+    y0: Option<u32>,
+    animate: bool,
+    output: PathBuf,
+    seed: u64,
+}
+
+impl Args {
+    fn parse() -> clap::Result<Self> {
+        let args = clap_app!((crate_name!()) =>
+            (version: crate_version!())
+            (author: crate_authors!())
+            (setting: AppSettings::ColoredHelp)
+            (@group source =>
+                (@arg DEPTH: -b --("bit-depth") +takes_value "Use all DEPTH-bit colors")
+                (@arg INPUT: -i --input +takes_value "Use colors from the INPUT image")
+            )
+            (@group order =>
+                (@arg HUE: -s --hue-sort "Sort colors by hue [default]")
+                (@arg RANDOM: -r --random "Randomize colors")
+                (@arg MORTON: -M --morton "Place colors in Morton order (Z-order)")
+                (@arg HILBERT: -H --hilbert "Place colors in Hilbert curve order")
+            )
+            (@group stripe =>
+                (@arg STRIPE: -t --stripe "Reduce artifacts by iterating through the colors in multiple stripes [default]")
+                (@arg NOSTRIPE: -T --("no-stripe") "Don't stripe")
+            )
+            (@group frontier =>
+                (@arg MODE: -l --selection +takes_value possible_value[min mean] "Specify the selection mode")
+                (@arg TARGET: -g --target +takes_value "Place colors on the closest pixels of the TARGET image")
+            )
+            (@arg SPACE: -c --("color-space") default_value("Lab") possible_value[RGB Lab Luv] "Use the given color space")
+            (@arg WIDTH: -w --width +takes_value "The width of the generated image")
+            (@arg HEIGHT: -h --height +takes_value "The height of the generated image")
+            (@arg X: -x +takes_value "The x coordinate of the first pixel")
+            (@arg Y: -y +takes_value "The y coordinate of the first pixel")
+            (@arg ANIMATE: -a --animate "Generate frames of an animation")
+            (@arg PATH: -o --output default_value("kd-forest.png") "Save the image to PATH")
+            (@arg SEED: -e --seed default_value("0") "Seed the random number generator")
+        )
+        .get_matches_safe()?;
+
+        let source = if let Some(input) = args.value_of("INPUT") {
+            SourceArg::Image(PathBuf::from(input))
+        } else {
+            SourceArg::AllRgb(parse_arg(args.value_of("DEPTH"))?.unwrap_or(24))
+        };
+
+        let order = if args.is_present("RANDOM") {
+            OrderArg::Random
+        } else if args.is_present("MORTON") {
+            OrderArg::Morton
+        } else if args.is_present("HILBERT") {
+            OrderArg::Hilbert
+        } else {
+            OrderArg::HueSort
+        };
+
+        let stripe = !args.is_present("NOSTRIPE") && order != OrderArg::Random;
+
+        let frontier = if let Some(target) = args.value_of("TARGET") {
+            FrontierArg::Image(PathBuf::from(target))
+        } else {
+            match args.value_of("MODE") {
+                Some("min") | None => FrontierArg::Min,
+                Some("mean") => FrontierArg::Mean,
+                _ => unreachable!(),
+            }
+        };
+
+        let space = match args.value_of("SPACE").unwrap() {
+            "RGB" => ColorSpaceArg::Rgb,
+            "Lab" => ColorSpaceArg::Lab,
+            "Luv" => ColorSpaceArg::Luv,
+            _ => unreachable!(),
+        };
+
+        let width = parse_arg(args.value_of("WIDTH"))?;
+        let height = parse_arg(args.value_of("HEIGHT"))?;
+        let x0 = parse_arg(args.value_of("X"))?;
+        let y0 = parse_arg(args.value_of("Y"))?;
+
+        let animate = args.is_present("ANIMATE");
+
+        let mut path = args.value_of("PATH").unwrap();
+        if animate && args.occurrences_of("PATH") == 0 {
+            path = "kd-frames";
+        }
+        let output = PathBuf::from(path);
+
+        let seed = parse_arg(args.value_of("SEED"))?.unwrap_or(0);
+
+        Ok(Self {
+            source,
+            order,
+            stripe,
+            frontier,
+            space,
+            width,
+            height,
+            x0,
+            y0,
+            animate,
+            output,
+            seed,
+        })
+    }
+}
+
+/// main() return type.
+type MainResult = Result<(), Box<dyn Error>>;
+
+/// The kd-forest application itself.
+#[derive(Debug)]
+struct App {
+    args: Args,
+    rng: Pcg64,
+    width: Option<u32>,
+    height: Option<u32>,
+    start_time: Instant,
+}
+
+impl App {
+    /// Make the App.
+    fn new(args: Args) -> Self {
+        let rng = Pcg64::seed_from_u64(args.seed);
+        let width = args.width;
+        let height = args.height;
+        let start_time = Instant::now();
+
+        Self {
+            args,
+            rng,
+            width,
+            height,
+            start_time,
+        }
+    }
+
+    fn run(&mut self) -> MainResult {
+        let colors = match self.args.source {
+            SourceArg::AllRgb(depth) => {
+                self.width.get_or_insert(1u32 << ((depth + 1) / 2));
+                self.height.get_or_insert(1u32 << (depth / 2));
+                self.get_colors(AllColors::new(depth as usize))
+            }
+            SourceArg::Image(ref path) => {
+                let img = image::open(path)?.into_rgb();
+                self.width.get_or_insert(img.width());
+                self.height.get_or_insert(img.height());
+                self.get_colors(ImageColors::from(img))
+            }
+        };
+
+        match self.args.space {
+            ColorSpaceArg::Rgb => self.paint::<RgbSpace>(colors),
+            ColorSpaceArg::Lab => self.paint::<LabSpace>(colors),
+            ColorSpaceArg::Luv => self.paint::<LuvSpace>(colors),
+        }
+    }
+
+    fn get_colors<S: ColorSource>(&mut self, source: S) -> Vec<Rgb8> {
+        let colors = match self.args.order {
+            OrderArg::HueSort => order::hue_sorted(source),
+            OrderArg::Random => order::shuffled(source, &mut self.rng),
+            OrderArg::Morton => order::morton(source),
+            OrderArg::Hilbert => order::hilbert(source),
+        };
+
+        if self.args.stripe {
+            order::striped(colors)
+        } else {
+            colors
+        }
+    }
+
+    fn paint<C: ColorSpace>(&mut self, colors: Vec<Rgb8>) -> MainResult {
+        let width = self.width.unwrap();
+        let height = self.height.unwrap();
+        let x0 = self.args.x0.unwrap_or(width / 2);
+        let y0 = self.args.x0.unwrap_or(height / 2);
+
+        match &self.args.frontier {
+            FrontierArg::Image(ref path) => {
+                let img = image::open(path)?.into_rgb();
+                self.paint_on(colors, ImageFrontier::<C>::new(&img))
+            }
+            FrontierArg::Min => {
+                let rng = Pcg64::from_rng(&mut self.rng)?;
+                self.paint_on(colors, MinFrontier::<C, _>::new(rng, width, height, x0, y0))
+            }
+            FrontierArg::Mean => {
+                self.paint_on(colors, MeanFrontier::<C>::new(width, height, x0, y0))
+            }
+        }
+    }
+
+    fn paint_on<F: Frontier>(&mut self, colors: Vec<Rgb8>, mut frontier: F) -> MainResult {
+        let width = frontier.width();
+        let height = frontier.height();
+        let mut output = RgbaImage::new(width, height);
+
+        let size = cmp::min((width * height) as usize, colors.len());
+        println!("Generating a {}x{} image ({} pixels)", width, height, size);
+
+        if self.args.animate {
+            fs::create_dir_all(&self.args.output)?;
+            output.save(&self.args.output.join("0000.png"))?;
+        }
+
+        let interval = cmp::max(width, height) as usize;
+
+        let mut max_frontier = frontier.len();
+
+        for (i, color) in colors.into_iter().enumerate() {
+            let pos = frontier.place(color);
+            if pos.is_none() {
+                break;
+            }
+
+            let (x, y) = pos.unwrap();
+            let rgba = Rgba([color[0], color[1], color[2], 255]);
+            output.put_pixel(x, y, rgba);
+
+            max_frontier = cmp::max(max_frontier, frontier.len());
+
+            if (i + 1) % interval == 0 {
+                if self.args.animate {
+                    let frame = (i + 1) / interval;
+                    output.save(&self.args.output.join(format!("{:04}.png", frame)))?;
+                }
+
+                if i + 1 < size {
+                    self.print_progress(i + 1, size, frontier.len())?;
+                }
+            }
+        }
+
+        if self.args.animate && size % interval != 0 {
+            let frame = size / interval;
+            output.save(&self.args.output.join(format!("{:04}.png", frame)))?;
+        }
+
+        self.print_progress(size, size, max_frontier)?;
+
+        if !self.args.animate {
+            output.save(&self.args.output)?;
+        }
+
+        Ok(())
+    }
+
+    fn print_progress(&self, i: usize, size: usize, frontier_len: usize) -> io::Result<()> {
+        let mut term = match term::stderr() {
+            Some(term) => term,
+            None => return Ok(()),
+        };
+
+        let progress = 100.0 * (i as f64) / (size as f64);
+        let mut rate = (i as f64) / self.start_time.elapsed().as_secs_f64();
+        let mut unit = "px/s";
+
+        if rate >= 10_000.0 {
+            rate /= 1_000.0;
+            unit = "Kpx/s";
+        }
+
+        if rate >= 10_000.0 {
+            rate /= 1_000.0;
+            unit = "Mpx/s";
+        }
+
+        if rate >= 10_000.0 {
+            rate /= 1_000.0;
+            unit = "Gpx/s";
+        }
+
+        let (frontier_label, newline) = if i == size {
+            ("max frontier size", "\n")
+        } else {
+            ("frontier size", "")
+        };
+
+        term.carriage_return()?;
+        term.delete_line()?;
+
+        write!(
+            term,
+            "{:>6.2}%  | {:4.0} {:>5}  | {}: {}{}",
+            progress, rate, unit, frontier_label, frontier_len, newline,
+        )
+    }
+}
+
+fn main() -> MainResult {
+    let args = match Args::parse() {
+        Ok(args) => args,
+        Err(e) => e.exit(),
+    };
+
+    App::new(args).run()
+}