Test models

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

crates/oblichey-cli/Cargo.toml

@@ -8,6 +8,11 @@ edition.workspace = true
 burn-import = { git = "https://github.com/tracel-ai/burn", rev = "f7639bd35a1120fccc849dcb94fbab162df7103a" }
 merkle_hash = "3.7.0"
 
+[dev-dependencies]
+burn = { git = "https://github.com/tracel-ai/burn", rev = "f7639bd35a1120fccc849dcb94fbab162df7103a", features = [
+  "ndarray",
+] }
+
 [dependencies]
 v4l = "0.14.0"
 chrono = "0.4.38"

crates/oblichey-cli/src/models/detector.rs

@@ -105,3 +105,197 @@ impl<B: Backend> FaceDetector<B> {
 		face_rectangles
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use super::FaceDetector;
+	use crate::{
+		camera::Frame,
+		geometry::{Rectangle, Vec2D},
+		models::detector::{CONFIDENCE_THRESHOLD, DETECTOR_INPUT_SIZE},
+		processors::frame_processor::BurnBackend,
+	};
+	use burn::{
+		backend::{ndarray::NdArrayDevice, NdArray},
+		tensor::{Tensor, TensorData},
+	};
+	use image::ImageBuffer;
+
+	const FRAME_CHANNEL_SIZE: usize = (DETECTOR_INPUT_SIZE.x * DETECTOR_INPUT_SIZE.y) as usize;
+	const FRAME_VEC_SIZE: usize = FRAME_CHANNEL_SIZE * 3;
+
+	fn get_device() -> NdArrayDevice {
+		NdArrayDevice::default()
+	}
+
+	fn get_face_detector() -> FaceDetector<NdArray<f32>> {
+		FaceDetector::new(&get_device())
+	}
+
+	fn get_frame(data: Vec<u8>) -> Frame {
+		ImageBuffer::from_vec(DETECTOR_INPUT_SIZE.x, DETECTOR_INPUT_SIZE.y, data)
+			.expect("Failed to construct frame")
+	}
+
+	fn check_channel(channel: &[f32], expected_value: f32) {
+		let mut current_expected_value: f32 = 0.0;
+		for (index, v) in channel.iter().enumerate() {
+			// Change expectation at the end of a row
+			if index as u32 % DETECTOR_INPUT_SIZE.x == 0 {
+				current_expected_value =
+					if (current_expected_value - expected_value).abs() < f32::EPSILON {
+						0.0
+					} else {
+						expected_value
+					}
+			}
+
+			println!("{v}");
+
+			assert!((*v - current_expected_value).abs() < f32::EPSILON);
+		}
+	}
+
+	#[test]
+	fn normalizes_input_between_values() {
+		let face_detector = get_face_detector();
+		let test_cases = vec![
+			(vec![255; FRAME_VEC_SIZE], 1.0),
+			(vec![191; FRAME_VEC_SIZE], 0.5),
+			(vec![63; FRAME_VEC_SIZE], -0.5),
+		];
+
+		for (frame_data, expected_result) in test_cases {
+			let frame = get_frame(frame_data);
+
+			let normalized = face_detector.normalize_input(&frame);
+			let vector = normalized.to_data().to_vec::<f32>().expect("Failed");
+
+			for element in vector {
+				assert!((element - expected_result).abs() < f32::EPSILON);
+			}
+		}
+	}
+
+	#[test]
+	fn permutes_dimensions_during_normalization() {
+		#[allow(clippy::identity_op, clippy::erasing_op)]
+		let zero = 0 * 128 + 127; // This will become 0 after normalization
+		let zero_value_by_channel = vec![zero; 3]; // This will normalize to [0.0, 0.0, 0.0]
+		let value_by_channel = vec![159, 191, 255]; // This will normalize to [0.25, 0.5, 1.0]
+		let value_by_channel_normalized: Vec<f32> = value_by_channel
+			.iter()
+			.map(|&value| (value as f32 - 127.0) / 128.0)
+			.collect();
+		let frame_data = {
+			let mut data = vec![];
+			for i in 0..DETECTOR_INPUT_SIZE.y {
+				for _ in 0..DETECTOR_INPUT_SIZE.x {
+					data.extend(if i % 2 == 0 {
+						value_by_channel.clone()
+					} else {
+						zero_value_by_channel.clone()
+					});
+				}
+			}
+			data
+		};
+		let face_recognizer = get_face_detector();
+		let frame = get_frame(frame_data);
+		let normalized = face_recognizer.normalize_input(&frame);
+		let vector = normalized.to_data().to_vec::<f32>().expect("Failed");
+
+		for (index, value) in value_by_channel_normalized.iter().enumerate() {
+			let channel = &vector[FRAME_CHANNEL_SIZE * index..FRAME_CHANNEL_SIZE * (index + 1)];
+			check_channel(channel, *value);
+		}
+	}
+
+	#[test]
+	fn interprets_output() {
+		fn create_confidences_tensor(
+			confidences: &[f32],
+			device: &NdArrayDevice,
+		) -> Tensor<BurnBackend, 3> {
+			let tensor_data_vec: Vec<f32> = confidences
+				.iter()
+				.flat_map(|&confidence| [0.0, confidence])
+				.collect();
+
+			Tensor::from_data(
+				TensorData::new(tensor_data_vec.clone(), [tensor_data_vec.len(), 1, 1]),
+				device,
+			)
+		}
+		fn create_rectangles_tensor(
+			rectangles: &[Rectangle<u32>],
+			device: &NdArrayDevice,
+		) -> Tensor<BurnBackend, 3> {
+			let tensor_data_vec: Vec<f32> = rectangles
+				.iter()
+				.flat_map(|rect| {
+					vec![
+						(rect.min.x as f32) / (DETECTOR_INPUT_SIZE.x as f32),
+						(rect.min.y as f32) / (DETECTOR_INPUT_SIZE.y as f32),
+						(rect.max.x as f32) / (DETECTOR_INPUT_SIZE.x as f32),
+						(rect.max.y as f32) / (DETECTOR_INPUT_SIZE.y as f32),
+					]
+				})
+				.collect();
+
+			Tensor::from_data(
+				TensorData::new(tensor_data_vec.clone(), [tensor_data_vec.len(), 1, 1]),
+				device,
+			)
+		}
+
+		let device = get_device();
+		let test_cases = vec![
+			(
+				vec![0.0, 1.0],
+				vec![
+					Rectangle::new(Vec2D::new(0, 0), Vec2D::new(10, 10)),
+					Rectangle::new(Vec2D::new(0, 0), Vec2D::new(10, 10)),
+				],
+			),
+			(
+				vec![0.0, 1.0, 0.0, 1.0],
+				vec![
+					Rectangle::new(Vec2D::new(0, 0), Vec2D::new(10, 10)),
+					Rectangle::new(Vec2D::new(0, 0), Vec2D::new(10, 10)),
+					Rectangle::new(Vec2D::new(0, 0), Vec2D::new(20, 20)),
+					Rectangle::new(Vec2D::new(0, 0), Vec2D::new(20, 20)),
+				],
+			),
+		];
+
+		for (confidences, rectangles) in test_cases {
+			let confidences_tensor = create_confidences_tensor(&confidences, &device);
+			let rectangles_tensor = create_rectangles_tensor(&rectangles, &device);
+			let valid_rectangle_count = confidences.iter().fold(0, |count, confidence| {
+				if *confidence > CONFIDENCE_THRESHOLD {
+					count + 1
+				} else {
+					count
+				}
+			});
+
+			let interpreted =
+				FaceDetector::interpret_output(&(confidences_tensor, rectangles_tensor));
+
+			let mut valid_rectangles = vec![];
+			for (index, rectangle) in rectangles.iter().enumerate() {
+				if confidences[index] < CONFIDENCE_THRESHOLD {
+					continue;
+				}
+
+				valid_rectangles.push(rectangle);
+			}
+
+			assert_eq!(interpreted.len(), valid_rectangle_count);
+			for i in 0..interpreted.len() {
+				assert_eq!(interpreted[i], *valid_rectangles[i]);
+			}
+		}
+	}
+}

crates/oblichey-cli/src/models/recognizer.rs

@@ -78,3 +78,129 @@ impl<B: Backend> FaceRecognizer<B> {
 		FaceRecognitionData { embedding }
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use super::{FaceRecognizer, RECOGNIZER_INPUT_SIZE};
+	use crate::{
+		camera::Frame,
+		processors::{
+			face::{FaceEmbedding, EMBEDDING_LENGTH},
+			frame_processor::BurnBackend,
+		},
+	};
+	use burn::{
+		backend::{ndarray::NdArrayDevice, NdArray},
+		tensor::{Tensor, TensorData},
+	};
+	use image::ImageBuffer;
+
+	const FRAME_CHANNEL_SIZE: usize = (RECOGNIZER_INPUT_SIZE.x * RECOGNIZER_INPUT_SIZE.y) as usize;
+	const FRAME_VEC_SIZE: usize = FRAME_CHANNEL_SIZE * 3;
+
+	fn get_device() -> NdArrayDevice {
+		NdArrayDevice::default()
+	}
+
+	fn get_face_recognizer() -> FaceRecognizer<NdArray<f32>> {
+		FaceRecognizer::new(&get_device())
+	}
+
+	fn get_frame(data: Vec<u8>) -> Frame {
+		ImageBuffer::from_vec(RECOGNIZER_INPUT_SIZE.x, RECOGNIZER_INPUT_SIZE.y, data)
+			.expect("Failed to construct frame")
+	}
+
+	fn check_channel(channel: &[f32], expected_value: f32) {
+		let mut current_expected_value: f32 = 0.0;
+		for (index, v) in channel.iter().enumerate() {
+			// Change expectation at the end of a row
+			if index as u32 % RECOGNIZER_INPUT_SIZE.x == 0 {
+				current_expected_value =
+					if (current_expected_value - expected_value).abs() < f32::EPSILON {
+						0.0
+					} else {
+						expected_value
+					}
+			}
+
+			println!("{v}");
+
+			assert!((*v - current_expected_value).abs() < f32::EPSILON);
+		}
+	}
+
+	#[test]
+	fn normalizes_input_between_values() {
+		let face_recognizer = get_face_recognizer();
+		let test_cases = vec![
+			(vec![255; FRAME_VEC_SIZE], 1.0),
+			(vec![191; FRAME_VEC_SIZE], 0.5),
+			(vec![63; FRAME_VEC_SIZE], -0.5),
+		];
+
+		for (frame_data, expected_result) in test_cases {
+			let frame = get_frame(frame_data);
+
+			let normalized = face_recognizer.normalize_input(&frame);
+			let vector = normalized.to_data().to_vec::<f32>().expect("Failed");
+
+			for element in vector {
+				assert!((element - expected_result).abs() < f32::EPSILON);
+			}
+		}
+	}
+
+	#[test]
+	fn permutes_dimensions_during_normalization() {
+		#[allow(clippy::identity_op, clippy::erasing_op)]
+		let zero = 0 * 128 + 127; // This will become 0 after normalization
+		let zero_value_by_channel = vec![zero; 3]; // This will normalize to [0.0, 0.0, 0.0]
+		let value_by_channel = vec![159, 191, 255]; // This will normalize to [0.25, 0.5, 1.0]
+		let value_by_channel_normalized: Vec<f32> = value_by_channel
+			.iter()
+			.map(|&value| (value as f32 - 127.0) / 128.0)
+			.collect();
+		let frame_data = {
+			let mut data = vec![];
+			for i in 0..RECOGNIZER_INPUT_SIZE.y {
+				for _ in 0..RECOGNIZER_INPUT_SIZE.x {
+					data.extend(if i % 2 == 0 {
+						value_by_channel.clone()
+					} else {
+						zero_value_by_channel.clone()
+					});
+				}
+			}
+			data
+		};
+		let face_recognizer = get_face_recognizer();
+		let frame = get_frame(frame_data);
+		let normalized = face_recognizer.normalize_input(&frame);
+		let vector = normalized.to_data().to_vec::<f32>().expect("Failed");
+
+		for (index, value) in value_by_channel_normalized.iter().enumerate() {
+			let channel = &vector[FRAME_CHANNEL_SIZE * index..FRAME_CHANNEL_SIZE * (index + 1)];
+			check_channel(channel, *value);
+		}
+	}
+
+	#[test]
+	fn interprets_output() {
+		let device = get_device();
+		let test_cases = vec![
+			[0.0; EMBEDDING_LENGTH],
+			[128.0; EMBEDDING_LENGTH],
+			[255.0; EMBEDDING_LENGTH],
+		];
+
+		for test_case in test_cases {
+			let expected_result = FaceEmbedding::new(&test_case);
+			let output: Tensor<BurnBackend, 2> =
+				Tensor::from_data(TensorData::new(test_case.to_vec(), [512, 1]), &device);
+			let result = FaceRecognizer::interpret_output(&output).embedding;
+
+			assert_eq!(result, expected_result);
+		}
+	}
+}

crates/oblichey-cli/src/processors/frame_processor.rs

@@ -7,12 +7,18 @@ use crate::{
 	camera::Frame, geometry::Rectangle, models::recognizer::RECOGNIZER_INPUT_SIZE,
 	processors::face::FaceRecognitionError,
 };
+#[cfg(test)]
+use burn::backend::{ndarray::NdArrayDevice, NdArray};
+#[cfg(not(test))]
 use burn::backend::{wgpu::WgpuDevice, Wgpu};
 use image::imageops::{crop, resize, FilterType};
 use log::trace;
 use mockall_double::double;
 
+#[cfg(not(test))]
 pub type BurnBackend = Wgpu<f32, i32>;
+#[cfg(test)]
+pub type BurnBackend = NdArray<f32>;
 
 /// Checks whether a `Rectangle` is large enough to be passed into the recognizer model. We would
 /// not want to pass an upscaled image to it
@@ -56,7 +62,10 @@ pub struct FrameProcessor {
 
 impl FrameProcessor {
 	pub fn new() -> Self {
+		#[cfg(not(test))]
 		let device = WgpuDevice::default();
+		#[cfg(test)]
+		let device = NdArrayDevice::default();
 
 		Self {
 			detector: FaceDetector::new(&device),