feat: recording and post-process automation

Cargo.toml

@@ -28,9 +28,11 @@ default-run = "viewer"
 [features]
 default = [
   "person_matting",
+  "pipeline",
 ]
 
 person_matting = ["bevy_ort", "ort", "ndarray"]
+pipeline = ["image", "rayon"]
 
 
 [dependencies]
@@ -41,8 +43,10 @@ bevy_ort = { version = "0.6", optional = true }
 bytes = "1.5"
 clap = { version = "4.4", features = ["derive"] }
 futures = "0.3"
+image = { version = "0.24", optional = true }
 ndarray = { version = "0.15", optional = true }
 openh264 = "0.5"
+rayon = { version = "1.8", optional = true }
 serde = "1.0"
 serde_json = "1.0"
 serde_qs = "0.12"
@@ -76,6 +80,10 @@ features = [
 ]
 
 
+[dev-dependencies]
+approx = "0.5"
+
+
 [profile.dev.package."*"]
 opt-level = 3
 

README.md

@@ -54,7 +54,7 @@ use bevy::{
 };
 
 use bevy_light_field::stream::{
-    RtspStreamDescriptor,
+    RtspStreamHandle,
     RtspStreamPlugin,
     StreamId,
 };
@@ -118,7 +118,7 @@ fn create_streams(
                 ..default()
             });
 
-            let rtsp_stream = RtspStreamDescriptor::new(
+            let rtsp_stream = RtspStreamHandle::new(
                 url.to_string(),
                 StreamId(index),
                 image,

assets/streams.json

@@ -1,19 +1,20 @@
 [
-    "rtsp://192.168.1.22/stream/main",
+    { "uri": "rtsp://192.168.1.21/stream/main", "transport": "Udp" },
+    { "uri": "rtsp://192.168.1.22/stream/main", "transport": "Udp" },
 
-    "rtsp://192.168.1.23/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.24/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.25/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.26/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.27/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.28/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.29/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.30/user=admin&password=admin123&channel=1&stream=0.sdp?",
+    { "uri": "rtsp://192.168.1.23/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.24/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.25/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.26/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.27/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.28/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.29/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.30/user=admin&password=admin123&channel=1&stream=0.sdp?" },
 
-    "rtsp://192.168.1.31/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.32/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.33/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.34/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.35/user=admin&password=admin123&channel=1&stream=0.sdp?",
-    "rtsp://192.168.1.36/user=admin&password=admin123&channel=1&stream=0.sdp?"
+    { "uri": "rtsp://192.168.1.31/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.32/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.33/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.34/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.35/user=admin&password=admin123&channel=1&stream=0.sdp?" },
+    { "uri": "rtsp://192.168.1.36/user=admin&password=admin123&channel=1&stream=0.sdp?" }
 ]

src/lib.rs

@@ -5,6 +5,8 @@ pub mod matting;
 
 pub mod materials;
 pub mod mp4;
+pub mod person_detect;
+pub mod pipeline;
 pub mod stream;
 
 

src/person_detect.rs

@@ -0,0 +1,193 @@
+use std::cmp::{max, min};
+
+use bevy::prelude::*;
+use image::DynamicImage;
+use rayon::prelude::*;
+
+use crate::{
+    matting::MattedStream,
+    stream::StreamId,
+};
+
+
+pub struct PersonDetectPlugin;
+
+impl Plugin for PersonDetectPlugin {
+    fn build(&self, app: &mut App) {
+        app.add_systems(Update, detect_person);
+    }
+
+}
+
+
+#[derive(Component)]
+pub struct DetectPersons;
+
+
+#[derive(Debug, Clone, Reflect, PartialEq)]
+pub struct BoundingBox {
+    pub x: i32,
+    pub y: i32,
+    pub width: i32,
+    pub height: i32,
+}
+
+#[derive(Event, Debug, Reflect, Clone)]
+pub struct PersonDetectedEvent {
+    pub stream_id: StreamId,
+    pub bounding_box: BoundingBox,
+    pub mask_sum: f32,
+}
+
+
+fn detect_person(
+    mut ev_asset: EventReader<AssetEvent<Image>>,
+    mut ev_person_detected: EventWriter<PersonDetectedEvent>,
+    person_detect_streams: Query<(
+        &MattedStream,
+        &DetectPersons,
+    )>,
+    images: Res<Assets<Image>>,
+) {
+    for ev in ev_asset.read() {
+        match ev {
+            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 bounding_box = masked_bounding_box(&image);
+                        let sum = sum_masked_pixels(&image);
+
+                        println!("bounding box: {:?}, sum: {}", bounding_box, sum);
+
+                        // TODO: add thresholds for detection
+                        let person_detected = false;
+                        if person_detected {
+                            ev_person_detected.send(PersonDetectedEvent {
+                                stream_id: matted_stream.stream_id,
+                                bounding_box: bounding_box.unwrap(),
+                                mask_sum: sum,
+                            });
+                        }
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+}
+
+
+
+pub fn masked_bounding_box(image: &DynamicImage) -> Option<BoundingBox> {
+    let img = image.as_luma8().unwrap();
+
+    let bounding_boxes = img.enumerate_pixels()
+        .par_bridge()
+        .filter_map(|(x, y, pixel)| {
+            if pixel[0] > 128 {
+                Some((x as i32, y as i32, x as i32, y as i32))
+            } else {
+                None
+            }
+        })
+        .reduce_with(|(
+            min_x1,
+            min_y1,
+            max_x1,
+            max_y1,
+        ), (
+            min_x2,
+            min_y2,
+            max_x2,
+            max_y2,
+        )| {
+            (
+                min(min_x1, min_x2),
+                min(min_y1, min_y2),
+                max(max_x1, max_x2),
+                max(max_y1, max_y2),
+            )
+        });
+
+    bounding_boxes.map(|(
+        min_x,
+        min_y,
+        max_x,
+        max_y
+    )| {
+        BoundingBox {
+            x: min_x,
+            y: min_y,
+            width: max_x - min_x + 1,
+            height: max_y - min_y + 1,
+        }
+    })
+}
+
+
+pub fn sum_masked_pixels(image: &DynamicImage) -> f32 {
+    let img = image.as_luma8().unwrap();
+    let pixels = img.pixels();
+
+    let count = pixels.par_bridge()
+        .map(|pixel| {
+            pixel.0[0] as f32 / 255.0
+        })
+        .sum();
+
+    count
+}
+
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use image::{ImageBuffer, Luma};
+    use approx::assert_relative_eq;
+
+
+    #[test]
+    fn test_masked_bounding_box() {
+        let width = 10;
+        let height = 10;
+        let mut img: ImageBuffer<Luma<u8>, Vec<u8>> = ImageBuffer::new(width, height);
+
+        for x in 2..=5 {
+            for y in 2..=5 {
+                img.put_pixel(x, y, Luma([200]));
+            }
+        }
+
+        let dynamic_img = DynamicImage::ImageLuma8(img);
+        let result = masked_bounding_box(&dynamic_img).expect("expected a bounding box");
+
+        let expected = BoundingBox {
+            x:2,
+            y: 2,
+            width: 4,
+            height: 4,
+        };
+        assert_eq!(result, expected, "the computed bounding box did not match the expected values.");
+    }
+
+
+    #[test]
+    fn test_sum_masked_pixels() {
+        let width = 4;
+        let height = 4;
+        let mut img: ImageBuffer<Luma<u8>, Vec<u8>> = ImageBuffer::new(width, height);
+
+        img.put_pixel(0, 0, Luma([255]));
+        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 expected = (255.0 + 127.0 + 63.0) / 255.0;
+        assert_relative_eq!(result, expected);
+    }
+}