lib.rs

//! `texture-synthesis` is a light API for Multiresolution Stochastic Texture Synthesis,
//! a non-parametric example-based algorithm for image generation.
//!
//! First, you build a `Session` via a `SessionBuilder`, which follows the builder pattern. Calling
//! `build` on the SessionBuilder loads all of the input images and checks for various errors.
//!
//! `Session` has a `run()` method that takes all of the parameters and inputs added in the session
//! builder to generated an image, which is returned as a `GeneratedImage`.
//!
//! You can save, stream, or inspect the image from `GeneratedImage`.
//!
//! ## Features
//!
//! 1. Single example generation
//! 2. Multi example generation
//! 3. Guided synthesis
//! 4. Style transfer
//! 5. Inpainting
//! 6. Tiling textures
//!
//! Please, refer to the examples folder in the [repository](https://github.com/EmbarkStudios/texture-synthesis) for the features usage examples.
//!
//! ## Usage
//! Session follows a "builder pattern" for defining parameters, meaning you chain functions together.
//!
//! ```no_run
//! // Create a new session with default parameters
//! let session = texture_synthesis::Session::builder()
//!     // Set some parameters
//!     .seed(10)
//!     .nearest_neighbors(20)
//!     // Specify example images
//!     .add_example(&"imgs/1.jpg")
//!     // Build the session
//!     .build().expect("failed to build session");
//!
//! // Generate a new image
//! let generated_img = session.run(None);
//!
//! // Save the generated image to disk
//! generated_img.save("my_generated_img.jpg").expect("failed to save generated image");
//! ```
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,
        }
    }
}

/// An image generated by a `Session::run()`
pub struct GeneratedImage {
    inner: multires_stochastic_texture_synthesis::Generator,
}

impl GeneratedImage {
    /// Saves the generated image to the specified path
    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(())
    }

    /// Writes the generated image to the specified stream
    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)?)
    }

    /// Saves debug information such as copied patches ids, map ids (if you have multi example generation)
    /// and a map indicating generated pixels the generator was "uncertain" of.
    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(())
    }

    /// Returns the generated output image
    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
    }
}

/// Method used for sampling an example image.
pub enum SampleMethod<'a> {
    /// All pixels in the example image can be sampled.
    All,
    /// No pixels in the example image will be sampled.
    Ignore,
    /// Pixels are selectively sampled based on an image.
    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()))
    }
}

/// Internal sample method
pub enum SamplingMethod {
    All,
    Ignore,
    Image(image::RgbaImage),
}

impl SamplingMethod {
    fn is_ignore(&self) -> bool {
        match self {
            Self::Ignore => true,
            _ => false,
        }
    }
}

/// A builder for an `Example`
pub struct ExampleBuilder<'a> {
    img: ImageSource<'a>,
    guide: Option<ImageSource<'a>>,
    sample_method: SampleMethod<'a>,
}

impl<'a> ExampleBuilder<'a> {
    /// Creates a new example builder from the specified image source
    pub fn new<I: Into<ImageSource<'a>>>(img: I) -> Self {
        Self {
            img: img.into(),
            guide: None,
            sample_method: SampleMethod::All,
        }
    }

    /// Use a guide map that describe a 'FROM' transformation.
    ///
    /// Note: If any one example has a guide, then they **all** must have
    /// a guide, otherwise a session will not be created.
    pub fn with_guide<G: Into<ImageSource<'a>>>(mut self, guide: G) -> Self {
        self.guide = Some(guide.into());
        self
    }

    /// Specify how the example image is sampled during texture generation.
    ///
    /// By default, all pixels in the example can be sampled.
    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,
        }
    }
}

/// An example to be used in texture generation
pub struct Example<'a> {
    img: ImageSource<'a>,
    guide: Option<ImageSource<'a>>,
    sample_method: SampleMethod<'a>,
}

impl<'a> Example<'a> {
    /// Creates a new example builder from the specified image source
    pub fn builder<I: Into<ImageSource<'a>>>(img: I) -> ExampleBuilder<'a> {
        ExampleBuilder::new(img)
    }

    /// Creates a new example input from the specified image source
    pub fn new<I: Into<ImageSource<'a>>>(img: I) -> Self {
        Self {
            img: img.into(),
            guide: None,
            sample_method: SampleMethod::All,
        }
    }

    /// Use a guide map that describe a 'FROM' transformation.
    ///
    /// Note: If any one example has a guide, then they **all** must have
    /// a guide, otherwise a session will not be created.
    pub fn with_guide<G: Into<ImageSource<'a>>>(&mut self, guide: G) -> &mut Self {
        self.guide = Some(guide.into());
        self
    }

    /// Specify how the example image is sampled during texture generation.
    ///
    /// By default, all pixels in the example can be sampled.
    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 => {
                        // if we do not have an example guide, create it as a b/w maps of the example
                        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)
    }
}

/// Builds a session by setting parameters and adding input images, calling
/// `build` will check all of the provided inputs to verify that texture
/// synthesis will provide valid output
#[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> {
    /// Creates a new `SessionBuilder`, can also be created via `Session::builder()`
    pub fn new() -> Self {
        Self::default()
    }

    /// Adds an `Example` from which a generator will synthesize a new image.
    ///
    /// See [examples/01_single_example_synthesis](https://github.com/EmbarkStudios/texture-synthesis/tree/master/lib/examples/01_single_example_synthesis.rs)
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let tex_synth = texture_synthesis::Session::builder()
    ///     .add_example(&"imgs/1.jpg")
    ///     .build().expect("failed to build session");
    /// ```
    pub fn add_example<E: Into<Example<'a>>>(mut self, example: E) -> Self {
        self.examples.push(example.into());
        self
    }

    /// Adds Examples from which a generator will synthesize a new image.
    ///
    /// See [examples/02_multi_example_synthesis](https://github.com/EmbarkStudios/texture-synthesis/tree/master/lib/examples/02_multi_example_synthesis.rs)
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let tex_synth = texture_synthesis::Session::builder()
    ///     .add_examples(&[&"imgs/1.jpg", &"imgs/2.jpg"])
    ///     .build().expect("failed to build session");
    /// ```
    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
    }

    /// Inpaints an example.
    ///
    /// To prevent sampling from the example, you can specify `SamplingMethod::Ignore` with `Example::set_sample_method`.
    ///
    /// See [examples/05_inpaint](https://github.com/EmbarkStudios/texture-synthesis/tree/master/lib/examples/05_inpaint.rs)
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let tex_synth = texture_synthesis::Session::builder()
    ///     .add_examples(&[&"imgs/1.jpg", &"imgs/3.jpg"])
    ///     .inpaint_example(
    ///         &"masks/inpaint.jpg",
    ///         // This will prevent sampling from the imgs/2.jpg, note that
    ///         // we *MUST* provide at least one example to source from!
    ///         texture_synthesis::Example::builder(&"imgs/2.jpg")
    ///             .set_sample_method(texture_synthesis::SampleMethod::Ignore)
    ///     )
    ///     .build().expect("failed to build session");
    /// ```
    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
    }

    /// Loads a target guide map.
    ///
    /// If no `Example` guide maps are provided, this will produce a style transfer effect,
    /// where the Examples are styles and the target guide is content.
    ///
    /// See [examples/03_guided_synthesis](https://github.com/EmbarkStudios/texture-synthesis/tree/master/lib/examples/03_guided_synthesis.rs),
    /// or [examples/04_style_transfer](https://github.com/EmbarkStudios/texture-synthesis/tree/master/lib/examples/04_style_transfer.rs),
    pub fn load_target_guide<I: Into<ImageSource<'a>>>(mut self, guide: I) -> Self {
        self.target_guide = Some(guide.into());
        self
    }

    /// Overwrite incoming images sizes
    pub fn resize_input(mut self, w: u32, h: u32) -> Self {
        self.params.resize_input = Some((w, h));
        self
    }

    /// Changes pseudo-deterministic seed.
    /// Global structures will stay same, if same seed is provided, but smaller details may change due to undeterministic nature of multithreading.
    pub fn seed(mut self, value: u64) -> Self {
        self.params.seed = value;
        self
    }

    /// Makes the generator output tiling image.
    /// Default: false.
    pub fn tiling_mode(mut self, is_tiling: bool) -> Self {
        self.params.tiling_mode = is_tiling;
        self
    }

    /// How many neighboring pixels each pixel is aware of during the generation (bigger number -> more global structures are captured).
    /// Default: 50
    pub fn nearest_neighbors(mut self, count: u32) -> Self {
        self.params.nearest_neighbors = count;
        self
    }

    /// How many random locations will be considered during a pixel resolution apart from its immediate neighbors.
    /// If unsure, keep same as nearest neighbors.
    /// Default: 50
    pub fn random_sample_locations(mut self, count: u64) -> Self {
        self.params.random_sample_locations = count;
        self
    }

    /// Make first X pixels to be randomly resolved and prevent them from being overwritten.
    /// Can be an enforcing factor of remixing multiple images together.
    pub fn random_init(mut self, count: u64) -> Self {
        self.params.random_resolve = Some(count);
        self
    }

    /// The distribution dispersion used for picking best candidate (controls the distribution 'tail flatness').
    /// Values close to 0.0 will produce 'harsh' borders between generated 'chunks'. Values closer to 1.0 will produce a smoother gradient on those borders.
    /// For futher reading, check out P.Harrison's "Image Texture Tools".
    /// Default: 1.0
    pub fn cauchy_dispersion(mut self, value: f32) -> Self {
        self.params.cauchy_dispersion = value;
        self
    }

    /// Controls the trade-off between guide and example map.
    /// If doing style transfer, set to about 0.8-0.6 to allow for more global structures of the style.
    /// If you'd like the guide maps to be considered through all generation stages, set to 1.0 (which would prevent guide maps weight "decay" during the score calculation).
    /// Default: 0.8
    pub fn guide_alpha(mut self, value: f32) -> Self {
        self.params.guide_alpha = value;
        self
    }

    /// The percentage of pixels to be backtracked during each p_stage. Range (0,1).
    /// Default: 0.5
    pub fn backtrack_percent(mut self, value: f32) -> Self {
        self.params.backtrack_percent = value;
        self
    }

    /// Controls the number of backtracking stages. Backtracking prevents 'garbage' generation.
    /// Right now, the depth of image pyramid for multiresolution synthesis
    /// depends on this parameter as well.
    /// Default: 5
    pub fn backtrack_stages(mut self, stages: u32) -> Self {
        self.params.backtrack_stages = stages;
        self
    }

    /// Specify size of the generated image.
    /// Default: 500x500
    pub fn output_size(mut self, w: u32, h: u32) -> Self {
        self.params.output_size = (w, h);
        self
    }

    /// Creates a `Session`, or returns an error if invalid parameters or input
    /// images were specified.
    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);
        }

        // Initialize generator based on availability of an inpaint_mask.
        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()
                ),
            )));
        }

        // We must have at least one example image to source pixels from
        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,
}

/// Texture synthesis session.
///
/// Calling `run()` will generate a new image and return it, consuming the session in the
/// process. You can provide a `GeneratorProgress` implementation to periodically get
/// updates with the currently generated image and the number of pixels that
/// have been resolved both in the current stage and globally
///
/// # Example
/// ```no_run
/// let tex_synth = texture_synthesis::Session::builder()
///     .seed(10)
///     .tiling_mode(true)
///     .add_example(&"imgs/1.jpg")
///     .build().expect("failed to build session");
///
/// let generated_img = tex_synth.run(None);
/// generated_img.save("my_generated_img.jpg").expect("failed to save image");
/// ```
pub struct Session {
    examples: Vec<ImagePyramid>,
    guides: Option<GuidesPyramidStruct>,
    sampling_methods: Vec<SamplingMethod>,
    generator: multires_stochastic_texture_synthesis::Generator,
    params: Parameters,
}

impl Session {
    /// Creates a new session with default parameters.
    pub fn builder<'a>() -> SessionBuilder<'a> {
        SessionBuilder::default()
    }

    /// Runs the generator and outputs a generated image.
    /// Now, only runs Multiresolution Stochastic Texture Synthesis.
    /// Might be interesting to include more algorithms in the future.
    pub fn run(mut self, progress: Option<Box<dyn GeneratorProgress>>) -> GeneratedImage {
        // random resolve
        // TODO: Instead of consuming the generator, we could instead make the
        // seed and random_resolve parameters, so that you could rerun the
        // generator with the same inputs
        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()) //take the blurriest image
                .collect();
            let imgs_ref = imgs.iter().collect::<Vec<_>>();
            self.generator
                .resolve_random_batch(count as usize, &imgs_ref, self.params.seed);
        }

        // run generator
        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,
        }
    }
}

/// Helper struct for passing progress information to external callers
pub struct ProgressStat {
    /// The current amount of work that has been done
    pub current: usize,
    /// The total amount of work to do
    pub total: usize,
}

/// The current state of the image generator
pub struct ProgressUpdate<'a> {
    /// The currenty resolved image
    pub image: &'a image::RgbaImage,
    /// The total progress for the final image
    pub total: ProgressStat,
    /// The progress for the current stage
    pub stage: ProgressStat,
}

/// Allows the generator to update external callers with the current
/// progress of the image synthesis
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)
    }
}