spectrust/
lib.rs

1// Importing necessary image processing and screenshot capturing modules.
2use image::{DynamicImage, GenericImageView, Pixel, Rgba};
3use screenshots::{DisplayInfo, Screen};
4
5// Function that takes a screenshot of a specified area.
6// It takes as parameters the x, y coordinates and the width, height of the desired area.
7fn screenshot(x: u16, y: u16, width: u16, height: u16) -> DynamicImage {
8    // Determine current display size
9    let display = size();
10    // Ensure the capture area is within the screen size
11    if !(x + width <= display.0 && y + height <= display.1) {
12        panic!("One or more specified parameter is not within the screen size. Use screen::size() to check.")
13    }
14    
15    // Retrieve screen based on specified coordinates
16    let screen = Screen::from_point(x.into(), y.into()).expect("Cannot get screen from specified x and y");
17
18    // Capture specified area of the screen
19    let capture = screen
20        .capture_area(x.into(), y.into(), width.into(), height.into())
21        .expect("Unable to screen capture.");
22
23    // Convert capture to image buffer
24    let buffer = capture.buffer();
25
26    // Load the image from memory buffer
27    let dynamic_image = image::load_from_memory(buffer).unwrap();
28
29    return dynamic_image;
30}
31
32
33
34// Function to get the size of the primary display.
35fn size() -> (u16, u16) {
36    // Retrieve all display info
37    let displays: Vec<DisplayInfo> = DisplayInfo::all().expect("Unable to get displays");
38    // Find primary display
39    let primary = displays
40        .iter()
41        .find(|display| display.is_primary == true)
42        .expect("Unable to find primary display");
43    // Return width and height of primary display
44    return (primary.width as u16, primary.height as u16);
45}
46
47// Function to locate an image on the screen with optional region, minimum confidence, and tolerance.
48// Returns coordinates, width, height and confidence if image is found, otherwise None.
49fn locate_on_screen(screen: &[Rgba<u8>], img: &[Rgba<u8>], screen_width: u32, screen_height: u32, img_width: u32, img_height: u32, min_confidence: f32, tolerance: u8) -> Option<(u32, u32, u32, u32, f32)> {
50    let step_size = 1;
51
52    for y in (0..screen_height - img_height).step_by(step_size) {
53        for x in (0..screen_width - img_width).step_by(step_size) {
54            let mut matching_pixels = 0;
55            let mut total_pixels = 0;
56
57            'outer: for dy in 0..img_height {
58                for dx in 0..img_width {
59                    let screen_idx: usize = ((y + dy) * screen_width + (x + dx)) as usize;
60                    let img_idx: usize = (dy * img_width + dx) as usize;
61
62                    let screen_pixel = screen[screen_idx];
63                    let img_pixel = img[img_idx];
64
65                    // Skip transparent pixels
66                    if img_pixel[3] < 128 {
67                        continue;
68                    }
69
70                    total_pixels += 1;
71
72                    if within_tolerance(screen_pixel[0], img_pixel[0], tolerance) &&
73                       within_tolerance(screen_pixel[1], img_pixel[1], tolerance) &&
74                       within_tolerance(screen_pixel[2], img_pixel[2], tolerance) {
75                        matching_pixels += 1;
76                    } else {
77                        break 'outer;
78                    }
79                }
80            }
81
82            let confidence = if total_pixels == 0 { 0.0 } else { matching_pixels as f32 / total_pixels as f32 };
83
84            if confidence >= min_confidence {
85                return Some((x, y, img_width, img_height, confidence));
86            }
87        }
88    }
89
90    None
91}
92
93
94// Function to locate center of an image within the screen with given parameters and tolerance.
95// Returns coordinates and confidence if image is found, otherwise None.
96fn locate_center_on_screen(screen: &[Rgba<u8>], img: &[Rgba<u8>], screen_width: u32, screen_height: u32, img_width: u32, img_height: u32, min_confidence: f32, tolerance: u8) -> Option<(u32, u32, f32)> {
97    let step_size = 1;
98
99    for y in (0..screen_height - img_height).step_by(step_size) {
100        for x in (0..screen_width - img_width).step_by(step_size) {
101            let mut matching_pixels = 0;
102            let mut total_pixels = 0;
103
104            'outer: for dy in 0..img_height {
105                for dx in 0..img_width {
106                    let screen_idx: usize = ((y + dy) * screen_width + (x + dx)) as usize;
107                    let img_idx: usize = (dy * img_width + dx) as usize;
108
109                    let screen_pixel = screen[screen_idx];
110                    let img_pixel = img[img_idx];
111
112                    // Skip transparent pixels
113                    if img_pixel[3] < 128 {
114                        continue;
115                    }
116
117                    total_pixels += 1;
118
119                    if within_tolerance(screen_pixel[0], img_pixel[0], tolerance) &&
120                       within_tolerance(screen_pixel[1], img_pixel[1], tolerance) &&
121                       within_tolerance(screen_pixel[2], img_pixel[2], tolerance) {
122                        matching_pixels += 1;
123                    } else {
124                        break 'outer;
125                    }
126                }
127            }
128
129            let confidence = if total_pixels == 0 { 0.0 } else { matching_pixels as f32 / total_pixels as f32 };
130
131            if confidence >= min_confidence {
132                return Some((x + img_width / 2, y + img_height / 2, confidence));
133            }
134        }
135    }   
136
137    None
138}
139
140
141
142// Helper function to check if a color value is within a tolerance range
143fn within_tolerance(value1: u8, value2: u8, tolerance: u8) -> bool {
144    let min_value = value2.saturating_sub(tolerance);
145    let max_value = value2.saturating_add(tolerance);
146    // Check if the color value is within tolerance range
147    value1 >= min_value && value1 <= max_value
148}
149
150
151
152// Function to locate the center of an image on the screen with optional region, minimum confidence, and tolerance.
153// Returns coordinates and confidence if image is found, otherwise None.
154pub fn locate_center_of_image(img: &DynamicImage, region: Option<(u16, u16, u16, u16)>, min_confidence: Option<f32>, tolerance: Option<u8>) -> Option<(u32, u32, f32)> {
155    // Default values
156    let (x, y, width, height) = region.unwrap_or((0, 0, size().0, size().1));
157    let min_confidence = min_confidence.unwrap_or(0.75);
158    let tolerance = tolerance.unwrap_or(25);
159
160    let img_pixels: Vec<_> = img.pixels().map(|p| p.2.to_rgba()).collect(); // use to_rgba
161    let img_width = img.width();
162
163    let img_height = img.height();
164
165    let screenshot = screenshot(x, y, width, height);
166    let screen_pixels: Vec<_> = screenshot.pixels().map(|p| p.2.to_rgba()).collect(); // use to_rgba
167    let screen_width = screenshot.width();
168    let screen_height = screenshot.height();
169
170    match locate_center_on_screen(
171        &screen_pixels,
172        &img_pixels,
173        screen_width,
174        screen_height,
175        img_width,
176        img_height,
177        min_confidence,
178        tolerance
179    ) {
180        Some((found_x, found_y, confidence)) => Some((found_x + x as u32, found_y + y as u32, confidence)), // Add region start position to the result
181        None => None,
182    }
183}
184
185
186// Function to locate an image on the screen with optional region, minimum confidence, and tolerance.
187// Returns coordinates, width, height and confidence if image is found, otherwise None.
188pub fn locate_image(img: &DynamicImage, region: Option<(u16, u16, u16, u16)>, min_confidence: Option<f32>, tolerance: Option<u8>) -> Option<(u32, u32, u32, u32, f32)> {
189    // Default values
190    let (x, y, width, height) = region.unwrap_or((0, 0, size().0, size().1));
191    let min_confidence = min_confidence.unwrap_or(0.75);
192    let tolerance = tolerance.unwrap_or(25);
193
194    let img_pixels: Vec<_> = img.pixels().map(|p| p.2.to_rgba()).collect();
195    let img_width = img.width();
196    let img_height = img.height();
197
198    let screenshot = screenshot(x, y, width, height);
199    let screen_pixels: Vec<_> = screenshot.pixels().map(|p| p.2.to_rgba()).collect();
200    let screen_width = screenshot.width();
201    let screen_height = screenshot.height();
202
203    locate_on_screen(
204        &screen_pixels,
205        &img_pixels,
206        screen_width,
207        screen_height,
208        img_width,
209        img_height,
210        min_confidence,
211        tolerance
212    )
213}