pbnify 0.2.1

// Copyright (C) 2019  Adam Gausmann
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use clap::{crate_authors, crate_version, App, Arg};
use image::imageops::{resize, FilterType};
use image::{DynamicImage, ImageResult};
use std::fs;
use std::io::{stdin, Read};

use pbnify::quantize::{build_palette, quantize};
use pbnify::PbnConfig;

fn load_read<R: Read>(mut read: R) -> ImageResult<DynamicImage> {
    let mut buf = Vec::new();
    read.read_to_end(&mut buf)?;
    image::load_from_memory(&buf)
}

fn main() {
    let default_config = PbnConfig::default();
    let default_x_offset = default_config.x_offset.to_string();
    let default_y_offset = default_config.y_offset.to_string();
    let default_max_length = default_config.max_length.to_string();

    let args = App::new("PBNify")
        .version(crate_version!())
        .author(crate_authors!(", "))
        .about("Converts images into a series of PBN (Paint By Numbers) commands.")
        .arg(
            Arg::with_name("input")
                .short("i")
                .long("input")
                .takes_value(true)
                .value_name("FILE")
                .help("Tells the program to accept the image from FILE. If this option is not present, standard input will be read and PNG is assumed.")
        )
        .arg(
            Arg::with_name("quantize")
                .short("n")
                .long("quantize")
                .takes_value(true)
                .value_name("SIZE")
                .help("Generates a palette of size SIZE, and quantizes the input image before PBNification.")
        )
        .arg(
            Arg::with_name("x-offset")
                .short("x")
                .long("x-offset")
                .takes_value(true)
                .value_name("NUM")
                .help("Offsets the output x coordinates by the given amount.")
                .default_value(&default_x_offset)
        )
        .arg(
            Arg::with_name("y-offset")
                .short("y")
                .long("y-offset")
                .takes_value(true)
                .value_name("NUM")
                .help("Offsets the output y coordinates by the given amount.")
                .default_value(&default_y_offset)
        )
        .arg(
            Arg::with_name("width")
                .short("w")
                .long("width")
                .takes_value(true)
                .value_name("WIDTH")
                .help("Rescales the picture to WIDTH pixels wide. Will preserve the original ratio unless --height is also provided.")
        )
        .arg(
            Arg::with_name("height")
                .short("H")
                .long("height")
                .takes_value(true)
                .value_name("HEIGHT")
                .help("Rescales the picture to HEIGHT pixels tall. Will preserve the original ratio unless --width is also provided.")
        )
        .arg(
            Arg::with_name("max-length")
                .short("l")
                .long("max-length")
                .takes_value(true)
                .value_name("NUM")
                .help("Sets the maximum length of a command.")
                .default_value(&default_max_length)
        )
        .arg(
            Arg::with_name("output")
                .short("o")
                .long("output")
                .takes_value(true)
                .value_name("FILE")
                .help("Tells the program to write the generated output to FILE. If this option is not present, it will be written to standard output.")
        )
        .get_matches();

    let mut image_buffer = match args.value_of_os("input") {
        Some(file) => image::open(file),
        None => load_read(stdin()),
    }
    .expect("An error occurred while loading the image")
    .to_rgba();

    let resize_params = match (
        args.value_of("width")
            .map(|x| x.parse().expect("width is not a valid integer")),
        args.value_of("height")
            .map(|x| x.parse().expect("height is not a valid integer")),
    ) {
        (None, None) => None,
        (Some(width), Some(height)) => Some((width, height)),
        (Some(width), None) => Some((width, width * image_buffer.height() / image_buffer.width())),
        (None, Some(height)) => Some((
            height * image_buffer.width() / image_buffer.height(),
            height,
        )),
    };
    if let Some((width, height)) = resize_params {
        image_buffer = resize(&image_buffer, width, height, FilterType::Lanczos3);
    }

    if let Some(size) = args.value_of("quantize") {
        let size = size.parse().expect("palette size is not a valid integer");

        let palette = build_palette(size, image_buffer.pixels().cloned());
        quantize(&mut image_buffer, &palette);
    }

    let x_offset: u32 = args
        .value_of("x-offset")
        .unwrap()
        .parse()
        .expect("x offset is not a valid integer");
    let y_offset: u32 = args
        .value_of("y-offset")
        .unwrap()
        .parse()
        .expect("y offset is not a valid integer");
    let max_length: usize = args
        .value_of("max-length")
        .unwrap()
        .parse()
        .expect("max length is not a valid integer");

    let commands = PbnConfig::new()
        .x_offset(x_offset)
        .y_offset(y_offset)
        .max_length(max_length)
        .run(&image_buffer);

    let output_string: String = commands
        .iter()
        .flat_map(|s| vec![&s[..], "\n"].into_iter())
        .collect();

    match args.value_of_os("output") {
        Some(file) => fs::write(file, &output_string).expect("Failed to write to file"),
        None => print!("{}", output_string),
    }
}