computer_use/
screenshot.rs

1use std::io::Cursor;
2
3use image::imageops::FilterType;
4use image::{ImageFormat, RgbaImage};
5
6use crate::error::{Error, Result};
7use crate::geometry::{Point, Rect};
8
9/// Uncompressed RGBA8 pixels returned by a backend.
10#[derive(Clone, PartialEq, Eq)]
11pub struct RawImage {
12    /// Width in pixels.
13    pub width: u32,
14    /// Height in pixels.
15    pub height: u32,
16    /// Row-major RGBA bytes, `width * height * 4` long, no row padding.
17    pub rgba: Vec<u8>,
18}
19
20impl std::fmt::Debug for RawImage {
21    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
22        f.debug_struct("RawImage")
23            .field("width", &self.width)
24            .field("height", &self.height)
25            .field("bytes", &self.rgba.len())
26            .finish()
27    }
28}
29
30/// A PNG screenshot of a window together with the screen area it shows.
31///
32/// Coordinates handed to [`Computer::click`](crate::Computer::click) and
33/// [`Computer::drag`](crate::Computer::drag) are pixels in this image; the crate converts
34/// them back to screen points with [`Screenshot::to_screen_point`].
35#[derive(Clone, PartialEq)]
36pub struct Screenshot {
37    png: Vec<u8>,
38    width: u32,
39    height: u32,
40    frame: Rect,
41}
42
43impl std::fmt::Debug for Screenshot {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("Screenshot")
46            .field("width", &self.width)
47            .field("height", &self.height)
48            .field("frame", &self.frame)
49            .field("png_bytes", &self.png.len())
50            .finish()
51    }
52}
53
54impl Screenshot {
55    /// Downscales `image` so neither side exceeds `max_dimension`, then encodes it as PNG.
56    ///
57    /// `frame` is the screen area, in points, the image covers.
58    ///
59    /// # Errors
60    ///
61    /// Returns an error if the buffer size does not match the dimensions or encoding fails.
62    pub fn encode(image: RawImage, frame: Rect, max_dimension: Option<u32>) -> Result<Self> {
63        let RawImage { width, height, rgba } = image;
64        let mut buffer = RgbaImage::from_raw(width, height, rgba)
65            .ok_or_else(|| Error::platform("screenshot buffer does not match its dimensions"))?;
66        if let Some(max) = max_dimension.filter(|&m| m > 0)
67            && (width > max || height > max)
68        {
69            let scale = f64::from(max) / f64::from(width.max(height));
70            let w = scaled(width, scale);
71            let h = scaled(height, scale);
72            buffer = image::imageops::resize(&buffer, w, h, FilterType::Triangle);
73        }
74        let (width, height) = buffer.dimensions();
75        let mut png = Vec::new();
76        buffer.write_to(&mut Cursor::new(&mut png), ImageFormat::Png)?;
77        Ok(Self { png, width, height, frame })
78    }
79
80    /// PNG-encoded bytes.
81    #[must_use]
82    pub fn png(&self) -> &[u8] {
83        &self.png
84    }
85
86    /// Consumes the screenshot and returns the PNG bytes.
87    #[must_use]
88    pub fn into_png(self) -> Vec<u8> {
89        self.png
90    }
91
92    /// Width in pixels.
93    #[must_use]
94    pub const fn width(&self) -> u32 {
95        self.width
96    }
97
98    /// Height in pixels.
99    #[must_use]
100    pub const fn height(&self) -> u32 {
101        self.height
102    }
103
104    /// Screen area covered, in points.
105    #[must_use]
106    pub const fn frame(&self) -> Rect {
107        self.frame
108    }
109
110    /// Converts a pixel position in this image to a screen point.
111    #[must_use]
112    pub fn to_screen_point(&self, x: f64, y: f64) -> Point {
113        Point::new(
114            (x * self.frame.size.width).mul_add(1.0 / f64::from(self.width.max(1)), self.frame.origin.x),
115            (y * self.frame.size.height).mul_add(1.0 / f64::from(self.height.max(1)), self.frame.origin.y),
116        )
117    }
118
119    /// Converts a screen point to a pixel position in this image.
120    #[must_use]
121    pub fn from_screen_point(&self, point: Point) -> (f64, f64) {
122        let sx = f64::from(self.width) / self.frame.size.width.max(f64::EPSILON);
123        let sy = f64::from(self.height) / self.frame.size.height.max(f64::EPSILON);
124        ((point.x - self.frame.origin.x) * sx, (point.y - self.frame.origin.y) * sy)
125    }
126
127    /// True when `(x, y)` lies inside the image.
128    #[must_use]
129    pub fn contains_pixel(&self, x: f64, y: f64) -> bool {
130        x.is_finite()
131            && y.is_finite()
132            && x >= 0.0
133            && y >= 0.0
134            && x <= f64::from(self.width)
135            && y <= f64::from(self.height)
136    }
137}
138
139#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
140fn scaled(value: u32, scale: f64) -> u32 {
141    ((f64::from(value) * scale).round() as u32).max(1)
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    fn solid(width: u32, height: u32) -> RawImage {
149        RawImage { width, height, rgba: vec![255; (width * height * 4) as usize] }
150    }
151
152    #[test]
153    fn downscales_and_maps_coordinates() -> Result<()> {
154        let frame = Rect::new(100.0, 50.0, 400.0, 200.0);
155        let shot = Screenshot::encode(solid(800, 400), frame, Some(200))?;
156        assert_eq!((shot.width(), shot.height()), (200, 100));
157        assert!(shot.png().starts_with(b"\x89PNG"));
158        assert_eq!(shot.to_screen_point(100.0, 50.0), Point::new(300.0, 150.0));
159        assert_eq!(shot.from_screen_point(Point::new(300.0, 150.0)), (100.0, 50.0));
160        Ok(())
161    }
162
163    #[test]
164    fn rejects_mismatched_buffer() {
165        let bad = RawImage { width: 10, height: 10, rgba: vec![0; 4] };
166        assert!(Screenshot::encode(bad, Rect::default(), None).is_err());
167    }
168}
computer_use/screenshot.rs

computer_use/
screenshot.rs
