You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

257 lines
9.1 KiB

use crate::{transform::{Transform, NCOLORS}, genetic::Genetic};
use core::fmt::Debug;
use std::sync::RwLock;
use rand::{Rng, distributions::{Uniform, WeightedIndex}};
use rand_distr::{Distribution, Bernoulli};
use image::{Rgb, GenericImage, ImageBuffer};
pub mod gpu;
#[derive(Debug)]
pub struct Map
{
transforms: Vec<Transform>,
cached_error: RwLock<Option<f32>>,
}
impl Clone for Map {
fn clone(&self) -> Self {
Self {
transforms: self.transforms.clone(),
cached_error: RwLock::new(None),
}
}
}
impl Map {
pub fn apply_add<I: GenericImage<Pixel=Rgb<f32>>, O: GenericImage<Pixel=Rgb<f32>>>(&self, input: &I, output: &mut O, dims: (u32, u32)) {
let offset = 1.0;
let c = (self.transforms.len() as f32).sqrt().recip() * 3.5;
for tf in self.transforms.iter() {
tf.apply_add(input, output, dims, c);
}
for x in 0..dims.0 {
for y in 0..dims.1 {
let Rgb([r, g, b]) = output.get_pixel(x, y);
output.put_pixel(x, y, Rgb([r + offset, g + offset, b + offset]));
}
}
}
pub fn error(&self, cfg: &MapGeneticConfig) -> f32 {
let cached = { *self.cached_error.read().unwrap() };
if let Some(c) = cached {
c
} else {
let mut input = ImageBuffer::from_fn(cfg.dims.0, cfg.dims.1, |_, _| Rgb([0.0, 0.0, 0.0]));
let mut output = input.clone();
for _ in 0..cfg.num_iters {
self.apply_add(&input, &mut output, cfg.dims);
core::mem::swap(&mut input, &mut output);
output = ImageBuffer::from_fn(input.width(), input.height(), |_, _| Rgb([0.0, 0.0, 0.0]));
}
// spooky least squares math, transcribed from my notebook where i solved
// this problem. basically, compute the minimum square error we can get
// if we are allowed to scale this image and add a constant to it.
let a_a = image_dot_product(&input, &input);
let sum_a = image_sum(&input);
let denom = a_a * cfg.ndof as f32 - sum_a * sum_a;
// not very stable, but i think that it doesn't matter here, since we
// just want to know if it's large, indicating deconvergence
let variance = a_a / cfg.ndof as f32 - sum_a * sum_a / (cfg.ndof * cfg.ndof) as f32;
if denom == 0.0 || variance > 50.0 {
// if the image is entirely zero or it doesn't appear to converge, return infinite cost
*self.cached_error.write().unwrap() = Some(f32::INFINITY);
f32::INFINITY
} else {
let a_x = image_dot_product(&input, &cfg.goal_image);
let numer = a_x * a_x * cfg.ndof as f32 - 2.0 * a_x * sum_a * cfg.goal_sum + cfg.goal_sum * cfg.goal_sum * a_a;
let cost = cfg.goal_dot - numer / denom;
*self.cached_error.write().unwrap() = Some(cost);
cost
}
}
}
pub fn cost(&self, cfg: &MapGeneticConfig) -> f32 {
self.error(cfg) + self.transforms.len() as f32 * cfg.cost_per_tf
}
}
#[derive(Debug)]
pub struct MapGeneticConfig
{
dims: (u32, u32),
pub initial_size: usize,
pub num_iters: usize,
pub mean_add: usize,
pub cost_per_tf: f32,
goal_image: ImageBuffer<Rgb<f32>, Vec<f32>>,
// dot product of goal image with itself
goal_dot: f32,
// goal image sum of all elements
goal_sum: f32,
// number of degrees of freedom (vector dimension)
ndof: usize,
}
impl Genetic for Map
{
type Configuration = MapGeneticConfig;
fn generate<R: Rng>(rng: &mut R, cfg: &Self::Configuration) -> Self {
cfg.create_map(rng, cfg.initial_size)
}
fn mutate<R: Rng>(&self, mut rng: &mut R, cfg: &Self::Configuration) -> Self {
let mut new = self.clone();
// first, remove elements
let mut accum = 1.0;
let threshold = (cfg.mean_add as f32).exp().recip();
accum *= rng.gen::<f32>();
while new.transforms.len() != 0 && accum > threshold {
let which = rng.gen_range(0..new.transforms.len());
new.transforms.swap_remove(which);
accum *= rng.gen::<f32>();
}
// now, change some of the matrices
let mut accum = 1.0;
let threshold = (cfg.mean_add as f32).exp().recip();
accum *= rng.gen::<f32>();
while accum > threshold {
let which = rng.gen_range(0..new.transforms.len());
let mut matrix_coeff_dist = Uniform::new(-1.0, 1.0).sample_iter(&mut rng);
for i in 0..NCOLORS * NCOLORS {
new.transforms[which].color_matrix[i] = matrix_coeff_dist.next().unwrap();
}
accum *= rng.gen::<f32>();
}
// add some new transforms
let mut accum = 1.0;
let threshold = (cfg.mean_add as f32).exp().recip();
accum *= rng.gen::<f32>();
while accum > threshold {
let transform = cfg.random_transform(rng);
new.transforms.push(transform);
accum *= rng.gen::<f32>();
}
new
}
fn crossover<R: Rng>(&self, rng: &mut R, other: &Self, _config: &Self::Configuration) -> Self {
let mut new = self.clone();
new.transforms.extend_from_slice(&other.transforms);
// some paper that i read said this works correctly for selecting (in this
// case, removing) a number of randomly selected elements
let mut num_remaining = self.transforms.len() + other.transforms.len();
let mut num_to_remove = num_remaining / 2;
let mut idx = 0;
for _ in 0..num_remaining {
let distr = Bernoulli::new(num_to_remove as f64 / num_remaining as f64).unwrap();
if distr.sample(rng) {
new.transforms.swap_remove(idx);
num_remaining -= 1;
num_to_remove -= 1;
} else {
idx += 1;
num_remaining -= 1;
}
}
new
}
fn compare<R: Rng>(&self, _rng: &mut R, other: &Self, config: &Self::Configuration) -> std::cmp::Ordering {
self.cost(&config).partial_cmp(&other.cost(&config)).unwrap_or(std::cmp::Ordering::Equal)
}
}
impl MapGeneticConfig {
pub fn new(goal_image: ImageBuffer<Rgb<f32>, Vec<f32>>, initial_size: usize, num_iters: usize, cost_per_tf: f32, mean_add: usize) -> Self {
let dims = goal_image.dimensions();
let goal_dot = image_dot_product(&goal_image, &goal_image);
let goal_sum = image_sum(&goal_image);
let ndof = (dims.0 * dims.1 * 3) as usize;
MapGeneticConfig {
dims,
initial_size,
num_iters,
cost_per_tf,
mean_add,
goal_image,
goal_dot,
goal_sum,
ndof,
}
}
pub fn create_map<R: Rng>(&self, rng: &mut R, tf_count: usize) -> Map {
let mut transforms = vec![];
for _ in 0..tf_count {
transforms.push(self.random_transform(rng));
}
Map {
transforms,
cached_error: RwLock::new(None),
}
}
pub fn random_transform<R: Rng>(&self, mut rng: &mut R) -> Transform {
let rect_xs = loop {
let first = rng.gen_range(0..self.dims.0);
let second = first as i32 + rng.gen_range(-(self.dims.0 as i32)..=self.dims.0 as i32);
let pair = (first, if second < 0 { 0 } else if second >= self.dims.0 as i32 { self.dims.0 - 1 } else { second as u32 });
if pair.1 < pair.0 {
break (pair.1, pair.0);
} else if pair.0 < pair.1 {
break pair;
}
};
let rect_ys = loop {
let first = rng.gen_range(0..self.dims.1);
let second = first as i32 + rng.gen_range(-(self.dims.1 as i32)..=self.dims.1 as i32);
let pair = (first, if second < 0 { 0 } else if second >= self.dims.1 as i32 { self.dims.1 - 1 } else { second as u32 });
if pair.1 < pair.0 {
break (pair.1, pair.0);
} else if pair.0 < pair.1 {
break pair;
}
};
let mut corner_dist = Uniform::new(0.0, 1.0).sample_iter(&mut rng);
let (cx, cy) = (corner_dist.next().unwrap(), corner_dist.next().unwrap());
let mut off_dist = Uniform::new(-0.5, 0.5).sample_iter(&mut rng);
let (e1x, e1y, e2x, e2y) = (off_dist.next().unwrap(), off_dist.next().unwrap(), off_dist.next().unwrap(), off_dist.next().unwrap());
let (to_corner, to_extents) = match WeightedIndex::new(&[1, 1, 1, 1]).unwrap().sample(rng) {
0 => ((cx, cy), [(e1x, e1y), (e2x, e2y)]),
1 => ((cx + e1x, cy + e1y), [(e2x, e2y), (-e1x, -e1y)]),
2 => ((cx + e1x + e2x, cy + e1y + e2y), [(-e1x, -e1y), (-e2x, -e2y)]),
3 => ((cx + e2x, cy + e2y), [(-e2x, -e2y), (e1x, e1y)]),
_ => unreachable!(),
};
let mut matrix_coeff_dist = Uniform::new(-1.0, 1.0).sample_iter(rng);
let mut color_matrix = [0.0f32; NCOLORS * NCOLORS];
for i in 0..NCOLORS * NCOLORS {
color_matrix[i] = matrix_coeff_dist.next().unwrap();
}
Transform {
color_matrix,
from_corner: (rect_xs.0, rect_ys.0),
from_extents: (rect_xs.1 - rect_xs.0, rect_ys.1 - rect_ys.0),
to_corner,
to_extents,
}
}
}
/// Treat the entire image as a vector, and find the dot product of two images.
pub fn image_dot_product(x: &ImageBuffer<Rgb<f32>, Vec<f32>>, y: &ImageBuffer<Rgb<f32>, Vec<f32>>) -> f32 {
// do the sum row-wise for numerical stability
x.rows().zip(y.rows()).map(|(x, y)| x.zip(y).map(|(&Rgb([xr, xg, xb]), &Rgb([yr, yg, yb]))| xr * yr + xg * yg + xb * yb).sum::<f32>()).sum()
}
/// Compute the sum of all the degrees of freedom (subpixel values) of an image. Equivalent to the dot product with the vector of all 1s.
pub fn image_sum(x: &ImageBuffer<Rgb<f32>, Vec<f32>>) -> f32 {
// do the sum row-wise for numerical stability
x.rows().map(|x| x.map(|&Rgb([r, g, b])| r + g + b).sum::<f32>()).sum()
}