#![warn(clippy::all)]
#![warn(rust_2018_idioms)]
mod errors;
mod img_pyramid;
use img_pyramid::*;
mod utils;
use utils::*;
mod multires_stochastic_texture_synthesis;
use multires_stochastic_texture_synthesis::*;
use std::path::Path;
mod unsync;
pub use image;
pub use utils::{load_dynamic_image, ChannelMask, ImageSource};
pub use errors::Error;
#[derive(Copy, Clone)]
pub struct Dims {
pub width: u32,
pub height: u32,
}
impl Dims {
pub fn square(size: u32) -> Self {
Self {
width: size,
height: size,
}
}
pub fn new(width: u32, height: u32) -> Self {
Self { width, height }
}
}
pub struct CoordinateTransform {
buffer: Vec<u32>,
dims: Dims,
max_map_id: u32, }
impl<'a> CoordinateTransform {
pub fn apply<E: Into<ImageSource<'a>>, I: IntoIterator<Item = E>>(
&self,
source: I,
) -> Result<image::RgbaImage, Error> {
let ref_maps: Vec<image::RgbaImage> = source
.into_iter()
.map(|t| load_image(t.into(), None))
.collect::<Result<Vec<_>, Error>>()?;
if ref_maps.len() as u32 != self.max_map_id {
return Err(Error::MapsCountMismatch(
ref_maps.len() as u32,
self.max_map_id,
));
}
let mut img = image::RgbaImage::new(self.dims.width, self.dims.height);
for (i, pix) in img.pixels_mut().enumerate() {
let x = self.buffer[i * 3];
let y = self.buffer[i * 3 + 1];
let map = self.buffer[i * 3 + 2];
*pix = *ref_maps[map as usize].get_pixel(x, y);
}
Ok(img)
}
}
struct Parameters {
tiling_mode: bool,
nearest_neighbors: u32,
random_sample_locations: u64,
cauchy_dispersion: f32,
backtrack_percent: f32,
backtrack_stages: u32,
resize_input: Option<Dims>,
output_size: Dims,
guide_alpha: f32,
random_resolve: Option<u64>,
max_thread_count: Option<usize>,
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: Dims::square(500),
guide_alpha: 0.8,
random_resolve: None,
max_thread_count: 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,
max_thread_count: self.max_thread_count.unwrap_or_else(num_cpus::get),
tiling_mode: self.tiling_mode,
}
}
}
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.as_ref().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 get_coordinate_transform(&self) -> CoordinateTransform {
self.inner.get_coord_transform()
}
pub fn into_image(self) -> image::DynamicImage {
image::DynamicImage::ImageRgba8(self.inner.color_map.into_inner())
}
}
impl AsRef<image::RgbaImage> for GeneratedImage {
fn as_ref(&self) -> &image::RgbaImage {
self.inner.color_map.as_ref()
}
}
pub enum GenericSampleMethod<Img> {
All,
Ignore,
Image(Img),
}
pub type SampleMethod<'a> = GenericSampleMethod<ImageSource<'a>>;
pub type SamplingMethod = GenericSampleMethod<image::RgbaImage>;
impl<Img> GenericSampleMethod<Img> {
#[inline]
fn is_ignore(&self) -> bool {
match self {
Self::Ignore => true,
_ => false,
}
}
}
impl<'a, IS> From<IS> for SampleMethod<'a>
where
IS: Into<ImageSource<'a>>,
{
fn from(is: IS) -> Self {
SampleMethod::Image(is.into())
}
}
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 image_source(&self) -> &ImageSource<'a> {
&self.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<Dims>,
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, IS> From<IS> for Example<'a>
where
IS: Into<ImageSource<'a>>,
{
fn from(is: IS) -> Self {
Example::new(is)
}
}
enum MaskOrImg<'a> {
Mask(utils::ChannelMask),
ImageSource(ImageSource<'a>),
}
struct InpaintMask<'a> {
src: MaskOrImg<'a>,
example_index: usize,
dims: Dims,
}
#[derive(Default)]
pub struct SessionBuilder<'a> {
examples: Vec<Example<'a>>,
target_guide: Option<ImageSource<'a>>,
inpaint_mask: Option<InpaintMask<'a>>,
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,
size: Dims,
) -> Self {
self.inpaint_mask = Some(InpaintMask {
src: MaskOrImg::ImageSource(inpaint_mask.into()),
example_index: self.examples.len(),
dims: size,
});
self.examples.push(example.into());
self
}
pub fn inpaint_example_channel<E: Into<Example<'a>>>(
mut self,
mask: utils::ChannelMask,
example: E,
size: Dims,
) -> Self {
self.inpaint_mask = Some(InpaintMask {
src: MaskOrImg::Mask(mask),
example_index: self.examples.len(),
dims: size,
});
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, dims: Dims) -> Self {
self.params.resize_input = Some(dims);
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, dims: Dims) -> Self {
self.params.output_size = dims;
self
}
pub fn max_thread_count(mut self, count: usize) -> Self {
self.params.max_thread_count = Some(count);
self
}
pub fn build(mut self) -> Result<Session, Error> {
self.check_parameters_validity()?;
self.check_images_validity()?;
struct InpaintExample {
inpaint_mask: image::RgbaImage,
color_map: image::RgbaImage,
example_index: usize,
}
let (inpaint, out_size, in_size) = match self.inpaint_mask {
Some(inpaint_mask) => {
let dims = inpaint_mask.dims;
let inpaint_img = match inpaint_mask.src {
MaskOrImg::ImageSource(img) => load_image(img, Some(dims))?,
MaskOrImg::Mask(mask) => {
let example_img = &mut self.examples[inpaint_mask.example_index].img;
let dynamic_img = utils::load_dynamic_image(example_img.clone())?;
let inpaint_src = ImageSource::Image(dynamic_img.clone());
*example_img = ImageSource::Image(dynamic_img);
let inpaint_mask = load_image(inpaint_src, Some(dims))?;
utils::apply_mask(inpaint_mask, mask)
}
};
let color_map = load_image(
self.examples[inpaint_mask.example_index].img.clone(),
Some(dims),
)?;
(
Some(InpaintExample {
inpaint_mask: inpaint_img,
color_map,
example_index: inpaint_mask.example_index,
}),
dims,
Some(dims),
)
}
None => (None, self.params.output_size, self.params.resize_input),
};
let target_guide = match self.target_guide {
Some(tg) => {
let tg_img = load_image(tg, Some(out_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,
};
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, in_size, &target_guide)?;
examples.push(resolved.image);
if let Some(guide) = resolved.guide {
guides.push(guide);
}
methods.push(resolved.method);
}
let generator = match inpaint {
None => Generator::new(out_size),
Some(inpaint) => Generator::new_from_inpaint(
out_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(Error::InvalidRange(errors::InvalidRange {
min: 0.0,
max: 1.0,
value: self.params.cauchy_dispersion,
name: "cauchy-dispersion",
}));
}
if self.params.backtrack_percent < 0.0 || self.params.backtrack_percent > 1.0 {
return Err(Error::InvalidRange(errors::InvalidRange {
min: 0.0,
max: 1.0,
value: self.params.backtrack_percent,
name: "backtrack-percent",
}));
}
if self.params.guide_alpha < 0.0 || self.params.guide_alpha > 1.0 {
return Err(Error::InvalidRange(errors::InvalidRange {
min: 0.0,
max: 1.0,
value: self.params.guide_alpha,
name: "guide-alpha",
}));
}
if let Some(max_count) = self.params.max_thread_count {
if max_count == 0 {
return Err(Error::InvalidRange(errors::InvalidRange {
min: 1.0,
max: 1024.0,
value: max_count as f32,
name: "max-thread-count",
}));
}
}
if self.params.random_sample_locations == 0 {
return Err(Error::InvalidRange(errors::InvalidRange {
min: 1.0,
max: 1024.0,
value: self.params.random_sample_locations as f32,
name: "m-rand",
}));
}
Ok(())
}
fn check_images_validity(&self) -> Result<(), Error> {
let input_count = self
.examples
.iter()
.filter(|ex| !ex.sample_method.is_ignore())
.count();
if input_count == 0 {
return Err(Error::NoExamples);
}
let num_guides = self.examples.iter().filter(|ex| ex.guide.is_some()).count();
if num_guides != 0 && self.examples.len() != num_guides {
return Err(Error::ExampleGuideMismatch(
self.examples.len() as u32,
num_guides as u32,
));
}
Ok(())
}
}
struct ResolvedExample {
image: ImagePyramid,
guide: Option<ImagePyramid>,
method: SamplingMethod,
}
pub struct Session {
examples: Vec<ImagePyramid>,
guides: Option<GuidesPyramidStruct>,
sampling_methods: Vec<SamplingMethod>,
generator: 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<_> = self
.examples
.iter()
.map(|a| ImageBuffer::from(&a.pyramid[lvl - 1])) .collect();
self.generator
.resolve_random_batch(count as usize, &imgs, self.params.seed);
}
self.generator.main_resolve_loop(
&self.params.to_generator_params(),
&self.examples,
progress,
&self.guides,
&self.sampling_methods,
);
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)
}
}