refactor: figured out issue

cs50victor · Jan 25, 2024 · c593f05 · c593f05
1 parent f99dd55
commit c593f05
Show file tree

Hide file tree

Showing 9 changed files with 178 additions and 208 deletions.
diff --git a/.env.example b/.env.example
@@ -1,4 +1 @@
-NEXT_PUBLIC_LK_TOKEN_ENDPOINT=""
-
-OPENAI_API_KEY=""
-OPENAI_ORG_ID=""
+# add later
diff --git a/crates/bevy_frame_capture/src/image_copy.rs b/crates/bevy_frame_capture/src/image_copy.rs
@@ -177,186 +177,179 @@
 //     }
 // }
 
-
-
-    use std::sync::Arc;
-
-    use bevy::prelude::*;
-    use bevy::render::render_asset::RenderAssets;
-    use bevy::render::render_graph::{self, NodeRunError, RenderGraph, RenderGraphContext};
-    use bevy::render::renderer::{RenderContext, RenderDevice, RenderQueue};
-    use bevy::render::{Extract, RenderApp};
-
-    use bevy::render::render_resource::{
-        Buffer, BufferDescriptor, BufferUsages, CommandEncoderDescriptor, Extent3d,
-        ImageCopyBuffer, ImageDataLayout, MapMode,
-    };
-    use pollster::FutureExt;
-    use wgpu::Maintain;
-
-    use std::sync::atomic::{AtomicBool, Ordering};
-
-    pub fn receive_images(
-        image_copiers: Query<&ImageCopier>,
-        mut images: ResMut<Assets<Image>>,
-        render_device: Res<RenderDevice>,
-    ) {
-        for image_copier in image_copiers.iter() {
-            if !image_copier.enabled() {
-                continue;
-            }
-            // Derived from: https://sotrh.github.io/learn-wgpu/showcase/windowless/#a-triangle-without-a-window
-            // We need to scope the mapping variables so that we can
-            // unmap the buffer
-            async {
-                let buffer_slice = image_copier.buffer.slice(..);
-
-                // NOTE: We have to create the mapping THEN device.poll() before await
-                // the future. Otherwise the application will freeze.
-                let (tx, rx) = futures_intrusive::channel::shared::oneshot_channel();
-                buffer_slice.map_async(MapMode::Read, move |result| {
-                    tx.send(result).unwrap();
-                });
-                render_device.poll(Maintain::Wait);
-                rx.receive().await.unwrap().unwrap();
-                if let Some(image) = images.get_mut(&image_copier.dst_image) {
-                    image.data = buffer_slice.get_mapped_range().to_vec();
-                }
-
-                image_copier.buffer.unmap();
+use std::sync::Arc;
+
+use bevy::{
+    prelude::*,
+    render::{
+        render_asset::RenderAssets,
+        render_graph::{self, NodeRunError, RenderGraph, RenderGraphContext},
+        renderer::{RenderContext, RenderDevice, RenderQueue},
+        Extract, RenderApp,
+    },
+};
+
+use bevy::render::render_resource::{
+    Buffer, BufferDescriptor, BufferUsages, CommandEncoderDescriptor, Extent3d, ImageCopyBuffer,
+    ImageDataLayout, MapMode,
+};
+use pollster::FutureExt;
+use wgpu::Maintain;
+
+use std::sync::atomic::{AtomicBool, Ordering};
+
+pub fn receive_images(
+    image_copiers: Query<&ImageCopier>,
+    mut images: ResMut<Assets<Image>>,
+    render_device: Res<RenderDevice>,
+) {
+    for image_copier in image_copiers.iter() {
+        if !image_copier.enabled() {
+            continue;
+        }
+        // Derived from: https://sotrh.github.io/learn-wgpu/showcase/windowless/#a-triangle-without-a-window
+        // We need to scope the mapping variables so that we can
+        // unmap the buffer
+        async {
+            let buffer_slice = image_copier.buffer.slice(..);
+
+            // NOTE: We have to create the mapping THEN device.poll() before await
+            // the future. Otherwise the application will freeze.
+            let (tx, rx) = futures_intrusive::channel::shared::oneshot_channel();
+            buffer_slice.map_async(MapMode::Read, move |result| {
+                tx.send(result).unwrap();
+            });
+            render_device.poll(Maintain::Wait);
+            rx.receive().await.unwrap().unwrap();
+            if let Some(image) = images.get_mut(&image_copier.dst_image) {
+                image.data = buffer_slice.get_mapped_range().to_vec();
             }
-            .block_on();
+
+            image_copier.buffer.unmap();
         }
+        .block_on();
     }
+}
 
-    pub const IMAGE_COPY: &str = "image_copy";
+pub const IMAGE_COPY: &str = "image_copy";
 
-    pub struct ImageCopyPlugin;
-    impl Plugin for ImageCopyPlugin {
-        fn build(&self, app: &mut App) {
-            let render_app = app
-                .add_systems(Update, receive_images)
-                .sub_app_mut(RenderApp);
+pub struct ImageCopyPlugin;
+impl Plugin for ImageCopyPlugin {
+    fn build(&self, app: &mut App) {
+        let render_app = app.add_systems(Update, receive_images).sub_app_mut(RenderApp);
 
-            render_app.add_systems(ExtractSchedule, image_copy_extract);
+        render_app.add_systems(ExtractSchedule, image_copy_extract);
 
-            let mut graph = render_app.world.get_resource_mut::<RenderGraph>().unwrap();
+        let mut graph = render_app.world.get_resource_mut::<RenderGraph>().unwrap();
 
-            graph.add_node(IMAGE_COPY, ImageCopyDriver);
+        graph.add_node(IMAGE_COPY, ImageCopyDriver);
 
-            graph.add_node_edge(IMAGE_COPY, bevy::render::main_graph::node::CAMERA_DRIVER);
-        }
+        graph.add_node_edge(IMAGE_COPY, bevy::render::main_graph::node::CAMERA_DRIVER);
     }
+}
+
+#[derive(Clone, Default, Resource, Deref, DerefMut)]
+pub struct ImageCopiers(pub Vec<ImageCopier>);
 
-    #[derive(Clone, Default, Resource, Deref, DerefMut)]
-    pub struct ImageCopiers(pub Vec<ImageCopier>);
+#[derive(Clone, Component)]
+pub struct ImageCopier {
+    buffer: Buffer,
+    enabled: Arc<AtomicBool>,
+    src_image: Handle<Image>,
+    dst_image: Handle<Image>,
+}
 
-    #[derive(Clone, Component)]
-    pub struct ImageCopier {
-        buffer: Buffer,
-        enabled: Arc<AtomicBool>,
+impl ImageCopier {
+    pub fn new(
         src_image: Handle<Image>,
         dst_image: Handle<Image>,
-    }
-
-    impl ImageCopier {
-        pub fn new(
-            src_image: Handle<Image>,
-            dst_image: Handle<Image>,
-            size: Extent3d,
-            render_device: &RenderDevice,
-        ) -> ImageCopier {
-            let padded_bytes_per_row =
-                RenderDevice::align_copy_bytes_per_row((size.width) as usize) * 4;
-
-            let cpu_buffer = render_device.create_buffer(&BufferDescriptor {
-                label: None,
-                size: padded_bytes_per_row as u64 * size.height as u64,
-                usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
-                mapped_at_creation: false,
-            });
-
-            ImageCopier {
-                buffer: cpu_buffer,
-                src_image,
-                dst_image,
-                enabled: Arc::new(AtomicBool::new(true)),
-            }
-        }
-
-        pub fn enabled(&self) -> bool {
-            self.enabled.load(Ordering::Relaxed)
+        size: Extent3d,
+        render_device: &RenderDevice,
+    ) -> ImageCopier {
+        let padded_bytes_per_row =
+            RenderDevice::align_copy_bytes_per_row((size.width) as usize) * 4;
+
+        let cpu_buffer = render_device.create_buffer(&BufferDescriptor {
+            label: None,
+            size: padded_bytes_per_row as u64 * size.height as u64,
+            usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
+
+        ImageCopier {
+            buffer: cpu_buffer,
+            src_image,
+            dst_image,
+            enabled: Arc::new(AtomicBool::new(true)),
         }
     }
 
-    pub fn image_copy_extract(
-        mut commands: Commands,
-        image_copiers: Extract<Query<&ImageCopier>>,
-    ) {
-        commands.insert_resource(ImageCopiers(
-            image_copiers.iter().cloned().collect::<Vec<ImageCopier>>(),
-        ));
+    pub fn enabled(&self) -> bool {
+        self.enabled.load(Ordering::Relaxed)
     }
+}
+
+pub fn image_copy_extract(mut commands: Commands, image_copiers: Extract<Query<&ImageCopier>>) {
+    commands
+        .insert_resource(ImageCopiers(image_copiers.iter().cloned().collect::<Vec<ImageCopier>>()));
+}
+
+#[derive(Default)]
+pub struct ImageCopyDriver;
+
+impl render_graph::Node for ImageCopyDriver {
+    fn run(
+        &self,
+        _graph: &mut RenderGraphContext,
+        render_context: &mut RenderContext,
+        world: &World,
+    ) -> Result<(), NodeRunError> {
+        let image_copiers = world.get_resource::<ImageCopiers>().unwrap();
+        let gpu_images = world.get_resource::<RenderAssets<Image>>().unwrap();
 
-    #[derive(Default)]
-    pub struct ImageCopyDriver;
-
-    impl render_graph::Node for ImageCopyDriver {
-        fn run(
-            &self,
-            _graph: &mut RenderGraphContext,
-            render_context: &mut RenderContext,
-            world: &World,
-        ) -> Result<(), NodeRunError> {
-            let image_copiers = world.get_resource::<ImageCopiers>().unwrap();
-            let gpu_images = world.get_resource::<RenderAssets<Image>>().unwrap();
-
-            for image_copier in image_copiers.iter() {
-                if !image_copier.enabled() {
-                    continue;
-                }
-
-                let src_image = gpu_images.get(&image_copier.src_image).unwrap();
-
-                let mut encoder = render_context
-                    .render_device()
-                    .create_command_encoder(&CommandEncoderDescriptor::default());
-
-                let block_dimensions = src_image.texture_format.block_dimensions();
-                let block_size = src_image.texture_format.block_size(None).unwrap();
-
-                let padded_bytes_per_row = RenderDevice::align_copy_bytes_per_row(
-                    (src_image.size.x as usize / block_dimensions.0 as usize)
-                        * block_size as usize,
-                );
-
-                let texture_extent = Extent3d {
-                    width: src_image.size.x as u32,
-                    height: src_image.size.y as u32,
-                    depth_or_array_layers: 1,
-                };
-
-                encoder.copy_texture_to_buffer(
-                    src_image.texture.as_image_copy(),
-                    ImageCopyBuffer {
-                        buffer: &image_copier.buffer,
-                        layout: ImageDataLayout {
-                            offset: 0,
-                            bytes_per_row: Some(
-                                std::num::NonZeroU32::new(padded_bytes_per_row as u32)
-                                    .unwrap()
-                                    .into(),
-                            ),
-                            rows_per_image: None,
-                        },
-                    },
-                    texture_extent,
-                );
-
-                let render_queue = world.get_resource::<RenderQueue>().unwrap();
-                render_queue.submit(std::iter::once(encoder.finish()));
+        for image_copier in image_copiers.iter() {
+            if !image_copier.enabled() {
+                continue;
             }
 
-            Ok(())
+            let src_image = gpu_images.get(&image_copier.src_image).unwrap();
+
+            let mut encoder = render_context
+                .render_device()
+                .create_command_encoder(&CommandEncoderDescriptor::default());
+
+            let block_dimensions = src_image.texture_format.block_dimensions();
+            let block_size = src_image.texture_format.block_size(None).unwrap();
+
+            let padded_bytes_per_row = RenderDevice::align_copy_bytes_per_row(
+                (src_image.size.x as usize / block_dimensions.0 as usize) * block_size as usize,
+            );
+
+            let texture_extent = Extent3d {
+                width: src_image.size.x as u32,
+                height: src_image.size.y as u32,
+                depth_or_array_layers: 1,
+            };
+
+            encoder.copy_texture_to_buffer(
+                src_image.texture.as_image_copy(),
+                ImageCopyBuffer {
+                    buffer: &image_copier.buffer,
+                    layout: ImageDataLayout {
+                        offset: 0,
+                        bytes_per_row: Some(
+                            std::num::NonZeroU32::new(padded_bytes_per_row as u32).unwrap().into(),
+                        ),
+                        rows_per_image: None,
+                    },
+                },
+                texture_extent,
+            );
+
+            let render_queue = world.get_resource::<RenderQueue>().unwrap();
+            render_queue.submit(std::iter::once(encoder.finish()));
         }
+
+        Ok(())
     }
+}
diff --git a/crates/bevy_frame_capture/src/lib.rs b/crates/bevy_frame_capture/src/lib.rs
@@ -5,7 +5,7 @@ mod scene;
 use bevy::app::{App, Plugin, PostUpdate};
 use image_copy::ImageCopyPlugin;
 use scene::update;
-pub use scene::{SceneInfo,setup_render_target, white_img_placeholder, CurrImageBase64};
+pub use scene::{setup_render_target, white_img_placeholder, CurrImageBase64, SceneInfo};
 
 pub struct FrameCapturePlugin;
 
@@ -16,4 +16,4 @@ impl Plugin for FrameCapturePlugin {
         app.init_resource::<scene::SceneInfo>();
         app.add_event::<scene::SceneInfo>();
     }
-}
+}
diff --git a/crates/bevy_frame_capture/src/scene.rs b/crates/bevy_frame_capture/src/scene.rs
@@ -130,21 +130,21 @@ pub fn update(
 
                 let rgba_img = img_bytes.clone().try_into_dynamic().unwrap().to_rgba8();
 
-                if let Some(base64_img) = curr_base64_img.as_mut(){
+                if let Some(base64_img) = curr_base64_img.as_mut() {
                     base64_img.0 = image_to_browser_base64(&rgba_img).unwrap();
+                    log::info!("updated current image");
                 } else {
                     let images_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("test_images");
                     log::info!("Saving image to: {images_dir:?}");
-                    
+
                     std::fs::create_dir_all(&images_dir).unwrap();
-    
+
                     let uuid = bevy::utils::Uuid::new_v4();
                     let image_path = images_dir.join(format!("{uuid}.png"));
                     if let Err(e) = rgba_img.save(image_path) {
                         panic!("Failed to save image: {}", e);
                     };
                 };
-
             }
         } else {
             scene_controller.state.decrement();

diff --git a/crates/bevy_ws_server/src/lib.rs b/crates/bevy_ws_server/src/lib.rs
@@ -21,8 +21,6 @@ use futures::sink::{Sink, SinkExt};
 use smol::{prelude::*, Async};
 use tungstenite::Message;
 
-// use tungstenite::Message;
-
 pub struct WsPlugin;
 
 impl Plugin for WsPlugin {