diff --git a/.gitignore b/.gitignore index e811ba0..0f4fba1 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ # Rust target/ +Cargo.lock # ggml binaries & other binaries in general *.bin diff --git a/Cargo.lock b/Cargo.lock index dc3e7e1..eb59288 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -690,6 +690,20 @@ dependencies = [ "encase_derive_impl", ] +[[package]] +name = "bevy_frame_capture" +version = "0.1.0" +dependencies = [ + "anyhow", + "base64", + "bevy", + "futures-intrusive", + "image", + "log", + "pollster", + "wgpu", +] + [[package]] name = "bevy_gaussian_splatting" version = "2.0.2" @@ -1189,18 +1203,20 @@ dependencies = [ [[package]] name = "bevy_ws_server" -version = "0.2.1" -source = "git+https://github.com/cs50victor/bevy-ws-server.git?branch=main#daaa533548c169a72c65270e74491d757c539273" +version = "0.1.0" dependencies = [ "anyhow", "async-channel 2.1.1", "async-native-tls", "async-net", "async-tungstenite", + "base64", "bevy", "crossbeam-channel", "futures", + "image", "log", + "pollster", "smol", "tungstenite", ] @@ -2751,20 +2767,17 @@ dependencies = [ "async-net", "base64", "bevy", + "bevy_frame_capture", "bevy_gaussian_splatting", "bevy_panorbit_camera", "bevy_ws_server", "crossbeam-channel", "dotenvy", - "futures-intrusive", - "image", "log", - "pollster", "pretty_env_logger", "serde", "serde_json", "tungstenite", - "wgpu", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 8e41ba1..fecb75a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [workspace] resolver = "2" -members = ["new_media"] +members = ["crates/*"] package.rust-version = "1.75.0" @@ -19,3 +19,20 @@ panic = "abort" # Abort on panic, instead of unwinding the stack. [profile.prod] inherits = "release" strip = true # Automatically strip symbols from the binary. + + +[workspace.dependencies] +anyhow = "1.0.75" +async-net = "2.0.0" +base64 = "0.21.5" +bevy = { version = "0.12.1", default-features = false, features = [ + "bevy_asset", + "bevy_core_pipeline", + "bevy_render", +] } +bevy_panorbit_camera = "0.10.0" +crossbeam-channel = "0.5.10" +log = "0.4.20" +pretty_env_logger = "0.5.0" +wgpu = { version = "0.17.1", features = ["vulkan-portability"] } +tungstenite = "0.21.0" diff --git a/crates/bevy_frame_capture/Cargo.toml b/crates/bevy_frame_capture/Cargo.toml new file mode 100644 index 0000000..0e86563 --- /dev/null +++ b/crates/bevy_frame_capture/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "bevy_frame_capture" +version = "0.1.0" +edition = "2021" +rust-version.workspace = true + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +anyhow = { workspace = true } +bevy = { workspace = true } +base64 = { workspace = true } +image = { version = "0.24.7", default-features = false, features = ["png"] } +pollster = "0.3.0" +wgpu = { workspace = true } +log = { workspace = true } +futures-intrusive = "0.5.0" diff --git a/crates/bevy_frame_capture/src/image_copy.rs b/crates/bevy_frame_capture/src/image_copy.rs new file mode 100644 index 0000000..a3ce5c0 --- /dev/null +++ b/crates/bevy_frame_capture/src/image_copy.rs @@ -0,0 +1,181 @@ +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>, + render_device: Res, +) { + 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(); + } + .block_on(); + } +} + +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); + + render_app.add_systems(ExtractSchedule, image_copy_extract); + + let mut graph = render_app.world.get_resource_mut::().unwrap(); + + graph.add_node(IMAGE_COPY, ImageCopyDriver); + + graph.add_node_edge(IMAGE_COPY, bevy::render::main_graph::node::CAMERA_DRIVER); + } +} + +#[derive(Clone, Default, Resource, Deref, DerefMut)] +pub struct ImageCopiers(pub Vec); + +#[derive(Clone, Component)] +pub struct ImageCopier { + buffer: Buffer, + enabled: Arc, + src_image: Handle, + dst_image: Handle, +} + +impl ImageCopier { + pub fn new( + src_image: Handle, + dst_image: Handle, + 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) + } +} + +pub fn image_copy_extract(mut commands: Commands, image_copiers: Extract>) { + commands.insert_resource(ImageCopiers( + image_copiers.iter().cloned().collect::>(), + )); +} + +#[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::().unwrap(); + let gpu_images = world.get_resource::>().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::().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 new file mode 100644 index 0000000..633a4fa --- /dev/null +++ b/crates/bevy_frame_capture/src/lib.rs @@ -0,0 +1,3 @@ +/// Derived from: https://github.com/bevyengine/bevy/pull/5550 +pub mod scene; +pub mod image_copy; diff --git a/crates/bevy_frame_capture/src/scene.rs b/crates/bevy_frame_capture/src/scene.rs new file mode 100644 index 0000000..6297737 --- /dev/null +++ b/crates/bevy_frame_capture/src/scene.rs @@ -0,0 +1,180 @@ +use std::io::Cursor; + +use anyhow::Result; +use base64::{engine::general_purpose, Engine}; +use bevy::{ + prelude::*, + render::{ + camera::RenderTarget, + renderer::RenderDevice, + }, +}; + +use image::{ImageBuffer, ImageOutputFormat, Rgba, RgbaImage}; +use wgpu::{Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages}; + +use super::image_copy::ImageCopier; + +#[derive(Component, Default)] +pub struct CaptureCamera; + +#[derive(Component, Deref, DerefMut)] +struct ImageToSave(Handle); + +#[derive(Resource)] +pub struct CurrImageBase64(pub String); + +#[derive(Resource)] +pub struct StreamingFrameData { + pixel_size: u32, +} + +pub struct CaptureFramePlugin; +impl Plugin for CaptureFramePlugin { + fn build(&self, app: &mut App) { + app.add_systems( + PostUpdate, + update.run_if(resource_exists::()), + ); + } +} + +#[derive(Debug, Default, Resource, Event)] +pub struct SceneController { + state: SceneState, + name: String, + width: u32, + height: u32, +} + +impl SceneController { + pub fn new(width: u32, height: u32) -> SceneController { + SceneController { state: SceneState::BuildScene, name: String::from(""), width, height } + } + + pub fn dimensions(&self) -> (u32, u32) { + (self.width, self.height) + } +} + +#[derive(Debug, Default)] +pub enum SceneState { + #[default] + BuildScene, + Render(u32), +} + +impl SceneState { + pub fn decrement(&mut self) { + if let SceneState::Render(n) = self { + *n -= 1; + } + } +} + +pub fn setup_render_target( + commands: &mut Commands, + images: &mut ResMut>, + render_device: &Res, + scene_controller: &mut ResMut, + pre_roll_frames: u32, + scene_name: String, +) -> RenderTarget { + let size = Extent3d { + width: scene_controller.width, + height: scene_controller.height, + ..Default::default() + }; + + // This is the texture that will be rendered to. + let mut render_target_image = Image { + texture_descriptor: TextureDescriptor { + label: None, + size, + dimension: TextureDimension::D2, + format: TextureFormat::Rgba8UnormSrgb, + mip_level_count: 1, + sample_count: 1, + usage: TextureUsages::COPY_SRC + | TextureUsages::COPY_DST + | TextureUsages::TEXTURE_BINDING + | TextureUsages::RENDER_ATTACHMENT, + view_formats: &[], + }, + ..Default::default() + }; + render_target_image.resize(size); + let render_target_image_handle = images.add(render_target_image); + + // This is the texture that will be copied to. + let mut cpu_image = Image { + texture_descriptor: TextureDescriptor { + label: None, + size, + dimension: TextureDimension::D2, + format: TextureFormat::Rgba8UnormSrgb, + mip_level_count: 1, + sample_count: 1, + usage: TextureUsages::COPY_DST | TextureUsages::TEXTURE_BINDING, + view_formats: &[], + }, + ..Default::default() + }; + cpu_image.resize(size); + let cpu_image_handle = images.add(cpu_image); + + commands.spawn(ImageCopier::new( + render_target_image_handle.clone(), + cpu_image_handle.clone(), + size, + render_device, + )); + + commands.spawn(ImageToSave(cpu_image_handle)); + + scene_controller.state = SceneState::Render(pre_roll_frames); + scene_controller.name = scene_name; + RenderTarget::Image(render_target_image_handle) +} + +fn update( + mut images: ResMut>, + images_to_save: Query<&ImageToSave>, + mut curr_base64_img: ResMut, + single_frame_data: ResMut, + mut scene_controller: ResMut, +) { + if let SceneState::Render(n) = scene_controller.state { + if n < 1 { + let single_frame_data = single_frame_data.into_inner(); + let _pixel_size = single_frame_data.pixel_size; + for image in images_to_save.iter() { + let img_bytes = images.get_mut(image.id()).unwrap(); + + let rgba_img = match img_bytes.clone().try_into_dynamic() { + Ok(img) => img.to_rgba8(), + Err(e) => panic!("Failed to create image buffer {e:?}"), + }; + + rgba_img.save("wtf.png"); + log::info!("saved"); + curr_base64_img.0 = image_to_browser_base64(&rgba_img).unwrap(); + } + } else { + scene_controller.state.decrement(); + } + } +} + +fn image_to_browser_base64(img: &ImageBuffer, Vec>) -> Result { + let mut image_data: Vec = Vec::new(); + img.write_to(&mut Cursor::new(&mut image_data), ImageOutputFormat::Png)?; + let res_base64 = general_purpose::STANDARD.encode(image_data); + Ok(format!("data:image/png;base64,{}", res_base64)) +} + +pub fn white_img_placeholder(w: u32, h: u32) -> String { + let img = RgbaImage::new(w, h); + // img.iter_mut().for_each(|pixel| *pixel = 255); + image_to_browser_base64(&img).unwrap() +} diff --git a/crates/bevy_ws_server/Cargo.toml b/crates/bevy_ws_server/Cargo.toml new file mode 100644 index 0000000..459b623 --- /dev/null +++ b/crates/bevy_ws_server/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "bevy_ws_server" +version = "0.1.0" +edition = "2021" +rust-version.workspace = true + +[dependencies] +anyhow = {workspace = true} +base64 = {workspace = true} +bevy = { version = "0.12.1", default-features = false, features = ["multi-threaded"] } +async-net = {workspace = true} +futures = "0.3.29" +async-tungstenite = "0.24.0" +log = {workspace = true} +crossbeam-channel = {workspace = true} +async-channel = "2.1.0" +async-native-tls = "0.5.0" +smol = "2.0.0" +tungstenite = {workspace = true} +pollster = "0.3.0" +image = { version = "0.24.7", default-features = false, features = ["png"] } diff --git a/crates/bevy_ws_server/src/lib.rs b/crates/bevy_ws_server/src/lib.rs new file mode 100644 index 0000000..753ddb7 --- /dev/null +++ b/crates/bevy_ws_server/src/lib.rs @@ -0,0 +1,216 @@ +use async_net::SocketAddr; + +use bevy::prelude::*; +use bevy::tasks::{IoTaskPool, Task}; +use crossbeam_channel::{Receiver, Sender}; +use futures::{pin_mut, select, FutureExt}; + +use std::net::{TcpListener, TcpStream}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use anyhow::{bail, Result}; +pub use async_channel::TryRecvError as ReceiveError; +use async_native_tls::{TlsAcceptor, TlsStream}; +use async_tungstenite::{tungstenite, WebSocketStream}; +use futures::sink::{Sink, SinkExt}; +use smol::{prelude::*, Async}; +use tungstenite::Message; + +// use tungstenite::Message; + +pub struct WsPlugin; + +impl Plugin for WsPlugin { + fn build(&self, app: &mut App) { + let (ws_tx, ws_rx) = crossbeam_channel::unbounded(); + app.insert_resource(WsListener::new(ws_tx)) + .insert_resource(WsAcceptQueue { ws_rx }) + .add_systems(Update, accept_ws_from_queue); + } +} + +#[derive(Resource)] +pub struct WsListener { + ws_tx: Sender, +} + +#[derive(Resource)] +pub struct WsAcceptQueue { + ws_rx: Receiver, +} + +/// A WebSocket or WebSocket+TLS connection. +enum WsStream { + /// A plain WebSocket connection. + Plain(WebSocketStream>), + + /// A WebSocket connection secured by TLS. + Tls(WebSocketStream>>), +} + +impl Sink for WsStream { + type Error = tungstenite::Error; + + fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match &mut *self { + WsStream::Plain(s) => Pin::new(s).poll_ready(cx), + WsStream::Tls(s) => Pin::new(s).poll_ready(cx), + } + } + + fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + match &mut *self { + WsStream::Plain(s) => Pin::new(s).start_send(item), + WsStream::Tls(s) => Pin::new(s).start_send(item), + } + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match &mut *self { + WsStream::Plain(s) => Pin::new(s).poll_flush(cx), + WsStream::Tls(s) => Pin::new(s).poll_flush(cx), + } + } + + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match &mut *self { + WsStream::Plain(s) => Pin::new(s).poll_close(cx), + WsStream::Tls(s) => Pin::new(s).poll_close(cx), + } + } +} + +impl Stream for WsStream { + type Item = tungstenite::Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match &mut *self { + WsStream::Plain(s) => Pin::new(s).poll_next(cx), + WsStream::Tls(s) => Pin::new(s).poll_next(cx), + } + } +} + +impl WsListener { + pub fn new(ws_tx: Sender) -> Self { + Self { ws_tx } + } + + pub fn listen(&self, bind_to: impl Into, tls: Option) -> Result { + let listener = Async::::bind(bind_to).expect("cannot bind to the address"); + + let host = match &tls { + None => format!("ws://{}", listener.get_ref().local_addr().unwrap()), + Some(_) => format!("wss://{}", listener.get_ref().local_addr().unwrap()), + }; + + let task_pool = IoTaskPool::get(); + let ws_tx = self.ws_tx.clone(); + let task = task_pool.spawn(async move { + let tls = tls; + loop { + match listener.accept().await { + Ok((stream, addr)) => { + log::debug!("new connection from {}", addr); + let ws_tx = ws_tx.clone(); + + let tls = tls.clone(); + + let accept = async move { + let ws_stream = match &tls { + None => { + let stream = match async_tungstenite::accept_async(stream).await { + Ok(s) => s, + Err(e) => { + bail!("Error accepting web socket connection | {e}"); + } + }; + + WsStream::Plain(stream) + }, + Some(tls) => { + let tls_stream = match tls.accept(stream).await { + Ok(s) => s, + Err(e) => { + bail!("Error accepting creating TLS web socket connection | {e}"); + } + }; + let stream = match async_tungstenite::accept_async(tls_stream).await { + Ok(s) => s, + Err(e) => { + bail!("Error accepting web socket connection | {e}"); + } + }; + WsStream::Tls(stream) + } + }; + let _ = ws_tx.send(ws_stream); + Ok(()) + }; + + task_pool.spawn(accept).detach(); + } + Err(e) => { + log::error!("error accepting a new connection: {}", e); + } + } + } + }); + + task.detach(); + Ok(host) + } +} + +#[derive(Component)] +pub struct WsConnection { + _io: Task<()>, + sender: async_channel::Sender, + receiver: async_channel::Receiver, +} + +impl WsConnection { + pub fn send(&self, message: Message) -> bool { + self.sender.try_send(message).is_ok() + } + + pub fn receive(&self) -> Result { + self.receiver.try_recv() + } +} + +pub fn accept_ws_from_queue(mut commands: Commands, queue: ResMut) { + for mut websocket in queue.ws_rx.try_iter() { + let (message_tx, io_message_rx) = async_channel::unbounded::(); + let (io_message_tx, message_rx) = async_channel::unbounded::(); + + let io = IoTaskPool::get().spawn(async move { + loop { + let from_channel = io_message_rx.recv().fuse(); + let from_ws = websocket.next().fuse(); + + pin_mut!(from_channel, from_ws); + + select! { + message = from_channel => if let Ok(message) = message { + let _ = websocket.send(message).await; + } else { + break; + }, + message = from_ws => if let Some(Ok(message)) = message { + let _ = io_message_tx.send(message).await; + } else { + break; + }, + complete => break, + } + } + }); + commands.spawn(WsConnection { + _io: io, + sender: message_tx, + receiver: message_rx, + }); + } +} diff --git a/crates/new_media/Cargo.toml b/crates/new_media/Cargo.toml new file mode 100644 index 0000000..2e9312d --- /dev/null +++ b/crates/new_media/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "new_media" +version = "0.0.1" +edition = "2021" +rust-version.workspace = true +default-run = "new_media" + +[dependencies] +anyhow = {workspace = true} +async-net = {workspace = true} +base64 = {workspace = true} +bevy = { workspace = true } +bevy_ws_server = { path = "../bevy_ws_server" } +bevy_frame_capture = { path = "../bevy_frame_capture" } +bevy_gaussian_splatting = { version = "2.0.2", default-features = true } +bevy_panorbit_camera = { workspace = true } +crossbeam-channel = { workspace = true } +dotenvy = "0.15.7" +log = { workspace = true } +pretty_env_logger = {workspace = true} +serde = { version = "1.0.192", features = ["derive"] } +serde_json = "1.0.108" +tungstenite = {workspace = true} + +# [target.i686-unknown-linux-gnu.dependencies] +# winit = { version = "0.29.10", default-features = false, features = ["x11"]} diff --git a/new_media/src/controls.rs b/crates/new_media/src/controls.rs similarity index 100% rename from new_media/src/controls.rs rename to crates/new_media/src/controls.rs diff --git a/new_media/src/main.rs b/crates/new_media/src/main.rs similarity index 84% rename from new_media/src/main.rs rename to crates/new_media/src/main.rs index c4b5588..a55a7c3 100644 --- a/new_media/src/main.rs +++ b/crates/new_media/src/main.rs @@ -1,6 +1,5 @@ #![feature(ascii_char, async_closure, slice_pattern)] mod controls; -mod frame_capture; mod server; use bevy_ws_server::WsPlugin; @@ -24,7 +23,7 @@ fn setup_gaussian_cloud( mut commands: Commands, _asset_server: Res, mut gaussian_assets: ResMut>, - mut scene_controller: ResMut, + mut scene_controller: ResMut, mut images: ResMut>, render_device: Res, ) { @@ -36,7 +35,7 @@ fn setup_gaussian_cloud( let cloud = gaussian_assets.add(GaussianCloud::test_model()); - let render_target = frame_capture::scene::setup_render_target( + let render_target = bevy_frame_capture::scene::setup_render_target( &mut commands, &mut images, &render_device, @@ -84,8 +83,8 @@ fn main() { let config = AppConfig { width: 1920, height: 1080 }; App::new() - .insert_resource(frame_capture::scene::SceneController::new(config.width, config.height)) - .insert_resource(frame_capture::scene::CurrImageBase64(frame_capture::scene::white_img_placeholder(config.width, config.height))) + .insert_resource(bevy_frame_capture::scene::SceneController::new(config.width, config.height)) + .insert_resource(bevy_frame_capture::scene::CurrImageBase64(bevy_frame_capture::scene::white_img_placeholder(config.width, config.height))) .insert_resource(ClearColor(Color::rgb_u8(0, 0, 0))) .add_plugins(( DefaultPlugins @@ -96,14 +95,14 @@ fn main() { close_when_requested: false, }).disable::(), WsPlugin, - frame_capture::image_copy::ImageCopyPlugin, - frame_capture::scene::CaptureFramePlugin, + bevy_frame_capture::image_copy::ImageCopyPlugin, + bevy_frame_capture::scene::CaptureFramePlugin, ScheduleRunnerPlugin::run_loop(std::time::Duration::from_secs_f64(1.0 / 60.0)), PanOrbitCameraPlugin, GaussianSplattingPlugin, )) - .init_resource::() - .add_event::() + .init_resource::() + .add_event::() .add_systems(Startup, (start_ws, setup_gaussian_cloud)) .add_systems(Update, ( move_camera, diff --git a/new_media/src/server.rs b/crates/new_media/src/server.rs similarity index 94% rename from new_media/src/server.rs rename to crates/new_media/src/server.rs index 93ead75..854e771 100644 --- a/new_media/src/server.rs +++ b/crates/new_media/src/server.rs @@ -7,15 +7,15 @@ use std::fmt::Debug; use bevy::prelude::*; use serde::{Deserialize, Serialize}; -use crate::frame_capture::scene::CurrImageBase64; +use bevy_frame_capture::scene::CurrImageBase64; #[derive(Debug, Serialize, Deserialize)] -pub struct ServerMsg { +pub struct HttpServerMsg { data: Option, error: Option, } -impl ServerMsg { +impl HttpServerMsg { pub fn data(data: T) -> Self { Self { data: Some(data), error: None } } diff --git a/new_media/Cargo.toml b/new_media/Cargo.toml deleted file mode 100644 index 8bbe5f1..0000000 --- a/new_media/Cargo.toml +++ /dev/null @@ -1,33 +0,0 @@ -[package] -name = "new_media" -version = "0.0.1" -edition = "2021" -rust-version.workspace = true -default-run = "new_media" - -[dependencies] -anyhow = "1.0.75" -async-net = "2.0.0" -base64 = "0.21.5" -bevy = { version = "0.12.0", default-features = false, features = [ - "bevy_asset", - "bevy_core_pipeline", - "bevy_render", -] } -bevy_ws_server = { git = "https://github.com/cs50victor/bevy-ws-server.git", branch = "main" } -bevy_gaussian_splatting = { version = "2.0.2", default-features = true } -bevy_panorbit_camera = "0.10.0" -crossbeam-channel = "0.5.10" -dotenvy = "0.15.7" -futures-intrusive = "0.5.0" -image = { version = "0.24.7", default-features = false, features = ["png"] } -log = "0.4.20" -pollster = "0.3.0" -pretty_env_logger = "0.5.0" -serde = { version = "1.0.192", features = ["derive"] } -serde_json = "1.0.108" -wgpu = "0.17.2" -tungstenite = "0.21.0" - -# [target.i686-unknown-linux-gnu.dependencies] -# winit = { version = "0.29.10", default-features = false, features = ["x11"]} diff --git a/new_media/src/frame_capture.rs b/new_media/src/frame_capture.rs deleted file mode 100644 index 8229599..0000000 --- a/new_media/src/frame_capture.rs +++ /dev/null @@ -1,364 +0,0 @@ -/// Derived from: https://github.com/bevyengine/bevy/pull/5550 -pub mod image_copy { - 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>, - render_device: Res, - ) { - 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(); - } - .block_on(); - } - } - - 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); - - render_app.add_systems(ExtractSchedule, image_copy_extract); - - let mut graph = render_app.world.get_resource_mut::().unwrap(); - - graph.add_node(IMAGE_COPY, ImageCopyDriver); - - graph.add_node_edge(IMAGE_COPY, bevy::render::main_graph::node::CAMERA_DRIVER); - } - } - - #[derive(Clone, Default, Resource, Deref, DerefMut)] - pub struct ImageCopiers(pub Vec); - - #[derive(Clone, Component)] - pub struct ImageCopier { - buffer: Buffer, - enabled: Arc, - src_image: Handle, - dst_image: Handle, - } - - impl ImageCopier { - pub fn new( - src_image: Handle, - dst_image: Handle, - 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) - } - } - - pub fn image_copy_extract(mut commands: Commands, image_copiers: Extract>) { - commands.insert_resource(ImageCopiers( - image_copiers.iter().cloned().collect::>(), - )); - } - - #[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::().unwrap(); - let gpu_images = world.get_resource::>().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::().unwrap(); - log::error!("0"); - // render_queue.submit(std::iter::once(encoder.finish())); - log::error!("1"); - } - - Ok(()) - } - } -} -pub mod scene { - - use std::io::Cursor; - - use anyhow::Result; - use base64::{engine::general_purpose, Engine}; - use bevy::{ - prelude::*, - render::{ - camera::RenderTarget, - renderer::RenderDevice, - }, - }; - - use image::{ImageBuffer, ImageOutputFormat, Rgba, RgbaImage}; - use wgpu::{Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages}; - - use super::image_copy::ImageCopier; - - #[derive(Component, Default)] - pub struct CaptureCamera; - - #[derive(Component, Deref, DerefMut)] - struct ImageToSave(Handle); - - #[derive(Resource)] - pub struct CurrImageBase64(pub String); - - pub struct CaptureFramePlugin; - impl Plugin for CaptureFramePlugin { - fn build(&self, app: &mut App) { - app.add_systems( - PostUpdate, - update.run_if(resource_exists::()), - ); - } - } - - #[derive(Debug, Default, Resource, Event)] - pub struct SceneController { - state: SceneState, - name: String, - width: u32, - height: u32, - } - - impl SceneController { - pub fn new(width: u32, height: u32) -> SceneController { - SceneController { state: SceneState::BuildScene, name: String::from(""), width, height } - } - - pub fn dimensions(&self) -> (u32, u32) { - (self.width, self.height) - } - } - - #[derive(Debug, Default)] - pub enum SceneState { - #[default] - BuildScene, - Render(u32), - } - - impl SceneState { - pub fn decrement(&mut self) { - if let SceneState::Render(n) = self { - *n -= 1; - } - } - } - - pub fn setup_render_target( - commands: &mut Commands, - images: &mut ResMut>, - render_device: &Res, - scene_controller: &mut ResMut, - pre_roll_frames: u32, - scene_name: String, - ) -> RenderTarget { - let size = Extent3d { - width: scene_controller.width, - height: scene_controller.height, - ..Default::default() - }; - - // This is the texture that will be rendered to. - let mut render_target_image = Image { - texture_descriptor: TextureDescriptor { - label: None, - size, - dimension: TextureDimension::D2, - format: TextureFormat::Rgba8UnormSrgb, - mip_level_count: 1, - sample_count: 1, - usage: TextureUsages::COPY_SRC - | TextureUsages::COPY_DST - | TextureUsages::TEXTURE_BINDING - | TextureUsages::RENDER_ATTACHMENT, - view_formats: &[], - }, - ..Default::default() - }; - render_target_image.resize(size); - let render_target_image_handle = images.add(render_target_image); - - // This is the texture that will be copied to. - let mut cpu_image = Image { - texture_descriptor: TextureDescriptor { - label: None, - size, - dimension: TextureDimension::D2, - format: TextureFormat::Rgba8UnormSrgb, - mip_level_count: 1, - sample_count: 1, - usage: TextureUsages::COPY_DST | TextureUsages::TEXTURE_BINDING, - view_formats: &[], - }, - ..Default::default() - }; - cpu_image.resize(size); - let cpu_image_handle = images.add(cpu_image); - - commands.spawn(ImageCopier::new( - render_target_image_handle.clone(), - cpu_image_handle.clone(), - size, - render_device, - )); - - commands.spawn(ImageToSave(cpu_image_handle)); - - scene_controller.state = SceneState::Render(pre_roll_frames); - scene_controller.name = scene_name; - RenderTarget::Image(render_target_image_handle) - } - - fn update( - mut images: ResMut>, - images_to_save: Query<&ImageToSave>, - mut curr_base64_img: ResMut, - single_frame_data: ResMut, - mut scene_controller: ResMut, - ) { - if let SceneState::Render(n) = scene_controller.state { - if n < 1 { - let single_frame_data = single_frame_data.into_inner(); - let _pixel_size = single_frame_data.pixel_size; - for image in images_to_save.iter() { - let img_bytes = images.get_mut(image.id()).unwrap(); - - let rgba_img = match img_bytes.clone().try_into_dynamic() { - Ok(img) => img.to_rgba8(), - Err(e) => panic!("Failed to create image buffer {e:?}"), - }; - - rgba_img.save("wtf.png"); - log::info!("saved"); - curr_base64_img.0 = image_to_browser_base64(&rgba_img).unwrap(); - } - } else { - scene_controller.state.decrement(); - } - } - } - - fn image_to_browser_base64(img: &ImageBuffer, Vec>) -> Result { - let mut image_data: Vec = Vec::new(); - img.write_to(&mut Cursor::new(&mut image_data), ImageOutputFormat::Png)?; - let res_base64 = general_purpose::STANDARD.encode(image_data); - Ok(format!("data:image/png;base64,{}", res_base64)) - } - - pub fn white_img_placeholder(w: u32, h: u32) -> String { - let img = RgbaImage::new(w, h); - // img.iter_mut().for_each(|pixel| *pixel = 255); - image_to_browser_base64(&img).unwrap() - } -}