#![warn(
clippy::all,
clippy::await_holding_lock,
clippy::dbg_macro,
clippy::debug_assert_with_mut_call,
clippy::doc_markdown,
clippy::empty_enum,
clippy::enum_glob_use,
clippy::explicit_into_iter_loop,
clippy::filter_map_next,
clippy::fn_params_excessive_bools,
clippy::if_let_mutex,
clippy::imprecise_flops,
clippy::inefficient_to_string,
clippy::large_types_passed_by_value,
clippy::let_unit_value,
clippy::linkedlist,
clippy::lossy_float_literal,
clippy::macro_use_imports,
clippy::map_err_ignore,
clippy::map_flatten,
clippy::map_unwrap_or,
clippy::match_on_vec_items,
clippy::match_same_arms,
clippy::match_wildcard_for_single_variants,
clippy::mem_forget,
clippy::mismatched_target_os,
clippy::needless_borrow,
clippy::needless_continue,
clippy::option_option,
clippy::pub_enum_variant_names,
clippy::ref_option_ref,
clippy::rest_pat_in_fully_bound_structs,
clippy::string_add_assign,
clippy::string_add,
clippy::string_to_string,
clippy::suboptimal_flops,
clippy::todo,
clippy::unimplemented,
clippy::unnested_or_patterns,
clippy::unused_self,
clippy::verbose_file_reads,
future_incompatible,
nonstandard_style,
rust_2018_idioms
)]
mod errors;
mod img_pyramid;
use img_pyramid::*;
mod utils;
use utils::*;
mod ms;
use ms::*;
pub mod session;
mod unsync;
pub use image;
use std::path::Path;
pub use errors::Error;
pub use session::{Session, SessionBuilder};
pub use utils::{load_dynamic_image, ChannelMask, ImageSource};
#[derive(Copy, Clone)]
#[cfg_attr(test, derive(Debug, PartialEq))]
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>,
pub output_size: Dims,
original_maps: Vec<Dims>,
}
const TRANSFORM_MAGIC: u32 = 0x1234_0001;
impl<'a> CoordinateTransform {
pub fn apply<E, I>(&self, source: I) -> Result<image::RgbaImage, Error>
where
I: IntoIterator<Item = E>,
E: Into<ImageSource<'a>>,
{
let ref_maps: Vec<image::RgbaImage> = source
.into_iter()
.zip(self.original_maps.iter())
.map(|(is, dims)| load_image(is.into(), Some(*dims)))
.collect::<Result<Vec<_>, Error>>()?;
if ref_maps.len() != self.original_maps.len() {
return Err(Error::MapsCountMismatch(
ref_maps.len() as u32,
self.original_maps.len() as u32,
));
}
let mut img = image::RgbaImage::new(self.output_size.width, self.output_size.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)
}
pub fn write<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<usize> {
use std::mem;
let mut written = 0;
if self.buffer.len()
!= self.output_size.width as usize * self.output_size.height as usize * 3
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"buffer length doesn't match dimensions",
));
}
let header = [
TRANSFORM_MAGIC,
self.output_size.width,
self.output_size.height,
self.original_maps.len() as u32,
];
fn cast(ina: &[u32]) -> &[u8] {
unsafe {
let p = ina.as_ptr();
let len = ina.len();
std::slice::from_raw_parts(p as *const u8, len * mem::size_of::<u32>())
}
}
w.write_all(cast(&header))?;
written += mem::size_of_val(&header);
for om in &self.original_maps {
let dims = [om.width, om.height];
w.write_all(cast(&dims))?;
written += mem::size_of_val(&dims);
}
w.write_all(cast(&self.buffer))?;
written += 4 * self.buffer.len();
Ok(written)
}
pub fn read<R: std::io::Read>(r: &mut R) -> std::io::Result<Self> {
use std::{
io::{Error, ErrorKind, Read},
mem,
};
fn do_read<R: Read>(r: &mut R, buf: &mut [u32]) -> std::io::Result<()> {
unsafe {
let p = buf.as_mut_ptr();
let len = buf.len();
let mut slice =
std::slice::from_raw_parts_mut(p as *mut u8, len * mem::size_of::<u32>());
r.read(&mut slice).map(|_| ())
}
}
let mut magic = [0u32];
do_read(r, &mut magic)?;
if magic[0] >> 16 != 0x1234 {
return Err(Error::new(ErrorKind::InvalidData, "invalid magic"));
}
let (output_size, original_maps) = match magic[0] & 0x0000_ffff {
0x1 => {
let mut header = [0u32; 3];
do_read(r, &mut header)?;
let mut omaps = Vec::with_capacity(header[2] as usize);
for _ in 0..header[2] {
let mut dims = [0u32; 2];
do_read(r, &mut dims)?;
omaps.push(Dims {
width: dims[0],
height: dims[1],
});
}
(
Dims {
width: header[0],
height: header[1],
},
omaps,
)
}
_ => return Err(Error::new(ErrorKind::InvalidData, "invalid version")),
};
let buffer = unsafe {
let len = output_size.width as usize * output_size.height as usize * 3;
let mut buffer = Vec::with_capacity(len);
buffer.set_len(len);
do_read(r, &mut buffer)?;
buffer
};
Ok(Self {
output_size,
original_maps,
buffer,
})
}
}
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: ms::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 {
matches!(self, Self::Ignore)
}
}
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
}
}
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> From<ExampleBuilder<'a>> for Example<'a> {
fn from(eb: ExampleBuilder<'a>) -> Self {
Self {
img: eb.img,
guide: eb.guide,
sample_method: eb.sample_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,
}
struct ResolvedExample {
image: ImagePyramid,
guide: Option<ImagePyramid>,
method: SamplingMethod,
}
#[cfg(test)]
mod test {
#[test]
fn coord_tx_serde() {
use super::CoordinateTransform as CT;
let fake_buffer = vec![1, 2, 3, 4, 5, 6];
let input = CT {
buffer: fake_buffer.clone(),
output_size: super::Dims {
width: 2,
height: 1,
},
original_maps: vec![
super::Dims {
width: 9001,
height: 9002,
},
super::Dims {
width: 20,
height: 5,
},
],
};
let mut buffer = Vec::new();
input.write(&mut buffer).unwrap();
let mut cursor = std::io::Cursor::new(&buffer);
let deserialized = CT::read(&mut cursor).unwrap();
assert_eq!(deserialized.buffer, fake_buffer);
assert_eq!(deserialized.output_size.width, 2);
assert_eq!(deserialized.output_size.height, 1);
assert_eq!(
super::Dims {
width: 9001,
height: 9002,
},
deserialized.original_maps[0]
);
assert_eq!(
super::Dims {
width: 20,
height: 5,
},
deserialized.original_maps[1]
);
}
}