Add scanning

Fixes #14

Signed-off-by: Šimon Brandner <[email protected]>

src/geometry.rs

@@ -32,6 +32,12 @@ where
 {
 }
 
+fn calculate_length<T: Vec2DNumber>(min: T, max: T) -> T {
+	let min_f32 = <f32 as NumCast>::from(min).unwrap_or(0.0);
+	let max_f32 = <f32 as NumCast>::from(max).unwrap_or(0.0);
+	<T as NumCast>::from((max_f32 - min_f32).abs()).unwrap_or(T::zero())
+}
+
 #[derive(Debug, Clone, Copy)]
 pub struct Vec2D<T: Vec2DNumber> {
 	pub x: T,
@@ -88,7 +94,7 @@ impl<T: Vec2DNumber + Neg<Output = T>> Neg for Vec2D<T> {
 	}
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub struct Rectangle<T: Vec2DNumber> {
 	pub min: Vec2D<T>,
 	pub max: Vec2D<T>,
@@ -103,6 +109,18 @@ impl<T: Vec2DNumber> Rectangle<T> {
 		Rectangle::new(self.min.to_i32(), self.max.to_i32())
 	}
 
+	pub fn width(&self) -> T {
+		calculate_length(self.min.x, self.max.x)
+	}
+
+	pub fn height(&self) -> T {
+		calculate_length(self.min.y, self.max.y)
+	}
+
+	pub fn size(&self) -> (T, T) {
+		(self.width(), self.height())
+	}
+
 	pub fn intersection_over_union(&self, other: &Rectangle<T>) -> f32 {
 		if self.min.x > other.max.x {
 			return 0.0;

src/gui/mod.rs

@@ -4,10 +4,10 @@ use crate::{
 	camera::{self, Frame},
 	geometry::{Rectangle, Vec2D, Vec2DNumber},
 	gui::poi::draw_poi_square,
-	processors::{frame_processor::DETECTOR_INPUT_SIZE, DetectedFace},
+	processors::{frame_processor::DETECTOR_INPUT_SIZE, FaceForGUI, FaceForGUIAnnotation},
 };
 use eframe::{
-	egui::{self, Color32, ColorImage, Pos2, Rect, Rounding, Vec2},
+	egui::{self, Align2, Color32, ColorImage, FontFamily, FontId, Pos2, Rect, Rounding, Vec2},
 	EventLoopBuilderHook, NativeOptions,
 };
 use num::NumCast;
@@ -21,6 +21,11 @@ use std::{
 };
 use winit::platform::unix::EventLoopBuilderExtUnix;
 
+const FACE_RECTANGLE_WHITE_COLOR: Color32 = Color32::from_rgb(255, 255, 255);
+const FACE_RECTANGLE_GREY_COLOR: Color32 = Color32::from_rgb(192, 192, 192);
+const FACE_RECTANGLE_YELLOW_COLOR: Color32 = Color32::from_rgb(255, 255, 0);
+const LABEL_SHIFT: Vec2 = Vec2::new(10.0, 0.0);
+
 trait ToVec2 {
 	fn to_pos2(&self) -> Pos2;
 }
@@ -51,8 +56,6 @@ impl<T: Vec2DNumber> ToRect for Rectangle<T> {
 	}
 }
 
-const FACE_RECTANGLE_COLOR: Color32 = Color32::from_rgb(255, 255, 0);
-
 pub enum Error {
 	Camera(camera::Error),
 }
@@ -73,7 +76,7 @@ impl Display for Error {
 
 pub fn start(
 	frame: Arc<Mutex<Option<Frame>>>,
-	detected_faces: Arc<Mutex<Vec<DetectedFace>>>,
+	faces_for_gui: Arc<Mutex<Vec<FaceForGUI>>>,
 	finished: Arc<AtomicBool>,
 ) {
 	let event_loop_builder: Option<EventLoopBuilderHook> = Some(Box::new(|event_loop_builder| {
@@ -91,25 +94,25 @@ pub fn start(
 			)),
 			..NativeOptions::default()
 		},
-		Box::new(|_| Box::new(GUI::new(frame, detected_faces, finished))),
+		Box::new(|_| Box::new(GUI::new(frame, faces_for_gui, finished))),
 	);
 }
 
 struct GUI {
 	frame: Arc<Mutex<Option<Frame>>>,
-	detected_faces: Arc<Mutex<Vec<DetectedFace>>>,
+	faces_for_gui: Arc<Mutex<Vec<FaceForGUI>>>,
 	finished: Arc<AtomicBool>,
 }
 
 impl GUI {
 	pub fn new(
 		frame: Arc<Mutex<Option<Frame>>>,
-		detected_faces: Arc<Mutex<Vec<DetectedFace>>>,
+		faces_for_gui: Arc<Mutex<Vec<FaceForGUI>>>,
 		finished: Arc<AtomicBool>,
 	) -> Self {
 		Self {
 			frame,
-			detected_faces,
+			faces_for_gui,
 			finished,
 		}
 	}
@@ -153,15 +156,15 @@ impl eframe::App for GUI {
 			"Image height does not match network requirements!"
 		);
 
-		let detected_faces_lock = match self.detected_faces.lock() {
+		let faces_for_gui_lock = match self.faces_for_gui.lock() {
 			Ok(l) => l,
 			Err(e) => {
 				self.finished.store(true, Ordering::SeqCst);
 				panic!("Failed to get detected faces lock: {e}");
 			}
 		};
-		let detected_faces = detected_faces_lock.clone();
-		drop(detected_faces_lock);
+		let faces_for_gui = faces_for_gui_lock.clone();
+		drop(faces_for_gui_lock);
 
 		egui::CentralPanel::default()
 			.frame(egui::Frame::none().inner_margin(0.0).outer_margin(0.0))
@@ -181,14 +184,33 @@ impl eframe::App for GUI {
 					&image_texture,
 					Vec2::new(DETECTOR_INPUT_SIZE.x as f32, DETECTOR_INPUT_SIZE.y as f32),
 				);
-				for face in detected_faces {
-					let rectangles = draw_poi_square(face.rectangle);
+				for face_for_gui in faces_for_gui {
+					let (text, color) = match face_for_gui.annotation {
+						FaceForGUIAnnotation::Name(n) => (n, FACE_RECTANGLE_YELLOW_COLOR),
+						FaceForGUIAnnotation::Warning(w) => (w, FACE_RECTANGLE_GREY_COLOR),
+						FaceForGUIAnnotation::ScanningState {
+							scanned_sample_count,
+							required_sample_count,
+						} => (
+							format!(
+								"Scanning: {}/{}",
+								scanned_sample_count, required_sample_count
+							),
+							FACE_RECTANGLE_WHITE_COLOR,
+						),
+					};
+					let (rectangles, top_right_position) = draw_poi_square(face_for_gui.rectangle);
+
+					ui.painter().text(
+						top_right_position.to_pos2() + LABEL_SHIFT,
+						Align2::LEFT_TOP,
+						text,
+						FontId::new(16.0, FontFamily::Monospace),
+						FACE_RECTANGLE_WHITE_COLOR,
+					);
 					for rectangle in rectangles {
-						ui.painter().rect_filled(
-							rectangle.to_rect(),
-							Rounding::default(),
-							FACE_RECTANGLE_COLOR,
-						);
+						ui.painter()
+							.rect_filled(rectangle.to_rect(), Rounding::default(), color);
 					}
 				}
 			});

src/gui/poi.rs

@@ -284,7 +284,7 @@ fn draw_side_rectangles(
 
 /// As a Person of Interest fan, I decided it would be fun to draw Person of Interest-like
 /// squares around the detected faces, so that is what this function does
-pub fn draw_poi_square(rectangle: Rectangle<u32>) -> Vec<Rectangle<i32>> {
+pub fn draw_poi_square(rectangle: Rectangle<u32>) -> (Vec<Rectangle<i32>>, Vec2D<i32>) {
 	let small_rectangles_offset = (CORNER_RECTANGLE_SIZE.y - SIDE_RECTANGLE_SIZE.y) / 2;
 	let (
 		square_size,
@@ -314,5 +314,5 @@ pub fn draw_poi_square(rectangle: Rectangle<u32>) -> Vec<Rectangle<i32>> {
 		&bottom_right_position,
 	));
 
-	rectangles
+	(rectangles, top_right_position)
 }

src/main.rs

@@ -7,7 +7,7 @@ mod processors;
 use camera::Frame;
 use clap::Parser;
 use core::panic;
-use processors::DetectedFace;
+use processors::FaceForGUI;
 use std::sync::atomic::AtomicBool;
 use std::sync::{Arc, Mutex};
 use std::thread::{self};
@@ -39,19 +39,19 @@ fn main() {
 	}
 
 	let frame: Arc<Mutex<Option<Frame>>> = Arc::new(Mutex::new(None));
-	let detected_faces: Arc<Mutex<Vec<DetectedFace>>> = Arc::new(Mutex::new(Vec::new()));
+	let faces_for_gui: Arc<Mutex<Vec<FaceForGUI>>> = Arc::new(Mutex::new(Vec::new()));
 	let finished: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
 
 	let frame_clone = frame.clone();
 	let finished_clone = finished.clone();
 	thread::spawn(move || camera::start(frame_clone, finished_clone));
 
 	if !args.no_gui {
-		let detected_faces_clone = detected_faces.clone();
+		let faces_for_gui_clone = faces_for_gui.clone();
 		let frame_clone = frame.clone();
 		let finished_clone = finished.clone();
-		thread::spawn(move || gui::start(frame_clone, detected_faces_clone, finished_clone));
+		thread::spawn(move || gui::start(frame_clone, faces_for_gui_clone, finished_clone));
 	}
 
-	let _ = thread::spawn(move || processors::start(frame, detected_faces, finished)).join();
+	let _ = thread::spawn(move || processors::start(frame, faces_for_gui, finished)).join();
 }

src/processors/face_processor.rs

@@ -1,24 +1,34 @@
-use super::DetectedFace;
+use crate::processors::FaceRecognitionError;
+
+use super::{FaceEmbedding, FaceForGUI, FaceForGUIAnnotation, FaceForProcessing};
 use std::fmt::Debug;
 
 const SIMILARITY_THRESHOLD: f32 = 0.51;
+const SCAN_SAMPLE_COUNT: usize = 16;
 
 pub trait FaceProcessor: Debug {
-	fn process_detected_faces(&mut self, detected_faces: &[DetectedFace]);
+	fn process_detected_faces(&mut self, detected_faces: Vec<FaceForProcessing>)
+		-> Vec<FaceForGUI>;
 	fn is_finished(&self) -> bool;
 }
 
-#[derive(Debug, Clone, Default)]
-pub struct ScanProcessorResult {}
+#[derive(Debug, Clone)]
+pub struct ScanProcessorResult {
+	face_embedding: FaceEmbedding,
+}
 
 #[derive(Debug)]
 pub struct ScanProcessor {
 	result: Option<ScanProcessorResult>,
+	embedding_samples: Vec<FaceEmbedding>,
 }
 
 impl ScanProcessor {
 	pub fn new() -> Self {
-		Self { result: None }
+		Self {
+			result: None,
+			embedding_samples: Vec::new(),
+		}
 	}
 
 	pub fn get_result(&self) -> Option<ScanProcessorResult> {
@@ -27,7 +37,74 @@ impl ScanProcessor {
 }
 
 impl FaceProcessor for ScanProcessor {
-	fn process_detected_faces(&mut self, detected_faces: &[DetectedFace]) {}
+	fn process_detected_faces(
+		&mut self,
+		faces_for_processing: Vec<FaceForProcessing>,
+	) -> Vec<FaceForGUI> {
+		// Handle edge-cases
+		if faces_for_processing.len() > 1 {
+			self.embedding_samples.clear();
+			return faces_for_processing
+				.into_iter()
+				.map(|f| FaceForGUI {
+					rectangle: f.rectangle,
+					annotation: super::FaceForGUIAnnotation::Warning(String::from(
+						"Too many faces for scanning",
+					)),
+				})
+				.collect();
+		};
+		let face_for_processing = match faces_for_processing.get(0) {
+			Some(f) => f,
+			None => {
+				self.embedding_samples.clear();
+				return vec![];
+			}
+		};
+		let embedding = match face_for_processing.face_data {
+			Ok(f) => f.embedding,
+			Err(e) => match e {
+				FaceRecognitionError::TooSmall => {
+					self.embedding_samples.clear();
+					return vec![FaceForGUI {
+						rectangle: face_for_processing.rectangle,
+						annotation: FaceForGUIAnnotation::Warning(String::from("Too small")),
+					}];
+				}
+			},
+		};
+
+		// We require multiple samples during the scan for safety purposes. All of these samples
+		// have to satisfy the similarity requirement. To avoid having to check each two samples we
+		// take and an average of all of them. We calculate the average in each run of this
+		// function to give all samples the same weight
+		if self.embedding_samples.len() > 0 {
+			let average_embedding = FaceEmbedding::average_embedding(&self.embedding_samples);
+			let similarity = average_embedding.cosine_similarity(&embedding);
+
+			if similarity < SIMILARITY_THRESHOLD {
+				self.embedding_samples.clear();
+				return vec![];
+			}
+		}
+		self.embedding_samples.push(embedding);
+
+		// If we have enough samples, we consider the scan to be successful, so we set the result
+		if self.embedding_samples.len() > SCAN_SAMPLE_COUNT {
+			self.result = Some(ScanProcessorResult {
+				face_embedding: FaceEmbedding::average_embedding(&self.embedding_samples),
+			})
+		};
+
+		// Return info to be display in the GUI
+		vec![FaceForGUI {
+			rectangle: face_for_processing.rectangle,
+			annotation: super::FaceForGUIAnnotation::ScanningState {
+				scanned_sample_count: self.embedding_samples.len(),
+				required_sample_count: SCAN_SAMPLE_COUNT,
+			},
+		}]
+	}
 
 	fn is_finished(&self) -> bool {
 		self.result.is_some()
@@ -53,7 +130,12 @@ impl AuthProcessor {
 }
 
 impl FaceProcessor for AuthProcessor {
-	fn process_detected_faces(&mut self, detected_faces: &[DetectedFace]) {}
+	fn process_detected_faces(
+		&mut self,
+		detected_faces: Vec<FaceForProcessing>,
+	) -> Vec<FaceForGUI> {
+		vec![]
+	}
 
 	fn is_finished(&self) -> bool {
 		self.result.is_some()