feat: automatic recording

Cargo.toml

@@ -39,7 +39,7 @@ pipeline = ["image", "rayon"]
 anyhow = "1.0"
 async-compat = "0.2"
 bevy_args = "1.3"
-bevy_ort = { version = "0.6", optional = true }
+bevy_ort = { version = "0.6.3", optional = true }
 bytes = "1.5"
 clap = { version = "4.4", features = ["derive"] }
 futures = "0.3"

src/lib.rs

@@ -17,6 +17,7 @@ pub struct LightFieldPlugin {
 impl Plugin for LightFieldPlugin {
     fn build(&self, app: &mut App) {
         app.add_plugins(materials::StreamMaterialsPlugin);
+        app.add_plugins(person_detect::PersonDetectPlugin);
         app.add_plugins(stream::RtspStreamPlugin {
             stream_config: self.stream_config.clone(),
         });

src/matting.rs

@@ -149,8 +149,8 @@ fn matting_inference(
                 command_queue.push(move |world: &mut World| {
                     world.resource_scope(|world, mut images: Mut<Assets<Image>>| {
                         world.resource_scope(|_world, mut foreground_materials: Mut<Assets<ForegroundMaterial>>| {
-                            outputs.into_iter().zip(mask_images).for_each(|((output, material), mask_image)| {
-                                images.insert(output, mask_image);
+                            outputs.into_iter().zip(mask_images).for_each(|((mask, material), mask_image)| {
+                                images.insert(mask, mask_image);
                                 foreground_materials.get_mut(&material).unwrap();
                             });
                         });

src/person_detect.rs

@@ -1,7 +1,7 @@
 use std::cmp::{max, min};
 
 use bevy::prelude::*;
-use image::DynamicImage;
+use image::{DynamicImage, GenericImageView, imageops::FilterType, ImageBuffer, Luma, RgbImage};
 use rayon::prelude::*;
 
 use crate::{
@@ -14,9 +14,9 @@ pub struct PersonDetectPlugin;
 
 impl Plugin for PersonDetectPlugin {
     fn build(&self, app: &mut App) {
+        app.add_event::<PersonDetectedEvent>();
         app.add_systems(Update, detect_person);
     }
-
 }
 
 
@@ -54,15 +54,20 @@ fn detect_person(
             AssetEvent::Modified { id } => {
                 for (matted_stream, _) in person_detect_streams.iter() {
                     if &matted_stream.output.id() == id {
-                        let image = images.get(&matted_stream.output).unwrap().clone().try_into_dynamic().unwrap();
+                        let image = images.get(&matted_stream.output).unwrap();
+
+                        let buffer = ImageBuffer::<Luma<u8>, Vec<u8>>::from_raw(
+                            image.width(),
+                            image.height(),
+                            image.data.clone(),
+                        ).unwrap();
 
-                        let bounding_box = masked_bounding_box(&image);
-                        let sum = sum_masked_pixels(&image);
+                        let bounding_box = masked_bounding_box(&buffer);
+                        let sum = sum_masked_pixels(&buffer);
 
-                        println!("bounding box: {:?}, sum: {}", bounding_box, sum);
+                        let masked_ratio = sum / (buffer.width() * buffer.height()) as f32;
+                        let person_detected = masked_ratio > 0.14;
 
-                        // TODO: add thresholds for detection
-                        let person_detected = false;
                         if person_detected {
                             ev_person_detected.send(PersonDetectedEvent {
                                 stream_id: matted_stream.stream_id,
@@ -80,19 +85,16 @@ fn detect_person(
 
 
 
-pub fn masked_bounding_box(image: &DynamicImage) -> Option<BoundingBox> {
-    let img = image.as_luma8().unwrap();
-
-    let bounding_boxes = img.enumerate_pixels()
-        .par_bridge()
+pub fn masked_bounding_box(buffer: &ImageBuffer<Luma<u8>, Vec<u8>>) -> Option<BoundingBox> {
+    let bounding_boxes = buffer.enumerate_pixels()
         .filter_map(|(x, y, pixel)| {
-            if pixel[0] > 128 {
+            if pixel.0[0] > 250 {
                 Some((x as i32, y as i32, x as i32, y as i32))
             } else {
                 None
             }
         })
-        .reduce_with(|(
+        .reduce(|(
             min_x1,
             min_y1,
             max_x1,
@@ -127,17 +129,12 @@ pub fn masked_bounding_box(image: &DynamicImage) -> Option<BoundingBox> {
 }
 
 
-pub fn sum_masked_pixels(image: &DynamicImage) -> f32 {
-    let img = image.as_luma8().unwrap();
-    let pixels = img.pixels();
-
-    let count = pixels.par_bridge()
+pub fn sum_masked_pixels(image: &ImageBuffer<Luma<u8>, Vec<u8>>) -> f32 {
+    image.pixels()
         .map(|pixel| {
             pixel.0[0] as f32 / 255.0
         })
-        .sum();
-
-    count
+        .sum()
 }
 
 
@@ -161,8 +158,7 @@ mod tests {
             }
         }
 
-        let dynamic_img = DynamicImage::ImageLuma8(img);
-        let result = masked_bounding_box(&dynamic_img).expect("expected a bounding box");
+        let result = masked_bounding_box(&img).expect("expected a bounding box");
 
         let expected = BoundingBox {
             x:2,
@@ -184,8 +180,7 @@ mod tests {
         img.put_pixel(1, 0, Luma([127]));
         img.put_pixel(2, 0, Luma([63]));
 
-        let dynamic_img = DynamicImage::ImageLuma8(img);
-        let result = sum_masked_pixels(&dynamic_img);
+        let result = sum_masked_pixels(&img);
 
         let expected = (255.0 + 127.0 + 63.0) / 255.0;
         assert_relative_eq!(result, expected);

src/pipeline.rs

@@ -58,7 +58,7 @@ pub struct Session {
 impl Session {
     pub fn new(directory: String) -> Self {
         let id = get_next_session_id(&directory);
-        let directory = format!("{}/{}", directory, id);
+        let directory = format!("{}/{}/raw", directory, id);
         std::fs::create_dir_all(&directory).unwrap();
 
         Self { id, directory }

src/stream.rs

@@ -28,15 +28,15 @@ use retina::{
 use serde::{Deserialize, Serialize};
 use tokio::{
     fs::File,
-    runtime::{
-        Handle,
-        Runtime,
-    },
+    runtime::Handle,
     sync::mpsc,
 };
 use url::Url;
 
-use crate::mp4::Mp4Writer;
+use crate::{
+    mp4::Mp4Writer,
+    pipeline::Session as PipelineSession,
+};
 
 
 pub struct RtspStreamPlugin {
@@ -192,7 +192,11 @@ pub struct RtspStreamManager {
 
 impl FromWorld for RtspStreamManager {
     fn from_world(_world: &mut World) -> Self {
-        let runtime = Runtime::new().unwrap();
+
+        // TODO: upgrade to [bevy-tokio-tasks](https://github.com/EkardNT/bevy-tokio-tasks) to share tokio runtime between rtsp and inference - waiting on: https://github.com/pykeio/ort/pull/174
+        let mut runtime = tokio::runtime::Builder::new_multi_thread();
+        runtime.enable_all();
+        let runtime = runtime.build().unwrap();
         let handle = runtime.handle().clone();
 
         std::thread::spawn(move || {
@@ -227,18 +231,24 @@ impl RtspStreamManager {
         });
     }
 
-    pub fn start_recording(&self, output_directory: &str) {
+    pub fn start_recording(&self, session: &PipelineSession) {
         let stream_handles = self.stream_handles.lock().unwrap();
         for descriptor in stream_handles.iter() {
             let filename = format!("{}.mp4", descriptor.id.0);
-            let filepath = format!("{}/{}", output_directory, filename);
+            let filepath = format!("{}/{}", session.directory, filename);
 
             let send_channel = descriptor.recording_sender.lock().unwrap();
+
+            if send_channel.is_none() {
+                println!("no recording sender for stream {}", descriptor.id.0);
+                continue;
+            }
+
             let sender_clone = send_channel.as_ref().unwrap().clone();
 
             self.handle.block_on(async move {
                 let file = File::create(&filepath).await.unwrap();
-                sender_clone.send(RecordingCommand::StartRecording(file)).await.unwrap();
+                let _ = sender_clone.send(RecordingCommand::StartRecording(file)).await;
             });
         }
     }
@@ -249,10 +259,16 @@ impl RtspStreamManager {
         let stream_handles = self.stream_handles.lock().unwrap();
         for descriptor in stream_handles.iter() {
             let send_channel = descriptor.recording_sender.lock().unwrap();
+
+            if send_channel.is_none() {
+                println!("no recording sender for stream {}", descriptor.id.0);
+                continue;
+            }
+
             let sender_clone = send_channel.as_ref().unwrap().clone();
 
             self.handle.block_on(async move {
-                sender_clone.send(RecordingCommand::StopRecording).await.unwrap();
+                let _ = sender_clone.send(RecordingCommand::StopRecording).await;
             });
 
             filepaths.push(format!("{}.mp4", descriptor.id.0));
@@ -364,6 +380,8 @@ impl RtspStream {
                                     data,
                                 };
 
+                                // TODO: write streams into a frame texture array (stream, channel, width, height)
+
                                 *locked_sink = Some(bgra);
                             },
                         }