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.
 
 

108 lines
4.8 KiB

#![allow(confusable_idents)]
#![allow(mixed_script_confusables)]
use image::{ImageBuffer, Rgb, GenericImage};
mod transform;
mod genetic;
mod map;
use crate::map::MapGeneticConfig;
//use transform::Transform;
fn main() {
let args: Vec<_> = std::env::args().collect();
let image_in_u16: ImageBuffer<Rgb<u16>, Vec<u16>> = image::open(args.get(1).expect("input file not provided")).expect("could not open input file").to_rgb16();
let out_name = args.get(2).expect("output file not provided");
let dims = image_in_u16.dimensions();
let image_in: ImageBuffer<Rgb<f32>, Vec<f32>> = ImageBuffer::from_fn(dims.0, dims.1, |x, y| {
let pix = image_in_u16.get_pixel(x, y);
Rgb([pix[0] as f32 / 65535.0, pix[1] as f32 / 65535.0, pix[2] as f32 / 65535.0])
});
let gpu_ctx = map::gpu::init(&image_in);
if let Some(_) = gpu_ctx {
println!("successfully created a GPU context");
}
let mut rng = rand::thread_rng();
let cfg = MapGeneticConfig::new(
image_in.clone(),
100,
20,
0.1,
10,
);
let mut arena = genetic::Arena::<map::Map>::new(&mut rng, 1000, 250, 250, cfg);
let mut i = 0;
while arena.set[0].error(&arena.config) > 50.0 {
println!("iteration {}", i);
arena.generate(&mut rng);
arena.par_process(|map, cfg| { map.cost(cfg); });
arena.par_cull();
let map = &arena.set[0];
println!("best cost and error: {}, {}", map.cost(&arena.config), map.error(&arena.config));
let last_map = arena.set.last().unwrap();
println!("worst cost and error: {}, {}", last_map.cost(&arena.config), last_map.error(&arena.config));
let mut image_in: ImageBuffer<Rgb<f32>, Vec<f32>> = ImageBuffer::from_fn(dims.0, dims.1, |_, _| Rgb([0.0, 0.0, 0.0]));
let mut image_out = image_in.clone();
for _ in 0..20 {
map.apply_add(&image_in, &mut image_out, dims);
// switch buffers; reuse as input
core::mem::swap(&mut image_out, &mut image_in);
image_out = ImageBuffer::from_fn(dims.0, dims.1, |_, _| Rgb([0.0, 0.0, 0.0]));
}
correct_image_brightness(&mut image_in, dims, 0.5, 0.15);
let image_out_u16: ImageBuffer<Rgb<u16>, Vec<u16>> = ImageBuffer::from_fn(dims.0, dims.1, |x, y| {
let pix = image_in.get_pixel(x, y);
let scale = 65535.0;
Rgb([(pix[0] * scale) as u16, (pix[1] * scale) as u16, (pix[2] * scale) as u16])
});
image_out_u16.save(format!("{}{}.png", out_name, i)).expect("could not write output file");
i += 1;
}
// let map = cfg.create_map(&mut rng, 50);
// for i in 0..50 {
// map.apply_add(&image_in, &mut image_out, dims);
//
// if true || i % 25 == 0 {
// // write intermediate to file
// let mut image_out_file = image_out.clone();
// correct_image_brightness(&mut image_out_file, dims, 0.5, 1.0);
// let image_out_u16: ImageBuffer<Rgb<u16>, Vec<u16>> = ImageBuffer::from_fn(dims.0, dims.1, |x, y| {
// let pix = image_out.get_pixel(x, y);
// let scale = 65535.0;
// Rgb([(pix[0] * scale) as u16, (pix[1] * scale) as u16, (pix[2] * scale) as u16])
// });
// image_out_u16.save(format!("{}{}.png", out_name, i)).expect("could not write output file");
// }
//
// // switch buffers; reuse as input
// core::mem::swap(&mut image_out, &mut image_in);
// image_out = ImageBuffer::from_fn(dims.0, dims.1, |_, _| Rgb([0.0, 0.0, 0.0]));
// }
//
// map.apply_add(&image_in, &mut image_out, dims);
// correct_image_brightness(&mut image_out, dims, 0.5, 0.15);
// let image_out_u16: ImageBuffer<Rgb<u16>, Vec<u16>> = ImageBuffer::from_fn(dims.0, dims.1, |x, y| {
// let pix = image_out.get_pixel(x, y);
// let scale = 65535.0;
// Rgb([(pix[0] * scale) as u16, (pix[1] * scale) as u16, (pix[2] * scale) as u16])
// });
// image_out_u16.save(format!("{}.png", out_name)).expect("could not write output file");
}
fn correct_image_brightness<I: GenericImage<Pixel=Rgb<f32>>>(img: &mut I, dims: (u32, u32), mean: f32, variance: f32) {
// add row-wise for numerical stability
let current_mean = (0..dims.1).map(|y| (0..dims.0).map(|x| { let c = img.get_pixel(x, y); c.0[0] as f64 + c.0[1] as f64 + c.0[2] as f64 }).sum::<f64>()).sum::<f64>() / (3 * dims.0 * dims.1) as f64;
let current_variance_recip = (3 * dims.0 * dims.1) as f64 / (0..dims.1).map(|y| (0..dims.0).map(|x| { let c = img.get_pixel(x, y); let (r, g, b) = (c.0[0] as f64 - current_mean, c.0[1] as f64 - current_mean, c.0[2] as f64 - current_mean); r * r + g * g + b * b }).sum::<f64>()).sum::<f64>();
println!("{} {}", current_mean, current_variance_recip);
let correct = |x| (x - current_mean as f32) * current_variance_recip as f32 * variance + mean;
for x in 0..dims.0 {
for y in 0..dims.1 {
let Rgb([r, g, b]) = img.get_pixel(x, y);
img.put_pixel(x, y, Rgb([correct(r), correct(g), correct(b)]));
}
}
}