frame_processor.rs

impl FrameProcessor {
	pub fn new() -> Self {
		let device = NdArrayDevice::default();
		let detector: detector::Model<Backend> = detector::Model::default();
		let recognizer: recognizer::Model<Backend> = recognizer::Model::default();

		Self {
			device,
			detector,
			recognizer,
		}
	}

	pub fn process_frame(&self, frame: &RgbImage) -> Vec<DetectedFace> {
		assert_eq!(
			frame.width(),
			DETECTOR_INPUT_SIZE.x,
			"Image width does not match network requirements!"
		);
		assert_eq!(
			frame.height(),
			DETECTOR_INPUT_SIZE.y,
			"Image height does not match network requirements!"
		);

		let detector_input = self.normalize_detector_input(frame);
		let detector_output = self.detector.forward(detector_input);
		let face_rectangles = self.interpret_detector_output(detector_output);

		let mut frame = frame.clone();
		for face_rectangle in &face_rectangles {
			let recognizer_input = self.normalize_recognizer_input(&mut frame, face_rectangle);
			let recognizer_output = self.recognizer.forward(recognizer_input);
			println!("{}", recognizer_output);
		}

		println!("Len: {}", face_rectangles.len());
		// For now we fill the `face` field with `default()`
		face_rectangles
			.into_iter()
			.map(|rectangle| DetectedFace {
				face: Face::default(),
				rectangle,
			})
			.collect()
	}

	fn normalize_recognizer_input(
		&self,
		frame: &mut Frame,
		face_rectangle: &Rectangle<u32>,
	) -> Tensor<Backend, 4> {
		let cropped = crop(
			frame,
			face_rectangle.min.x,
			face_rectangle.min.y,
			face_rectangle.max.x + face_rectangle.min.x,
			face_rectangle.max.y + face_rectangle.min.y,
		);

		let resized = resize(
			&cropped.to_image(),
			RECOGNIZER_INPUT_SIZE.x,
			RECOGNIZER_INPUT_SIZE.y,
			FilterType::CatmullRom,
		);

		// Shape of the image: height, width, channels
		let shape = [
			RECOGNIZER_INPUT_SIZE.y as usize,
			RECOGNIZER_INPUT_SIZE.x as usize,
			3 as usize,
		];

		// Make into a tensor
		let tensor = Tensor::from_data(TensorData::new(resized.to_vec(), shape), &self.device);

		// Normalize between [-1, 1]
		let normalized = (tensor - Tensor::full(shape, 127, &self.device)) / 128;

		// Reorder dimension to have: channels, height, width
		let permutated = normalized.permute([2, 0, 1]);

		// Make the tensor the correct shape: batch, channels, height, width
		let unsqueezed = permutated.unsqueeze::<4>();

		unsqueezed
	}
}