depth_analyzer/
lib.rs

1//! # Depth Analyzer
2//!
3//! Program that analyzes an image processed by AI depth detection models.
4//!
5//! ## Usage
6//!
7//! `$ depth-analyzer /path/to/image.[jpg | png | webp]`
8//!
9//! ### Options
10//!
11//! `-h`, `--help`                      Displays this help menu.
12//!
13//! `-v`, `--version`                   Displays version.
14//!
15//! `-c`, `--color` **[ RED | WHITE ]**     Specifies which color to 
16//!                                 use as an indicator for proximity.
17//!
18//! `-t`, `--threshold` **[ 0 .. 255 ]**    Specifies the value a pixel must have in order to be 
19//!                                 considered to be of the proximity color.
20//!
21//! `-w`, `--watch` **[ /path/to/images ]**     Analyze images as they come in. If no path is
22//!                                 provided, current directory is used.
23//!
24//! ## Possible Results
25//!
26//! In order of precedence:
27//!
28//! - `FORWARD`
29//! - `RIGHT`
30//! - `LEFT`
31//! - `STOP`
32
33
34use std::{fmt::Display, process, fs};
35
36use image::GenericImageView;
37
38pub mod config;
39
40#[derive(PartialEq, Debug)]
41pub enum Instruction {
42    Left,
43    Forward,
44    Right,
45    Stop
46}
47
48impl Display for Instruction {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        match self {
51            Self::Right => write!(f, "{}", "RIGHT"),
52            Self::Left => write!(f, "{}", "LEFT"),
53            Self::Forward => write!(f, "{}", "FORWARD"),
54            Self::Stop => write!(f, "{}", "STOP"),
55        }
56    }
57}
58
59/// Each sector of the image has a **total red pixels** and a **total pixel** count.
60/// A **red pixel** is a pixel that has a Red value higher than 150.
61#[derive(PartialEq, Debug)]
62pub struct DangerSectors {
63    left :(u32, u32),
64    center :(u32, u32),
65    right :(u32, u32),
66}
67
68impl DangerSectors {
69    pub fn new() -> DangerSectors {
70        DangerSectors {
71            left: (0, 0),
72            center: (0, 0),
73            right: (0, 0),
74        }
75    }
76
77    /// Updates sectors by analyzing the image.
78    pub fn analyze(&mut self, img_config :&mut config::ImageConfig) {
79
80        if let Some(watch_dir) = img_config.watch_dir.take() {
81            loop {
82                if fs::read_dir(&watch_dir).unwrap().count() == 0 {
83                    continue;
84                }
85
86                for image_option in fs::read_dir(&watch_dir).unwrap() {
87                    let image = image_option.expect("Error: Image does not exist!");
88
89                    match &image.path().extension() {
90
91                        Some(ext) if ext.to_str().unwrap() == "png" ||
92                                     ext.to_str().unwrap() == "jpg" ||
93                                     ext.to_str().unwrap() == "jpeg" ||
94                                     ext.to_str().unwrap() == "webp" => {
95
96                            match image::open(&image.path()) {
97                                Ok(res) => {
98                                    img_config.img = Some(res);
99                                    self.count_pixels(img_config);
100                                    println!("{} : {}", image
101                                             .file_name()
102                                             .to_str()
103                                             .unwrap(), 
104                                             self.get_instruction().to_string());
105                                },
106
107                                _ => {
108                                    eprintln!("Error: Could not open file: {}", 
109                                              image
110                                              .path()
111                                              .to_str()
112                                              .unwrap()); 
113                                    process::exit(1);
114                                }
115                            }
116                        },
117
118                        _ => continue,
119
120                    };
121
122                }
123            }
124        }
125
126        self.count_pixels(img_config);
127
128    }
129
130    /// Counts the pixels that meet the threshold in each sector.
131    fn count_pixels(&mut self, img_config :&mut config::ImageConfig) {
132
133        let img = match img_config.img.take() {
134            Some(image) => image,
135            None => {
136                eprintln!("Error no image specified!\n");
137                process::exit(1);
138            }
139        };
140
141        let image_width = img.dimensions().0;
142
143        for (x, _, rgba) in img.pixels() {
144            
145            if x < image_width / 3 {
146
147                if config::check_threshold(&img_config.proximity_color, img_config.threshold, &rgba) {
148                    self.left.0 += 1;
149                }
150                self.left.1 += 1;
151
152            } else if x < 2 * image_width / 3 {
153
154                if config::check_threshold(&img_config.proximity_color, img_config.threshold, &rgba) {
155                    self.center.0 += 1;
156                }
157                self.center.1 += 1;
158
159            } else {
160
161                if config::check_threshold(&img_config.proximity_color, img_config.threshold, &rgba) {
162                    self.right.0 += 1;
163                }
164                self.right.1 += 1;
165
166            }
167        }
168    }
169
170    /// Gets the ratio of *red pixels* versus *total pixels* for each sector.
171    fn get_ratios(&self) -> (f32, f32, f32) {
172
173        (
174            self.left.0 as f32 / self.left.1 as f32,
175            self.center.0  as f32 / self.center.1 as f32,
176            self.right.0 as f32 / self.right.1 as f32,
177        )
178    }
179
180    /// Based on the dangers in each sector, determines the outcome for the user.
181    ///
182    /// NOTE: The precedence is **Forward**, **Right**, **Left**
183    /// NOTE: If the ratio of Red pixels in each sector is >= 50% then **Stop**.
184    pub fn get_instruction(&self) -> Instruction {
185        let (left, center, right) = self.get_ratios();
186
187        if left >= 0.5 && center >= 0.5 && right >= 0.5 {
188            Instruction::Stop
189
190        } else if center <= right && center <= left {
191            Instruction::Forward
192
193        } else if right < center && right <= left {
194            Instruction::Right
195
196        } else {
197            Instruction::Left
198
199        }
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206
207    #[test]
208    fn room() {
209        let mut img_config = config::ImageConfig {
210            proximity_color: config::ModelProximityColor::RED,
211            threshold: 150,
212            img: Some( image::open("test_images/empty_room.png").unwrap() ),
213            watch_dir: None,
214        };
215
216        let mut sectors = DangerSectors::new();
217        sectors.analyze(&mut img_config);
218
219        assert_eq!(Instruction::Forward, sectors.get_instruction());
220    }
221
222    #[test]
223    fn chair() {
224        let mut img_config = config::ImageConfig {
225            proximity_color: config::ModelProximityColor::RED,
226            threshold: 150,
227            img: Some( image::open("test_images/chair.png").unwrap() ),
228            watch_dir: None,
229        };
230
231
232        let mut sectors = DangerSectors::new();
233        sectors.analyze(&mut img_config);
234
235        assert_eq!(Instruction::Right, sectors.get_instruction());
236    }
237
238    #[test]
239    fn library() {
240        let mut img_config = config::ImageConfig {
241            proximity_color: config::ModelProximityColor::WHITE,
242            threshold: 200,
243            img: Some( image::open("test_images/chair.png").unwrap() ),
244            watch_dir: None,
245        };
246
247
248        let mut sectors = DangerSectors::new();
249        sectors.analyze(&mut img_config);
250
251        assert_eq!(Instruction::Forward, sectors.get_instruction());
252    }
253}
depth_analyzer/lib.rs

depth_analyzer/
lib.rs
