feat: add support for BGR888 image format (#82)

* feat: add support for BGR888 image format * fix: don't coerce images to rgba
waycrate · Jan 20, 2024 · cb6bd68 · cb6bd68
1 parent edd1f14
commit cb6bd68
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 23 deletions.
diff --git a/libwayshot/src/convert.rs b/libwayshot/src/convert.rs
@@ -15,6 +15,9 @@ struct ConvertNone {}
 #[derive(Default)]
 struct ConvertRGB8 {}
 
+#[derive(Default)]
+struct ConvertBGR888 {}
+
 const SHIFT10BITS_1: u32 = 20;
 const SHIFT10BITS_2: u32 = 10;
 
@@ -27,6 +30,7 @@ pub fn create_converter(format: wl_shm::Format) -> Option<Box<dyn Convert>> {
         wl_shm::Format::Xbgr2101010 | wl_shm::Format::Abgr2101010 => {
             Some(Box::<ConvertBGR10>::default())
         }
+        wl_shm::Format::Bgr888 => Some(Box::<ConvertBGR888>::default()),
         _ => None,
     }
 }
@@ -69,3 +73,9 @@ impl Convert for ConvertBGR10 {
         ColorType::Rgba8
     }
 }
+
+impl Convert for ConvertBGR888 {
+    fn convert_inplace(&self, _data: &mut [u8]) -> ColorType {
+        ColorType::Rgb8
+    }
+}
diff --git a/libwayshot/src/image_util.rs b/libwayshot/src/image_util.rs
@@ -1,40 +1,41 @@
-use image::RgbaImage;
+use image::{DynamicImage, GenericImageView};
 use wayland_client::protocol::wl_output::Transform;
 
 pub(crate) fn rotate_image_buffer(
-    image: RgbaImage,
+    image: DynamicImage,
     transform: Transform,
     width: u32,
     height: u32,
-) -> RgbaImage {
-    let final_buffer = match transform {
-        Transform::_90 => image::imageops::rotate90(&image),
-        Transform::_180 => image::imageops::rotate180(&image),
-        Transform::_270 => image::imageops::rotate270(&image),
-        Transform::Flipped => image::imageops::flip_horizontal(&image),
+) -> DynamicImage {
+    let final_image = match transform {
+        Transform::_90 => image::imageops::rotate90(&image).into(),
+        Transform::_180 => image::imageops::rotate180(&image).into(),
+        Transform::_270 => image::imageops::rotate270(&image).into(),
+        Transform::Flipped => image::imageops::flip_horizontal(&image).into(),
         Transform::Flipped90 => {
             let flipped_buffer = image::imageops::flip_horizontal(&image);
-            image::imageops::rotate90(&flipped_buffer)
+            image::imageops::rotate90(&flipped_buffer).into()
         }
         Transform::Flipped180 => {
             let flipped_buffer = image::imageops::flip_horizontal(&image);
-            image::imageops::rotate180(&flipped_buffer)
+            image::imageops::rotate180(&flipped_buffer).into()
         }
         Transform::Flipped270 => {
             let flipped_buffer = image::imageops::flip_horizontal(&image);
-            image::imageops::rotate270(&flipped_buffer)
+            image::imageops::rotate270(&flipped_buffer).into()
         }
         _ => image,
     };
 
-    if final_buffer.dimensions() == (width, height) {
-        return final_buffer;
+    if final_image.dimensions() == (width, height) {
+        return final_image;
     }
 
     image::imageops::resize(
-        &final_buffer,
+        &final_image,
         width,
         height,
         image::imageops::FilterType::Gaussian,
     )
+    .into()
 }
diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs
@@ -19,7 +19,7 @@ use std::{
     thread,
 };
 
-use image::{imageops::overlay, RgbaImage};
+use image::{imageops::overlay, DynamicImage};
 use memmap2::MmapMut;
 use wayland_client::{
     globals::{registry_queue_init, GlobalList},
@@ -254,6 +254,7 @@ impl WayshotConnection {
                         | wl_shm::Format::Argb8888
                         | wl_shm::Format::Xrgb8888
                         | wl_shm::Format::Xbgr8888
+                        | wl_shm::Format::Bgr888
                 )
             })
             .copied();
@@ -443,7 +444,7 @@ impl WayshotConnection {
         &self,
         capture_region: CaptureRegion,
         cursor_overlay: bool,
-    ) -> Result<RgbaImage> {
+    ) -> Result<DynamicImage> {
         let (frame_copies, (width, height)) =
             self.create_frame_copy(capture_region, cursor_overlay)?;
 
@@ -499,7 +500,7 @@ impl WayshotConnection {
         &self,
         output_info: &OutputInfo,
         cursor_overlay: bool,
-    ) -> Result<RgbaImage> {
+    ) -> Result<DynamicImage> {
         let frame_copy = self.capture_output_frame(
             cursor_overlay,
             &output_info.wl_output,
@@ -514,7 +515,7 @@ impl WayshotConnection {
         &self,
         outputs: &Vec<OutputInfo>,
         cursor_overlay: bool,
-    ) -> Result<RgbaImage> {
+    ) -> Result<DynamicImage> {
         if outputs.is_empty() {
             return Err(Error::NoOutputs);
         }
@@ -549,7 +550,7 @@ impl WayshotConnection {
     }
 
     /// Take a screenshot from all accessible outputs.
-    pub fn screenshot_all(&self, cursor_overlay: bool) -> Result<RgbaImage> {
+    pub fn screenshot_all(&self, cursor_overlay: bool) -> Result<DynamicImage> {
         self.screenshot_outputs(self.get_all_outputs(), cursor_overlay)
     }
 }
diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs
@@ -4,7 +4,7 @@ use std::{
     time::{SystemTime, UNIX_EPOCH},
 };
 
-use image::{ColorType, ImageBuffer, Pixel, RgbaImage};
+use image::{ColorType, DynamicImage, ImageBuffer, Pixel};
 use memmap2::MmapMut;
 use nix::{
     fcntl,
@@ -46,13 +46,16 @@ pub struct FrameCopy {
     pub transform: wl_output::Transform,
 }
 
-impl TryFrom<FrameCopy> for RgbaImage {
+impl TryFrom<FrameCopy> for DynamicImage {
     type Error = Error;
 
     fn try_from(value: FrameCopy) -> Result<Self> {
         Ok(match value.frame_color_type {
-            ColorType::Rgb8 | ColorType::Rgba8 => {
-                create_image_buffer(&value.frame_format, &value.frame_mmap)?
+            ColorType::Rgb8 => {
+                Self::ImageRgb8(create_image_buffer(&value.frame_format, &value.frame_mmap)?)
+            }
+            ColorType::Rgba8 => {
+                Self::ImageRgba8(create_image_buffer(&value.frame_format, &value.frame_mmap)?)
             }
             _ => return Err(Error::InvalidColor),
         })