mod img_pyramid;
use img_pyramid::*;
mod utils;
use utils::*;
mod multires_stochastic_texture_synthesis;
use multires_stochastic_texture_synthesis::*;
use std::path::Path;
pub use image;
pub use utils::ImageSource;
pub type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
struct Parameters {
tiling_mode: bool,
nearest_neighbors: u32,
random_sample_locations: u64,
cauchy_dispersion: f32,
backtrack_percent: f32,
backtrack_stages: u32,
resize_input: Option<(u32, u32)>,
output_size: (u32, u32),
guide_alpha: f32,
random_resolve: Option<u64>,
seed: u64,
}
impl Default for Parameters {
fn default() -> Self {
Self {
tiling_mode: false,
nearest_neighbors: 50,
random_sample_locations: 50,
cauchy_dispersion: 1.0,
backtrack_percent: 0.5,
backtrack_stages: 5,
resize_input: None,
output_size: (500, 500),
guide_alpha: 0.8,
random_resolve: None,
seed: 0,
}
}
}
impl Parameters {
fn to_generator_params(&self) -> GeneratorParams {
GeneratorParams {
nearest_neighbors: self.nearest_neighbors,
random_sample_locations: self.random_sample_locations,
cauchy_dispersion: self.cauchy_dispersion,
p: self.backtrack_percent,
p_stages: self.backtrack_stages as i32,
seed: self.seed,
alpha: self.guide_alpha,
}
}
}
pub struct GeneratedImage {
inner: multires_stochastic_texture_synthesis::Generator,
}
impl GeneratedImage {
pub fn save<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
let path = path.as_ref();
if let Some(parent_path) = path.parent() {
std::fs::create_dir_all(&parent_path)?;
}
self.inner.color_map.save(&path)?;
Ok(())
}
pub fn write<W: std::io::Write>(
self,
writer: &mut W,
fmt: image::ImageOutputFormat,
) -> Result<(), Error> {
let dyn_img = self.into_image();
Ok(dyn_img.write_to(writer, fmt)?)
}
pub fn save_debug<P: AsRef<Path>>(&self, dir: P) -> Result<(), Error> {
let dir = dir.as_ref();
std::fs::create_dir_all(&dir)?;
self.inner
.get_uncertainty_map()
.save(&dir.join("uncertainty.png"))?;
let id_maps = self.inner.get_id_maps();
id_maps[0].save(&dir.join("patch_id.png"))?;
id_maps[1].save(&dir.join("map_id.png"))?;
Ok(())
}
pub fn into_image(self) -> image::DynamicImage {
image::DynamicImage::ImageRgba8(self.inner.color_map)
}
}
impl AsRef<image::RgbaImage> for GeneratedImage {
fn as_ref(&self) -> &image::RgbaImage {
&self.inner.color_map
}
}
pub enum SampleMethod<'a> {
All,
Ignore,
Image(ImageSource<'a>),
}
impl<'a> SampleMethod<'a> {
fn is_ignore(&self) -> bool {
match self {
Self::Ignore => true,
_ => false,
}
}
}
impl<'a, S> From<&'a S> for SampleMethod<'a>
where
S: AsRef<Path> + 'a,
{
fn from(path: &'a S) -> Self {
Self::Image(ImageSource::Path(path.as_ref()))
}
}
pub enum SamplingMethod {
All,
Ignore,
Image(image::RgbaImage),
}
impl SamplingMethod {
fn is_ignore(&self) -> bool {
match self {
Self::Ignore => true,
_ => false,
}
}
}
pub struct ExampleBuilder<'a> {
img: ImageSource<'a>,
guide: Option<ImageSource<'a>>,
sample_method: SampleMethod<'a>,
}
impl<'a> ExampleBuilder<'a> {
pub fn new<I: Into<ImageSource<'a>>>(img: I) -> Self {
Self {
img: img.into(),
guide: None,
sample_method: SampleMethod::All,
}
}
pub fn with_guide<G: Into<ImageSource<'a>>>(mut self, guide: G) -> Self {
self.guide = Some(guide.into());
self
}
pub fn set_sample_method<M: Into<SampleMethod<'a>>>(mut self, method: M) -> Self {
self.sample_method = method.into();
self
}
}
impl<'a> Into<Example<'a>> for ExampleBuilder<'a> {
fn into(self) -> Example<'a> {
Example {
img: self.img,
guide: self.guide,
sample_method: self.sample_method,
}
}
}
pub struct Example<'a> {
img: ImageSource<'a>,
guide: Option<ImageSource<'a>>,
sample_method: SampleMethod<'a>,
}
impl<'a> Example<'a> {
pub fn builder<I: Into<ImageSource<'a>>>(img: I) -> ExampleBuilder<'a> {
ExampleBuilder::new(img)
}
pub fn new<I: Into<ImageSource<'a>>>(img: I) -> Self {
Self {
img: img.into(),
guide: None,
sample_method: SampleMethod::All,
}
}
pub fn with_guide<G: Into<ImageSource<'a>>>(&mut self, guide: G) -> &mut Self {
self.guide = Some(guide.into());
self
}
pub fn set_sample_method<M: Into<SampleMethod<'a>>>(&mut self, method: M) -> &mut Self {
self.sample_method = method.into();
self
}
fn resolve(
self,
backtracks: u32,
resize: Option<(u32, u32)>,
target_guide: &Option<ImagePyramid>,
) -> Result<ResolvedExample, Error> {
let image = ImagePyramid::new(load_image(self.img, resize)?, Some(backtracks));
let guide = match target_guide {
Some(tg) => {
Some(match self.guide {
Some(exguide) => {
let exguide = load_image(exguide, resize)?;
ImagePyramid::new(exguide, Some(backtracks))
}
None => {
let mut gm = transform_to_guide_map(image.bottom().clone(), resize, 2.0);
match_histograms(&mut gm, tg.bottom());
ImagePyramid::new(gm, Some(backtracks))
}
})
}
None => None,
};
let method = match self.sample_method {
SampleMethod::All => SamplingMethod::All,
SampleMethod::Ignore => SamplingMethod::Ignore,
SampleMethod::Image(src) => {
let img = load_image(src, resize)?;
SamplingMethod::Image(img)
}
};
Ok(ResolvedExample {
image,
guide,
method,
})
}
}
impl<'a, S> From<&'a S> for Example<'a>
where
S: AsRef<Path> + 'a,
{
fn from(path: &'a S) -> Self {
Example::new(path)
}
}
#[derive(Default)]
pub struct SessionBuilder<'a> {
examples: Vec<Example<'a>>,
target_guide: Option<ImageSource<'a>>,
inpaint_mask: Option<(ImageSource<'a>, usize)>,
params: Parameters,
}
impl<'a> SessionBuilder<'a> {
pub fn new() -> Self {
Self::default()
}
pub fn add_example<E: Into<Example<'a>>>(mut self, example: E) -> Self {
self.examples.push(example.into());
self
}
pub fn add_examples<E: Into<Example<'a>>, I: IntoIterator<Item = E>>(
mut self,
examples: I,
) -> Self {
self.examples.extend(examples.into_iter().map(|e| e.into()));
self
}
pub fn inpaint_example<I: Into<ImageSource<'a>>, E: Into<Example<'a>>>(
mut self,
inpaint_mask: I,
example: E,
) -> Self {
self.inpaint_mask = Some((inpaint_mask.into(), self.examples.len()));
self.examples.push(example.into());
self
}
pub fn load_target_guide<I: Into<ImageSource<'a>>>(mut self, guide: I) -> Self {
self.target_guide = Some(guide.into());
self
}
pub fn resize_input(mut self, w: u32, h: u32) -> Self {
self.params.resize_input = Some((w, h));
self
}
pub fn seed(mut self, value: u64) -> Self {
self.params.seed = value;
self
}
pub fn tiling_mode(mut self, is_tiling: bool) -> Self {
self.params.tiling_mode = is_tiling;
self
}
pub fn nearest_neighbors(mut self, count: u32) -> Self {
self.params.nearest_neighbors = count;
self
}
pub fn random_sample_locations(mut self, count: u64) -> Self {
self.params.random_sample_locations = count;
self
}
pub fn random_init(mut self, count: u64) -> Self {
self.params.random_resolve = Some(count);
self
}
pub fn cauchy_dispersion(mut self, value: f32) -> Self {
self.params.cauchy_dispersion = value;
self
}
pub fn guide_alpha(mut self, value: f32) -> Self {
self.params.guide_alpha = value;
self
}
pub fn backtrack_percent(mut self, value: f32) -> Self {
self.params.backtrack_percent = value;
self
}
pub fn backtrack_stages(mut self, stages: u32) -> Self {
self.params.backtrack_stages = stages;
self
}
pub fn output_size(mut self, w: u32, h: u32) -> Self {
self.params.output_size = (w, h);
self
}
pub fn build(self) -> Result<Session, Error> {
self.check_parameters_validity()?;
self.check_images_validity()?;
let target_guide = match self.target_guide {
Some(tg) => {
let tg_img = load_image(tg, Some(self.params.output_size))?;
let num_guides = self.examples.iter().filter(|ex| ex.guide.is_some()).count();
let tg_img = if num_guides == 0 {
transform_to_guide_map(tg_img, None, 2.0)
} else {
tg_img
};
Some(ImagePyramid::new(
tg_img,
Some(self.params.backtrack_stages as u32),
))
}
None => None,
};
struct InpaintExample {
inpaint_mask: image::RgbaImage,
color_map: image::RgbaImage,
example_index: usize,
}
let inpaint = match self.inpaint_mask {
Some((src, ind)) => {
let inpaint_mask = load_image(src, Some(self.params.output_size))?;
let color_map = load_image(
self.examples[ind].img.clone(),
Some(self.params.output_size),
)?;
Some(InpaintExample {
inpaint_mask,
color_map,
example_index: ind,
})
}
None => None,
};
let example_len = self.examples.len();
let mut examples = Vec::with_capacity(example_len);
let mut guides = if target_guide.is_some() {
Vec::with_capacity(example_len)
} else {
Vec::new()
};
let mut methods = Vec::with_capacity(example_len);
for example in self.examples {
let resolved = example.resolve(
self.params.backtrack_stages,
self.params.resize_input,
&target_guide,
)?;
examples.push(resolved.image);
if let Some(guide) = resolved.guide {
guides.push(guide);
}
methods.push(resolved.method);
}
let generator = match inpaint {
None => multires_stochastic_texture_synthesis::Generator::new(self.params.output_size),
Some(inpaint) => multires_stochastic_texture_synthesis::Generator::new_from_inpaint(
self.params.output_size,
inpaint.inpaint_mask,
inpaint.color_map,
inpaint.example_index,
),
};
let session = Session {
examples,
guides: target_guide.map(|tg| GuidesPyramidStruct {
target_guide: tg,
example_guides: guides,
}),
sampling_methods: methods,
params: self.params,
generator,
};
Ok(session)
}
fn check_parameters_validity(&self) -> Result<(), Error> {
if self.params.cauchy_dispersion < 0.0 || self.params.cauchy_dispersion > 1.0 {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Invalid parameter range: cauchy dispersion. Make sure it is within 0.0-1.0 range",
)));
}
if self.params.backtrack_percent < 0.0 || self.params.backtrack_percent > 1.0 {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Invalid parameter range: backtrack percent. Make sure it is within 0.0-1.0 range",
)));
}
if self.params.guide_alpha < 0.0 || self.params.guide_alpha > 1.0 {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Invalid parameter range: guide alpha. Make sure it is within 0.0-1.0 range",
)));
}
if self.inpaint_mask.is_some() {
if let Some(resize_input) = self.params.resize_input {
if resize_input.0 != self.params.output_size.0
|| resize_input.1 != self.params.output_size.1
{
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Input and output sizes dont match. Make sure resize_input = output_size if using inpaint",
)));
}
} else {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Input and output sizes dont match. Make sure resize_input = output_size if using inpaint",
)));
}
}
Ok(())
}
fn check_images_validity(&self) -> Result<(), Error> {
if self.examples.is_empty() {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Missing input: example image(s)",
)));
}
let num_guides = self.examples.iter().filter(|ex| ex.guide.is_some()).count();
if num_guides != 0 && self.examples.len() != num_guides {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"Mismatch of maps: {} example guide(s) vs {} example(s)",
num_guides,
self.examples.len()
),
)));
}
let input_count = self
.examples
.iter()
.filter(|ex| !ex.sample_method.is_ignore())
.count();
if input_count == 0 {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"At least 1 example image must be sampled to generate an output image",
)));
}
Ok(())
}
}
struct ResolvedExample {
image: ImagePyramid,
guide: Option<ImagePyramid>,
method: SamplingMethod,
}
pub struct Session {
examples: Vec<ImagePyramid>,
guides: Option<GuidesPyramidStruct>,
sampling_methods: Vec<SamplingMethod>,
generator: multires_stochastic_texture_synthesis::Generator,
params: Parameters,
}
impl Session {
pub fn builder<'a>() -> SessionBuilder<'a> {
SessionBuilder::default()
}
pub fn run(mut self, progress: Option<Box<dyn GeneratorProgress>>) -> GeneratedImage {
if let Some(count) = self.params.random_resolve {
let lvl = self.examples[0].pyramid.len();
let imgs: Vec<image::RgbaImage> = self
.examples
.iter()
.map(|a| a.pyramid[lvl - 1].clone())
.collect();
let imgs_ref = imgs.iter().collect::<Vec<_>>();
self.generator
.resolve_random_batch(count as usize, &imgs_ref, self.params.seed);
}
self.generator.main_resolve_loop(
&self.params.to_generator_params(),
&self.examples,
progress,
&self.guides,
&self.sampling_methods,
self.params.tiling_mode,
);
GeneratedImage {
inner: self.generator,
}
}
}
pub struct ProgressStat {
pub current: usize,
pub total: usize,
}
pub struct ProgressUpdate<'a> {
pub image: &'a image::RgbaImage,
pub total: ProgressStat,
pub stage: ProgressStat,
}
pub trait GeneratorProgress {
fn update(&mut self, info: ProgressUpdate<'_>);
}
impl<G> GeneratorProgress for G
where
G: FnMut(ProgressUpdate<'_>) + Send,
{
fn update(&mut self, info: ProgressUpdate<'_>) {
self(info)
}
}