From 35e95f7d00e1169ee3db5cc805e79d73aa805c58 Mon Sep 17 00:00:00 2001
From: arongeo <me@arongeo.com>
Date: Tue, 25 Apr 2023 21:05:01 +0200
Subject: [PATCH] Added face detection

---
 src/state.rs |  1 +
 src/video.rs | 63 ++++++++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/src/state.rs b/src/state.rs
index d1555d6..6572c75 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -46,6 +46,7 @@ impl TelloState {
 
             let str_num = String::from_iter(num);
 
+            // Very ugly parsing solution, pretty sure there's something better out there
             match i {
                 0 => tellostate.mid = str_num.parse().unwrap(),
                 (1..=3) => tellostate.xyz[i-1] = str_num.parse().unwrap(),
diff --git a/src/video.rs b/src/video.rs
index b14e198..d9a84e1 100644
--- a/src/video.rs
+++ b/src/video.rs
@@ -2,8 +2,10 @@
 use opencv as cv;
 use cv::{
     highgui,
-    objdetect::CascadeClassifier,
-    types::VectorOfRect,
+    objdetect::{CascadeClassifier, self},
+    types::VectorOfRect, prelude::CascadeClassifierTrait,
+    core::Size,
+    imgproc::rectangle,
 };
 
 mod decoding;
@@ -24,6 +26,14 @@ pub fn spawn_video_thread(vrx: std::sync::mpsc::Receiver<crate::ThreadMsg>) -> s
 
         let mut decoder = decoding::Decoder::new().unwrap();
 
+        let mut face_detector = match CascadeClassifier::new("./cascades/haarcascade_frontalface_alt.xml") {
+            Ok(cc) => cc,
+            Err(e) => panic!("Failed to create face detector: {:?}", e),
+        };
+
+        let mut faces = VectorOfRect::new();
+        let mut find_face_in = 3;
+
         loop {
             let mut video_buffer = [0; 2048];
             let msg_len = match video_socket.recv(&mut video_buffer) {
@@ -32,23 +42,18 @@ pub fn spawn_video_thread(vrx: std::sync::mpsc::Receiver<crate::ThreadMsg>) -> s
             };
 
             frame_packet.extend_from_slice(&video_buffer[..msg_len]);
-
+            
             match decoding::check_for_valid_packet(&frame_packet) {
                 Some(frame_borders) => {
                     match decoding::h264_decode(&frame_packet[(frame_borders.0)..(frame_borders.1)], &mut decoder) {
                         Ok(mut result) => {
                             frame_packet = frame_packet[(frame_borders.1)..(frame_packet.len())].to_vec();
                             
-                            let bgra_frame = match conversions::vec_to_bgra_mat(&mut result.0, result.1) {
+                            let mut bgra_frame = match conversions::vec_to_bgra_mat(&mut result.0, result.1) {
                                 Ok(m) => m,
                                 Err(()) => continue,
                             };
 
-                            match highgui::imshow("tello", &bgra_frame) {
-                                Ok(_) => {},
-                                Err(e) => println!("Error: {:?}", e),
-                            };
-
                             let gray_buf = match conversions::mat_to_grayscale(&bgra_frame, result.1) {
                                 Ok(m) => m,
                                 Err(_) => continue,
@@ -59,6 +64,46 @@ pub fn spawn_video_thread(vrx: std::sync::mpsc::Receiver<crate::ThreadMsg>) -> s
                                 Err(e) => println!("Error: {:?}", e),
                             };
 
+                            find_face_in -= 1;
+
+                            // We need to limit how much the classifier is ran, for performance
+                            // reasons.
+                            if find_face_in == 0 {
+                                faces = VectorOfRect::new();
+                                match face_detector.detect_multi_scale(
+                                    &gray_buf,
+                                    &mut faces,
+                                    1.1,
+                                    2,
+                                    objdetect::CASCADE_SCALE_IMAGE,
+                                    Size::new(30, 30),
+                                    Size::new(0, 0)
+                                ) {
+                                    Ok(_) => {},
+                                    Err(e) => println!("Error with face detector: {}", e),
+                                };
+
+                                find_face_in = 3;
+                            }
+
+                            // We can still display the last frames rectangle though, since it's
+                            // quite impossible to notice changes in that little amount of time.
+                            for face in &faces {
+                                rectangle(
+                                    &mut bgra_frame,
+                                    face,
+                                    cv::core::Scalar::new(255.0, 0.0, 0.0, 0.0),
+                                    2,
+                                    cv::imgproc::LINE_8,
+                                    0
+                                ).unwrap();
+                            }
+
+                            match highgui::imshow("tello", &bgra_frame) {
+                                Ok(_) => {},
+                                Err(e) => println!("Error: {:?}", e),
+                            };
+
                             highgui::poll_key();
 
                             match vrx.try_recv() {