structix/themis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: use better calculations for distance and average color

Browse Source
This commit is contained in:
PurpleMyst
2022-07-15 01:22:01 +02:00
parent d7c5bf5121
commit 5a9b206e5b
6 changed files with 393 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

162
src/main.rs
View File

@@ -2,10 +2,11 @@ use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::{Path, PathBuf};
use anyhow::Result;
use image::{DynamicImage, GenericImage, GenericImageView, Pixel, Rgba};
use indicatif::{ParallelProgressIterator, ProgressBar, ProgressIterator, ProgressStyle};
use noisy_float::prelude::*;
use eyre::Result;
use image::{DynamicImage, GenericImage, GenericImageView, Rgba};
use indicatif::{
ParallelProgressIterator, ProgressBar, ProgressFinish, ProgressIterator, ProgressStyle,
};
use rayon::prelude::*;
use structopt::StructOpt;
@@ -15,28 +16,26 @@ fn average_color(image: &DynamicImage) -> Rgba<u8> {
let (mut r, mut g, mut b, mut a) = (0., 0., 0., 0.);
for (_x, _y, Rgba([pr, pg, pb, pa])) in image.pixels() {
r += pr as f64;
g += pg as f64;
b += pb as f64;
a += pa as f64;
r += pr as f64 * pr as f64;
g += pg as f64 * pg as f64;
b += pb as f64 * pb as f64;
a += pa as f64 * pa as f64;
}
let r = (r / pixel_count) as u8;
let g = (g / pixel_count) as u8;
let b = (b / pixel_count) as u8;
let a = (a / pixel_count) as u8;
let r = (r / pixel_count).sqrt() as u8;
let g = (g / pixel_count).sqrt() as u8;
let b = (b / pixel_count).sqrt() as u8;
let a = (a / pixel_count).sqrt() as u8;
Rgba([r, g, b, a])
}
/// Calculate the euclidean distance between two pixels
fn distance(pixel1: Rgba<u8>, pixel2: Rgba<u8>) -> R64 {
r64(pixel1
.map2(&pixel2, |l, r| if l < r { r - l } else { l - r })
.channels()
.iter()
.map(|&n| (n as f64).powi(2))
.sum::<f64>()
.sqrt())
/// Calculate the distance (squared) between two colors
/// Code adapted from https://stackoverflow.com/a/9085524/13204109
fn distance(Rgba([r1, g1, b1, _]): Rgba<u8>, Rgba([r2, g2, b2, _]): Rgba<u8>) -> i64 {
let rmean = (i64::from(r1) + i64::from(r2)) / 2;
let r = i64::from(r1) - i64::from(r2);
let g = i64::from(g1) - i64::from(g2);
let b = i64::from(b1) - i64::from(b2);
(((512 + rmean) * r * r) >> 8) + 4 * g * g + (((767 - rmean) * b * b) >> 8)
}
/// Choose the image in the given tileset whose average color is closest to the given pixel
@@ -54,6 +53,7 @@ fn pick_image_for_pixel(
fn load_images<P: AsRef<Path>>(dir: P, tile_side: u32) -> Result<Vec<(DynamicImage, Rgba<u8>)>> {
let dir = fs::read_dir(dir)?.collect::<Result<Vec<_>, _>>()?;
let len = dir.len();
Ok(dir
.into_par_iter()
.progress_with(make_pbar("images loaded", len as _))
@@ -68,37 +68,52 @@ fn load_images<P: AsRef<Path>>(dir: P, tile_side: u32) -> Result<Vec<(DynamicIma
}
/// Create a styled progress bar
fn make_pbar(msg: &str, len: u64) -> ProgressBar {
fn make_pbar(msg: &'static str, len: u64) -> ProgressBar {
let bar = ProgressBar::new(len);
bar.set_message(msg);
bar.set_style(
ProgressStyle::default_bar()
.template("[{elapsed_precise}/{eta_precise}] {bar:40.green/red} {pos:>4}/{len:4} {msg}")
.progress_chars("##-"),
.progress_chars("##-")
.on_finish(ProgressFinish::AndLeave),
);
bar
}
/// Create a styled spinner
fn make_spinner(msg: &'static str, done_msg: &'static str) -> ProgressBar {
let bar = ProgressBar::new(0);
bar.set_message(msg);
bar.set_style(
ProgressStyle::default_bar()
.template("{spinner:.yellow} {msg}")
.progress_chars("##-")
.on_finish(ProgressFinish::WithMessage(done_msg)),
);
bar.enable_steady_tick(100);
bar
}
#[derive(StructOpt)]
struct Opt {
/// The image to turn into a mosaic
#[structopt(parse(from_os_str))]
image: PathBuf,
#[structopt(short, long, parse(from_os_str))]
input_dir: PathBuf,
/// The directory containing the tiles to utilize
#[structopt(parse(from_os_str))]
tiles_directory: PathBuf,
#[structopt(short, long, parse(from_os_str))]
tiles_dir: PathBuf,
/// Where to save the finished mosaic
#[structopt(short, long, parse(from_os_str), default_value = "mosaic.png")]
output: PathBuf,
#[structopt(short, long, parse(from_os_str), default_value = "output")]
output_dir: PathBuf,
/// The side length that the target image'll be resized to.
#[structopt(short, long, default_value = "128")]
mosaic_size: u32,
/// The side length of each tile.
#[structopt(short, long, default_value = "26")]
#[structopt(long, default_value = "26")]
tile_size: u32,
/// Keep the image's aspect ratio
@@ -108,46 +123,67 @@ struct Opt {
fn main() -> Result<()> {
let Opt {
image,
tiles_directory,
input_dir,
tiles_dir,
mosaic_size,
tile_size,
keep_aspect_ratio,
output,
output_dir,
} = Opt::from_args();
let possible_tiles = load_images(tiles_directory, tile_size)?;
let image = image::open(image)?;
let image = if keep_aspect_ratio {
image.thumbnail(mosaic_size, mosaic_size)
} else {
image.thumbnail_exact(mosaic_size, mosaic_size)
};
fs::create_dir_all(&output_dir)?;
// For every unique pixel in the image, find its most appropiate tile
let unique_pixels = image
.pixels()
.map(|(_x, _y, pixel)| pixel)
.collect::<HashSet<_>>();
let pbar = make_pbar("pixels", unique_pixels.len() as _);
let tiles = unique_pixels
.into_par_iter()
.progress_with(pbar)
.filter_map(|pixel| {
let tile = pick_image_for_pixel(pixel, &possible_tiles)?;
Some((pixel, tile))
})
.collect::<HashMap<_, _>>();
let possible_tiles = load_images(tiles_dir, tile_size)?;
// Apply the mapping previously calculated and save the mosaic
let mut mosaic = DynamicImage::new_rgba8(image.width() * tile_size, image.height() * tile_size);
for (x, y, pixel) in image.pixels().progress_with(make_pbar(
"actual pixels",
u64::from(image.width() * image.height()),
)) {
mosaic.copy_from(&**tiles.get(&pixel).unwrap(), x * tile_size, y * tile_size)?;
let input_dir = fs::read_dir(input_dir)?.collect::<Result<Vec<_>, _>>()?;
for image in input_dir {
let input_path = image.path();
eprintln!("Processing {}", input_path.display());
let output = output_dir.join(format!(
"{}.mosaic{mosaic_size}.png",
input_path.file_stem().unwrap().to_string_lossy()
));
if output.exists() {
continue;
}
let img = image::open(&input_path)?;
let img = if keep_aspect_ratio {
img.thumbnail(mosaic_size, mosaic_size)
} else {
img.thumbnail_exact(mosaic_size, mosaic_size)
};
// For every unique pixel in the image, find its most appropiate tile
let unique_pixels = img
.pixels()
.map(|(_x, _y, pixel)| pixel)
.collect::<HashSet<_>>();
let len = unique_pixels.len();
let tiles = unique_pixels
.into_par_iter()
.progress_with(make_pbar("pixels", len as _))
.filter_map(|pixel| {
let tile = pick_image_for_pixel(pixel, &possible_tiles)?;
Some((pixel, tile))
})
.collect::<HashMap<_, _>>();
// Apply the mapping previously calculated and save the mosaic
let mut mosaic = DynamicImage::new_rgba8(img.width() * tile_size, img.height() * tile_size);
for (x, y, pixel) in img.pixels().progress_with(make_pbar(
"actual pixels",
u64::from(img.width() * img.height()),
)) {
mosaic.copy_from(&**tiles.get(&pixel).unwrap(), x * tile_size, y * tile_size)?;
}
let spinner = make_spinner("Saving", "Saved!");
mosaic.save(output)?;
spinner.finish_using_style();
}
mosaic.save(output)?;
Ok(())
}