Add --inpaint-alpha to CLI and API changes to derive mask from channels

cli/src/main.rs

@@ -5,7 +5,8 @@ use structopt::StructOpt;
 
 use std::path::PathBuf;
 use texture_synthesis::{
-    image::ImageOutputFormat as ImgFmt, Dims, Error, Example, ImageSource, SampleMethod, Session,
+    image::ImageOutputFormat as ImgFmt, Dims, Error, Example, ImageSource, Mask, SampleMethod,
+    Session,
 };
 
 fn parse_size(input: &str) -> Result<Dims, std::num::ParseIntError> {
@@ -138,6 +139,9 @@ struct Opt {
     /// Path to an inpaint map image, where black pixels are resolved, and white pixels are kept
     #[structopt(long, parse(from_os_str))]
     inpaint: Option<PathBuf>,
+    /// Flag to extract inpaint from the example's alpha channel
+    #[structopt(long)]
+    inpaint_channel: Option<Mask>,
     /// Size of the generated image, in `width x height`, or a single number for both dimensions
     #[structopt(
         long,
@@ -230,7 +234,7 @@ fn real_main() -> Result<(), Error> {
             match mask.as_str() {
                 "ALL" => example.set_sample_method(SampleMethod::All),
                 "IGNORE" => example.set_sample_method(SampleMethod::Ignore),
-                path => example.set_sample_method(SampleMethod::Image(ImageSource::Path(
+                path => example.set_sample_method(SampleMethod::Image(ImageSource::from_path(
                     &std::path::Path::new(path),
                 ))),
             };
@@ -240,7 +244,20 @@ fn real_main() -> Result<(), Error> {
     let mut sb = Session::builder();
 
     // TODO: Make inpaint work with multiple examples
-    if let Some(ref inpaint) = args.inpaint {
+    if let (Some(_), Some(_)) = (args.inpaint_channel, &args.inpaint) {
+        return Err(Error::UnsupportedArgPair(
+            "channel".to_owned(),
+            "inpaint".to_owned(),
+        ));
+    } else if let Some(channel) = args.inpaint_channel {
+        let mut inpaint_example = examples.remove(0);
+        let inpaint = inpaint_example.image_source().clone().mask(channel);
+        if args.sample_masks.is_empty() {
+            inpaint_example.set_sample_method(SampleMethod::Image(inpaint.clone()));
+        }
+
+        sb = sb.inpaint_example(inpaint, inpaint_example, args.out_size);
+    } else if let Some(ref inpaint) = args.inpaint {
         let mut inpaint_example = examples.remove(0);
 
         // If the user hasn't explicitly specified sample masks, assume they

lib/src/errors.rs

@@ -56,6 +56,10 @@ pub enum Error {
     NoExamples,
     ///
     MapsCountMismatch(u32, u32),
+    /// An invalid string has that can't be parsed into a `Mask`
+    ParseMask(String),
+    /// The user specified an unsupported pair of arguments
+    UnsupportedArgPair(String, String),
 }
 
 impl fmt::Display for Error {
@@ -92,6 +96,16 @@ impl fmt::Display for Error {
                 "{} map(s) were provided, but {} is/are required",
                 input, required
             ),
+            Self::ParseMask(mask_string) => write!(
+                f,
+                "couldn't parse mask '{}', not one of: 'r', 'g', 'b', 'a'",
+                mask_string
+            ),
+            Self::UnsupportedArgPair(arg_a, arg_b) => write!(
+                f,
+                "the arguments '{}' and '{}' can't be used together, pick one of them",
+                arg_a, arg_b,
+            ),
         }
     }
 }

lib/src/lib.rs

@@ -54,7 +54,7 @@ use std::path::Path;
 mod unsync;
 
 pub use image;
-pub use utils::ImageSource;
+pub use utils::{DataSource, ImageSource, Mask};
 
 pub use errors::Error;
 
@@ -325,7 +325,9 @@ 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
+    }
     /// Creates a new example input from the specified image source
     pub fn new<I: Into<ImageSource<'a>>>(img: I) -> Self {
         Self {

lib/src/utils.rs

@@ -1,9 +1,31 @@
 use crate::{Dims, Error};
 use std::path::Path;
+use std::str::FromStr;
 
 /// Helper type used to pass image data to the Session
 #[derive(Clone)]
-pub enum ImageSource<'a> {
+pub struct ImageSource<'a> {
+    data_source: DataSource<'a>,
+    mask: Option<Mask>,
+}
+
+impl<'a> ImageSource<'a> {
+    pub fn from_path(path: &'a Path) -> ImageSource<'a> {
+        ImageSource {
+            data_source: DataSource::Path(path),
+            mask: None,
+        }
+    }
+
+    pub fn mask(mut self, mask: Mask) -> Self {
+        self.mask = Some(mask);
+        self
+    }
+}
+
+/// Helper type used to define the source of `ImageSource`'s data
+#[derive(Clone)]
+pub enum DataSource<'a> {
     /// A raw buffer of image data, see `image::load_from_memory` for details
     /// on what is supported
     Memory(&'a [u8]),
@@ -14,13 +36,25 @@ pub enum ImageSource<'a> {
     Image(image::DynamicImage),
 }
 
-impl<'a> From<image::DynamicImage> for ImageSource<'a> {
+impl<'a> From<image::DynamicImage> for DataSource<'a> {
     fn from(img: image::DynamicImage) -> Self {
-        ImageSource::Image(img)
+        DataSource::Image(img)
     }
 }
 
 impl<'a, S> From<&'a S> for ImageSource<'a>
+where
+    S: AsRef<Path> + 'a,
+{
+    fn from(path: &'a S) -> Self {
+        Self {
+            data_source: DataSource::Path(path.as_ref()),
+            mask: None,
+        }
+    }
+}
+
+impl<'a, S> From<&'a S> for DataSource<'a>
 where
     S: AsRef<Path> + 'a,
 {
@@ -29,17 +63,40 @@ where
     }
 }
 
+/// Helper type used to mask `ImageSource`'s channels
+#[derive(Clone, Copy)]
+pub enum Mask {
+    R,
+    G,
+    B,
+    A,
+}
+
+impl FromStr for Mask {
+    type Err = Error;
+
+    fn from_str(mask: &str) -> Result<Self, Self::Err> {
+        match &mask.to_lowercase()[..] {
+            "r" => Ok(Mask::R),
+            "g" => Ok(Mask::G),
+            "b" => Ok(Mask::B),
+            "a" => Ok(Mask::A),
+            mask => Err(Error::ParseMask(mask.to_string())),
+        }
+    }
+}
+
 pub(crate) fn load_image(
     src: ImageSource<'_>,
     resize: Option<Dims>,
 ) -> Result<image::RgbaImage, Error> {
-    let img = match src {
-        ImageSource::Memory(data) => image::load_from_memory(data),
-        ImageSource::Path(path) => image::open(path),
-        ImageSource::Image(img) => Ok(img),
+    let img = match src.data_source {
+        DataSource::Memory(data) => image::load_from_memory(data),
+        DataSource::Path(path) => image::open(path),
+        DataSource::Image(img) => Ok(img),
     }?;
 
-    Ok(match resize {
+    let img = match resize {
         None => img.to_rgba(),
         Some(ref size) => {
             use image::GenericImageView;
@@ -55,7 +112,36 @@ pub(crate) fn load_image(
                 img.to_rgba()
             }
         }
-    })
+    };
+
+    let img = if let Some(mask) = src.mask {
+        apply_mask(&img, mask)
+    } else {
+        img
+    };
+
+    Ok(img)
+}
+
+pub(crate) fn apply_mask(original_image: &image::RgbaImage, mask: Mask) -> image::RgbaImage {
+    let mut image = original_image.clone();
+
+    let channel = match mask {
+        Mask::R => 0,
+        Mask::G => 1,
+        Mask::B => 2,
+        Mask::A => 3,
+    };
+
+    for pixel_iter in image.enumerate_pixels_mut() {
+        let pixel = pixel_iter.2;
+        pixel[0] = pixel[channel];
+        pixel[1] = pixel[channel];
+        pixel[2] = pixel[channel];
+        pixel[3] = 255;
+    }
+
+    image
 }
 
 pub(crate) fn transform_to_guide_map(