Add alpha mask utilities to ImageBuffer

src/error.rs

@@ -139,6 +139,8 @@ pub enum ParameterErrorKind {
         /// The cicp that was found.
         found: Cicp,
     },
+    /// The operation is only applicable to pixels with an alpha channel.
+    NoAlphaChannel,
 }
 
 /// An error was encountered while decoding an image.
@@ -450,6 +452,12 @@ impl fmt::Display for ParameterError {
                     "The color space {found:?} does not match the expected {expected:?}",
                 )
             }
+            ParameterErrorKind::NoAlphaChannel => {
+                write!(
+                    fmt,
+                    "The operation requires an alpha channel but the pixel type does not have one",
+                )
+            }
         }?;
 
         if let Some(underlying) = &self.underlying {

src/images/buffer.rs

@@ -888,6 +888,25 @@ where
             color_hint: None, // TODO: the pixel type might contain P::COLOR_TYPE if it satisfies PixelWithColorType
         }
     }
+
+    /// Extract the alpha channel as a Luma image.
+    ///
+    /// If the pixel does not have an alpha channel, the value is filled with a fully opaque mask
+    /// using the maximum value of the corresponding subpixel type.
+    pub fn to_alpha_mask(&self) -> ImageBuffer<Luma<P::Subpixel>, Vec<P::Subpixel>> {
+        let pixels = self.inner_pixels().chunks_exact(P::CHANNEL_COUNT.into());
+        let mut mask = vec![<P::Subpixel as crate::Primitive>::DEFAULT_MAX_VALUE; pixels.len()];
+
+        if P::HAS_ALPHA {
+            for (p, alpha) in pixels.zip(mask.iter_mut()) {
+                // If the pixel has an alpha channel, use it.
+                *alpha = *p.last().unwrap();
+            }
+        }
+
+        ImageBuffer::from_vec(self.width, self.height, mask)
+            .expect("used the right pixel and channel count")
+    }
 }
 
 impl<P, Container> ImageBuffer<P, Container>
@@ -989,6 +1008,42 @@ where
     pub fn put_pixel(&mut self, x: u32, y: u32, pixel: P) {
         *self.get_pixel_mut(x, y) = pixel;
     }
+
+    /// Fill the alpha channel of this image from a Luma mask.
+    ///
+    /// Returns an [`ImageError::Parameter`] if the mask dimensions do not match the image
+    /// dimensions. Otherwise, if the pixel type does not have an alpha channel this is a no-op.
+    pub fn apply_alpha_channel<RhsContainer>(
+        &mut self,
+        mask: &ImageBuffer<Luma<P::Subpixel>, RhsContainer>,
+    ) -> ImageResult<()>
+    where
+        RhsContainer: Deref<Target = [P::Subpixel]>,
+    {
+        if (self.width, self.height) != (mask.width(), mask.height()) {
+            return Err(ImageError::Parameter(ParameterError::from_kind(
+                ParameterErrorKind::DimensionMismatch,
+            )));
+        }
+
+        if !P::HAS_ALPHA {
+            return Err(ImageError::Parameter(ParameterError::from_kind(
+                ParameterErrorKind::NoAlphaChannel,
+            )));
+        }
+
+        let pixels = self
+            .inner_pixels_mut()
+            .chunks_exact_mut(P::CHANNEL_COUNT.into());
+
+        let mask = mask.inner_pixels();
+        for (p, alpha) in pixels.zip(mask.iter()) {
+            // If the pixel has an alpha channel, use it.
+            *p.last_mut().unwrap() = *alpha;
+        }
+
+        Ok(())
+    }
 }
 
 impl<P: Pixel, Container> ImageBuffer<P, Container> {
@@ -2114,6 +2169,65 @@ mod test {
         let result = target.copy_from_color_space(&source, options);
         assert!(matches!(result, Err(crate::ImageError::Parameter(_))));
     }
+
+    #[test]
+    fn alpha_mask_of_gray() {
+        let image: GrayImage = ImageBuffer::new(4, 4);
+        let mask = image.to_alpha_mask();
+        assert_eq!(mask.as_raw(), &[255; 16]);
+    }
+
+    #[test]
+    #[rustfmt::skip]
+    fn alpha_mask_extraction() {
+        let image: ImageBuffer<LumaA<u8>, _> = ImageBuffer::from_raw(4, 4, vec![
+            0, 1, 0, 2, 0, 3, 0, 4,
+            0, 5, 0, 6, 0, 7, 0, 8,
+            0, 9, 0, 10, 0, 11, 0, 12,
+            0, 13, 0, 14, 0, 15, 0, 16,
+        ]).unwrap();
+
+        let mask = image.to_alpha_mask();
+        assert_eq!(mask.as_raw(), &(1u8..17).collect::<Vec<_>>());
+    }
+
+    #[test]
+    fn apply_alpha_mask() {
+        let mut image: ImageBuffer<LumaA<u8>, _> = ImageBuffer::new(4, 4);
+
+        let alpha = ImageBuffer::from_pixel(4, 4, Luma([255]));
+        image.apply_alpha_channel(&alpha).expect("can apply");
+
+        for pixel in image.pixels() {
+            assert_eq!(pixel.0, [0, 255]);
+        }
+    }
+
+    #[test]
+    fn apply_alpha_mask_rgb() {
+        let mut image: ImageBuffer<Rgba<u8>, _> = ImageBuffer::new(4, 4);
+
+        let alpha = ImageBuffer::from_pixel(4, 4, Luma([255]));
+        image.apply_alpha_channel(&alpha).expect("can apply");
+
+        for pixel in image.pixels() {
+            assert_eq!(pixel.0, [0, 0, 0, 255]);
+        }
+    }
+
+    #[test]
+    fn can_not_apply_alpha_mask() {
+        ImageBuffer::<LumaA<u8>, _>::new(4, 4)
+            .apply_alpha_channel(&ImageBuffer::new(1, 1))
+            .expect_err("can not apply with wrong dimensions");
+
+        ImageBuffer::<Luma<u8>, _>::new(4, 4)
+            .apply_alpha_channel(&ImageBuffer::new(4, 4))
+            .expect_err("can not apply without alpha channel");
+        ImageBuffer::<Rgb<u8>, _>::new(4, 4)
+            .apply_alpha_channel(&ImageBuffer::new(4, 4))
+            .expect_err("can not apply without alpha channel");
+    }
 }
 
 #[cfg(test)]