feat: add support for BGR888 image format

waycrate · Jan 8, 2024 · 3eead85 · 3eead85
1 parent edd1f14
commit 3eead85
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 9 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,8 +1,8 @@
-use image::RgbaImage;
+use image::{DynamicImage, RgbaImage};
 use wayland_client::protocol::wl_output::Transform;
 
 pub(crate) fn rotate_image_buffer(
-    image: RgbaImage,
+    image: DynamicImage,
     transform: Transform,
     width: u32,
     height: u32,
@@ -24,7 +24,7 @@ pub(crate) fn rotate_image_buffer(
             let flipped_buffer = image::imageops::flip_horizontal(&image);
             image::imageops::rotate270(&flipped_buffer)
         }
-        _ => image,
+        _ => image.into_rgba8(),
     };
 
     if final_buffer.dimensions() == (width, height) {

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, RgbaImage};
 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();
@@ -506,7 +507,8 @@ impl WayshotConnection {
             output_info.transform,
             None,
         )?;
-        frame_copy.try_into()
+        let image = DynamicImage::try_from(frame_copy)?;
+        Ok(image.into_rgba8())
     }
 
     /// Take a screenshot from all of the specified outputs.

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),
         })