feat: define streams in config

mosure · Mar 11, 2024 · d54291c · d54291c
1 parent 6c55522
commit d54291c
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 39 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "bevy_light_field"
 description = "rust bevy light field array tooling"
-version = "0.3.0"
+version = "0.4.0"
 edition = "2021"
 authors = ["mosure <mitchell@mosure.me>"]
 license = "MIT"
@@ -44,6 +44,7 @@ futures = "0.3"
 ndarray = { version = "0.15", optional = true }
 openh264 = "0.5"
 serde = "1.0"
+serde_json = "1.0"
 serde_qs = "0.12"
 retina = "0.4"
 tokio = { version = "1.36", features = ["full"] }

diff --git a/assets/streams.json b/assets/streams.json
@@ -0,0 +1,10 @@
+[
+    "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?"
+]
diff --git a/src/lib.rs b/src/lib.rs
@@ -8,10 +8,15 @@ pub mod mp4;
 pub mod stream;
 
 
-pub struct LightFieldPlugin;
+pub struct LightFieldPlugin {
+    pub stream_config: String,
+}
+
 impl Plugin for LightFieldPlugin {
     fn build(&self, app: &mut App) {
         app.add_plugins(materials::StreamMaterialsPlugin);
-        app.add_plugins(stream::RtspStreamPlugin);
+        app.add_plugins(stream::RtspStreamPlugin {
+            stream_config: self.stream_config.clone(),
+        });
     }
 }
diff --git a/src/matting.rs b/src/matting.rs
@@ -103,11 +103,13 @@ fn matting_inference(
 
     let thread_pool = AsyncComputeTaskPool::get();
 
-    let inputs = matted_streams.iter()
+    let (inputs, outputs): (Vec<_>, Vec<_>) = matted_streams.iter()
         .map(|(_, matted_stream)| {
-            images.get(matted_stream.input.clone()).unwrap()
+            let input = images.get(matted_stream.input.clone()).unwrap();
+            let output = (matted_stream.output.clone(), matted_stream.material.clone());
+            (input, output)
         })
-        .collect::<Vec<_>>();
+        .unzip();
 
     let uninitialized = inputs.iter().any(|image| image.size() == (32, 32).into());
     if uninitialized {
@@ -126,10 +128,6 @@ fn matting_inference(
     let onnx = onnx_assets.get(&modnet.onnx).unwrap();
     let session_arc = onnx.session.clone();
 
-    let outputs = matted_streams.iter()
-        .map(|(_, matted_stream)| (matted_stream.output.clone(), matted_stream.material.clone()))
-        .collect::<Vec<_>>();
-
     let task = thread_pool.spawn(async move {
         let mask_images: Result<Vec<Image>, String> = (|| {
             let session_lock = session_arc.lock().map_err(|e| e.to_string())?;
@@ -145,18 +143,18 @@ fn matting_inference(
         })();
 
         match mask_images {
-            Ok(mut mask_images) => {
+            Ok(mask_images) => {
                 let mut command_queue = CommandQueue::default();
 
                 command_queue.push(move |world: &mut World| {
-                    outputs.iter()
-                        .for_each(|(mask, material)| {
-                            let mut images = world.get_resource_mut::<Assets<Image>>().unwrap();
-                            images.insert(mask, mask_images.pop().unwrap());
-
-                            let mut foreground_materials = world.get_resource_mut::<Assets<ForegroundMaterial>>().unwrap();
-                            foreground_materials.get_mut(material).unwrap();
+                    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);
+                                foreground_materials.get_mut(&material).unwrap();
+                            });
                         });
+                    });
                 });
 
                 command_queue

diff --git a/src/stream.rs b/src/stream.rs
@@ -24,6 +24,7 @@ use retina::{
     },
     codec::VideoFrame,
 };
+use serde::{Deserialize, Serialize};
 use tokio::{
     fs::File,
     runtime::{
@@ -37,10 +38,17 @@ use url::Url;
 use crate::mp4::Mp4Writer;
 
 
-pub struct RtspStreamPlugin;
+pub struct RtspStreamPlugin {
+    pub stream_config: String,
+}
+
 impl Plugin for RtspStreamPlugin {
     fn build(&self, app: &mut App) {
+        let config = std::fs::File::open(&self.stream_config).unwrap();
+        let stream_uris = serde_json::from_reader::<_, StreamUris>(config).unwrap();
+
         app
+            .insert_resource(stream_uris)
             .init_resource::<RtspStreamManager>()
             .add_systems(Update, create_streams_from_descriptors)
             .add_systems(Update, apply_decode);
@@ -426,3 +434,8 @@ fn convert_h264(data: &mut [u8]) -> Result<(), Error> {
 
     Ok(())
 }
+
+
+
+#[derive(Resource, Clone, Debug, Default, Reflect, Serialize, Deserialize)]
+pub struct StreamUris(pub Vec<String>);
diff --git a/tools/viewer.rs b/tools/viewer.rs
@@ -29,7 +29,10 @@ use bevy_light_field::{
     LightFieldPlugin,
     materials::foreground::ForegroundMaterial,
     stream::{
-        RtspStreamDescriptor, RtspStreamManager, StreamId
+        RtspStreamDescriptor,
+        RtspStreamManager,
+        StreamId,
+        StreamUris,
     },
 };
 
@@ -40,20 +43,6 @@ use bevy_light_field::matting::{
 };
 
 
-// TODO: move to config file
-const RTSP_URIS: [&str; 8] = [
-    "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?",
-];
-
-
 #[derive(
     Default,
     Debug,
@@ -64,6 +53,9 @@ const RTSP_URIS: [&str; 8] = [
 )]
 #[command(about = "bevy_light_field viewer", version)]
 pub struct LightFieldViewer {
+    #[arg(long, default_value = "assets/streams.json")]
+    pub config: String,
+
     #[arg(long, default_value = "false")]
     pub show_fps: bool,
 
@@ -113,7 +105,9 @@ fn main() {
                     primary_window,
                     ..default()
                 }),
-            LightFieldPlugin,
+            LightFieldPlugin {
+                stream_config: args.config.clone(),
+            },
 
             #[cfg(feature = "person_matting")]
             MattingPlugin::new((
@@ -148,9 +142,10 @@ fn create_streams(
     primary_window: Query<&Window, With<PrimaryWindow>>,
     args: Res<LightFieldViewer>,
     mut foreground_materials: ResMut<Assets<ForegroundMaterial>>,
+    stream_uris: Res<StreamUris>,
 ) {
     let window = primary_window.single();
-    let elements = RTSP_URIS.len();
+    let elements = stream_uris.0.len();
 
     let (
         columns,
@@ -169,15 +164,15 @@ fn create_streams(
         ..default()
     };
 
-    let input_images: Vec<Handle<Image>> = RTSP_URIS.iter()
+    let input_images: Vec<Handle<Image>> = stream_uris.0.iter()
         .enumerate()
-        .map(|(index, &url)| {
+        .map(|(index, url)| {
             let entity = commands.spawn_empty().id();
 
             let mut image = Image {
                 asset_usage: RenderAssetUsages::all(),
                 texture_descriptor: TextureDescriptor {
-                    label: Some(url),
+                    label: None,
                     size,
                     dimension: TextureDimension::D2,
                     format: TextureFormat::Rgba8UnormSrgb,