diff --git a/Cargo.toml b/Cargo.toml index d9b84b2..7c8c3f4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,7 +20,6 @@ panic = "abort" # Abort on panic, instead of unwinding the stack. inherits = "release" strip = true # Automatically strip symbols from the binary. - [workspace.dependencies] anyhow = "1.0.75" async-net = "2.0.0" @@ -29,10 +28,11 @@ bevy = { version = "0.12.1", default-features = false, features = [ "bevy_asset", "bevy_core_pipeline", "bevy_render", + "multi-threaded" ] } 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"] } +wgpu = "0.17.2" tungstenite = "0.21.0" diff --git a/crates/bevy_frame_capture/src/image_copy.rs b/crates/bevy_frame_capture/src/image_copy.rs deleted file mode 100644 index 88a71a8..0000000 --- a/crates/bevy_frame_capture/src/image_copy.rs +++ /dev/null @@ -1,176 +0,0 @@ -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 deleted file mode 100644 index 0bdd9dd..0000000 --- a/crates/bevy_frame_capture/src/lib.rs +++ /dev/null @@ -1,19 +0,0 @@ -/// Derived from: https://github.com/bevyengine/bevy/pull/5550 -mod image_copy; -mod scene; - -use bevy::app::{App, Plugin, PostUpdate}; -use image_copy::ImageCopyPlugin; -use scene::update; -pub use scene::{setup_render_target, white_img_placeholder, CurrImageBase64, SceneInfo}; - -pub struct FrameCapturePlugin; - -impl Plugin for FrameCapturePlugin { - fn build(&self, app: &mut App) { - app.add_plugins(ImageCopyPlugin); - app.add_systems(PostUpdate, update); - app.init_resource::(); - app.add_event::(); - } -} diff --git a/crates/bevy_frame_capture/src/scene.rs b/crates/bevy_frame_capture/src/scene.rs deleted file mode 100644 index 9675a5e..0000000 --- a/crates/bevy_frame_capture/src/scene.rs +++ /dev/null @@ -1,167 +0,0 @@ -use std::{io::Cursor, path::PathBuf}; - -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(Resource)] -pub struct CurrImageBase64(pub String); - -#[derive(Component, Deref, DerefMut)] -pub struct ImageToSave(Handle); - -#[derive(Debug, Default, Resource, Event)] -pub struct SceneInfo { - state: SceneState, - name: String, - width: u32, - height: u32, -} - -impl SceneInfo { - pub fn new(width: u32, height: u32) -> SceneInfo { - SceneInfo { 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) -} - -pub fn update( - images_to_save: Query<&ImageToSave>, - mut images: ResMut>, - mut curr_base64_img: Option>, - mut scene_controller: ResMut, -) { - if let SceneState::Render(n) = scene_controller.state { - if n < 1 { - for image in images_to_save.iter() { - let img_bytes = images.get_mut(image.id()).unwrap(); - - let rgba_img = img_bytes.clone().try_into_dynamic().unwrap().to_rgba8(); - - 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(); - } - } -} - -// useful utils -fn image_to_browser_base64(img: &ImageBuffer, Vec>) -> anyhow::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_frame_capture/Cargo.toml b/crates/bevy_headless/Cargo.toml similarity index 52% rename from crates/bevy_frame_capture/Cargo.toml rename to crates/bevy_headless/Cargo.toml index f5d7364..e90b288 100644 --- a/crates/bevy_frame_capture/Cargo.toml +++ b/crates/bevy_headless/Cargo.toml @@ -1,15 +1,17 @@ [package] -name = "bevy_frame_capture" +name = "bevy_headless" version = "0.1.0" edition = "2021" rust-version.workspace = true [dependencies] anyhow = { workspace = true } -bevy = { workspace = true } base64 = { workspace = true } -image = { version = "0.24.7", default-features = false, features = ["png"] } -pollster = "0.3.0" +bevy = { workspace = true } +bytemuck = "1.14.1" +futures = "0.3.30" +futures-lite = "2.2.0" +image = { version = "0.24.8", features = ["exr", "png", "jpeg", "webp"], default-features = false } wgpu = { workspace = true } log = { workspace = true } -futures-intrusive = "0.5.0" +parking_lot = "0.12.1" diff --git a/crates/bevy_headless/src/lib.rs b/crates/bevy_headless/src/lib.rs new file mode 100644 index 0000000..7808132 --- /dev/null +++ b/crates/bevy_headless/src/lib.rs @@ -0,0 +1,11 @@ +// Derived from https://github.com/paulkre/bevy_image_export +mod node; +mod plugin; +mod utils; + +pub use plugin::{ + CurrImageContainer, GpuImageExportSource, ImageExportBundle, ImageExportPlugin, + ImageExportSettings, ImageExportSource, ImageExportSystems, +}; + +pub use utils::{setup_render_target, SceneInfo}; diff --git a/crates/bevy_headless/src/node.rs b/crates/bevy_headless/src/node.rs new file mode 100644 index 0000000..d8a5c48 --- /dev/null +++ b/crates/bevy_headless/src/node.rs @@ -0,0 +1,45 @@ +use crate::ImageExportSource; + +use bevy::{ + ecs::world::World, + render::{ + render_asset::RenderAssets, + render_graph::{Node, NodeRunError, RenderGraphContext}, + render_resource::{ImageCopyBuffer, ImageDataLayout}, + renderer::RenderContext, + texture::Image, + }, +}; + +pub const NODE_NAME: &str = "image_export"; + +pub struct ImageExportNode; +impl Node for ImageExportNode { + fn run( + &self, + _: &mut RenderGraphContext, + render_context: &mut RenderContext, + world: &World, + ) -> Result<(), NodeRunError> { + for (_, source) in world.resource::>().iter() { + if let Some(gpu_image) = + world.resource::>().get(&source.source_handle) + { + render_context.command_encoder().copy_texture_to_buffer( + gpu_image.texture.as_image_copy(), + ImageCopyBuffer { + buffer: &source.buffer, + layout: ImageDataLayout { + offset: 0, + bytes_per_row: Some(source.padded_bytes_per_row), + rows_per_image: None, + }, + }, + source.source_size, + ); + } + } + + Ok(()) + } +} diff --git a/crates/bevy_headless/src/plugin.rs b/crates/bevy_headless/src/plugin.rs new file mode 100644 index 0000000..c189df6 --- /dev/null +++ b/crates/bevy_headless/src/plugin.rs @@ -0,0 +1,318 @@ +use std::sync::Arc; + +use crate::{ + node::{ImageExportNode, NODE_NAME}, + utils::CurrImage, +}; +use bevy::{ + app::{App, Plugin, PluginGroup, PostUpdate}, + asset::{Asset, AssetApp, Handle}, + ecs::{ + bundle::Bundle, + component::Component, + entity::Entity, + query::{QueryItem, With, Without}, + schedule::{apply_deferred, IntoSystemConfigs, IntoSystemSetConfigs, SystemSet}, + system::{ + lifetimeless::SRes, Commands, Local, Query, Res, ResMut, Resource, SystemParamItem, + }, + }, + log::LogPlugin, + reflect::{Reflect, TypeUuid}, + render::{ + camera::CameraUpdateSystem, + extract_component::{ExtractComponent, ExtractComponentPlugin}, + main_graph::node::CAMERA_DRIVER, + render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin, RenderAssets}, + render_graph::RenderGraph, + render_resource::{Buffer, BufferDescriptor, BufferUsages, Extent3d, MapMode}, + renderer::RenderDevice, + texture::{Image, ImagePlugin}, + Render, RenderApp, RenderSet, + }, + window::WindowPlugin, + DefaultPlugins, +}; +use bytemuck::AnyBitPattern; +use futures::channel::oneshot; +use image::{EncodableLayout, ImageBuffer, Pixel, PixelWithColorType, Rgba}; + +use parking_lot::Mutex; +use wgpu::Maintain; +use ImageExportSystems::{SetupImageExport, SetupImageExportFlush}; + +#[derive(Asset, Clone, TypeUuid, Default, Reflect)] +#[uuid = "d619b2f8-58cf-42f6-b7da-028c0595f7aa"] +pub struct ImageExportSource(pub Handle); + +impl From> for ImageExportSource { + fn from(value: Handle) -> Self { + Self(value) + } +} + +#[derive(Component, Clone)] +pub struct ImageExportSettings { + /// The image file extension. "png", "jpeg", "webp", or "exr". + pub extension: String, +} + +pub struct GpuImageExportSource { + pub buffer: Buffer, + pub source_handle: Handle, + pub source_size: Extent3d, + pub bytes_per_row: u32, + pub padded_bytes_per_row: u32, +} + +impl GpuImageExportSource { + fn get_bps(&self) -> (usize, usize, Extent3d) { + (self.bytes_per_row as usize, self.padded_bytes_per_row as usize, self.source_size) + } +} + +impl RenderAsset for ImageExportSource { + type ExtractedAsset = Self; + type Param = (SRes, SRes>); + type PreparedAsset = GpuImageExportSource; + + fn extract_asset(&self) -> Self::ExtractedAsset { + self.clone() + } + + fn prepare_asset( + extracted_asset: Self::ExtractedAsset, + (device, images): &mut SystemParamItem, + ) -> Result> { + let gpu_image = images.get(&extracted_asset.0).unwrap(); + + let size = gpu_image.texture.size(); + let format = &gpu_image.texture_format; + let bytes_per_row = + (size.width / format.block_dimensions().0) * format.block_size(None).unwrap(); + let padded_bytes_per_row = + RenderDevice::align_copy_bytes_per_row(bytes_per_row as usize) as u32; + + let source_size = gpu_image.texture.size(); + + Ok(GpuImageExportSource { + buffer: device.create_buffer(&BufferDescriptor { + label: Some("Image Export Buffer"), + size: (source_size.height * padded_bytes_per_row) as u64, + usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ, + mapped_at_creation: false, + }), + source_handle: extracted_asset.0.clone(), + source_size, + bytes_per_row, + padded_bytes_per_row, + }) + } +} + +#[derive(Component, Clone)] +pub struct ImageExportStartFrame(u64); + +impl Default for ImageExportSettings { + fn default() -> Self { + Self { extension: "png".into() } + } +} + +impl ExtractComponent for ImageExportSettings { + type Filter = (); + type Out = (Self, Handle, ImageExportStartFrame); + type Query = + (&'static Self, &'static Handle, &'static ImageExportStartFrame); + + fn extract_component( + (settings, source_handle, start_frame): QueryItem<'_, Self::Query>, + ) -> Option { + Some((settings.clone(), source_handle.clone_weak(), start_frame.clone())) + } +} + +fn setup_exporters( + mut commands: Commands, + exporters: Query, Without)>, + mut frame_id: Local, +) { + *frame_id = frame_id.wrapping_add(1); + for entity in &exporters { + commands.entity(entity).insert(ImageExportStartFrame(*frame_id)); + } +} + +#[derive(Bundle, Default)] +pub struct ImageExportBundle { + pub source: Handle, + pub settings: ImageExportSettings, +} + +fn save_buffer_as_resource( + export_bundles: Query<( + &Handle, + &ImageExportSettings, + &ImageExportStartFrame, + )>, + sources: Res>, + render_device: Res, + mut curr_img: ResMut, + mut frame_id: Local, +) { + *frame_id = frame_id.wrapping_add(1); + + log::debug!("num of export bundles {}", export_bundles.iter().len()); + + for (source_handle, settings, start_frame) in &export_bundles { + if let Some(gpu_source) = sources.get(source_handle) { + let mut image_bytes = { + let slice = gpu_source.buffer.slice(..); + + { + let (mapping_tx, mapping_rx) = oneshot::channel(); + + render_device.map_buffer(&slice, MapMode::Read, move |res| { + mapping_tx.send(res).unwrap(); + }); + + render_device.poll(Maintain::Wait); + futures_lite::future::block_on(mapping_rx).unwrap().unwrap(); + } + + slice.get_mapped_range().to_vec() + }; + + gpu_source.buffer.unmap(); + + let settings = settings.clone(); + let frame_id = *frame_id - start_frame.0 + 1; + let (bytes_per_row, padded_bytes_per_row, source_size) = gpu_source.get_bps(); + + if bytes_per_row != padded_bytes_per_row { + let mut unpadded_bytes = + Vec::::with_capacity(source_size.height as usize * bytes_per_row); + + for padded_row in image_bytes.chunks(padded_bytes_per_row) { + unpadded_bytes.extend_from_slice(&padded_row[..bytes_per_row]); + } + + image_bytes = unpadded_bytes; + } + + let extension = settings.extension.as_str(); + match extension { + "exr" => { + capture_img_bytes::>( + bytemuck::cast_slice(&image_bytes), + &source_size, + &mut curr_img, + frame_id, + extension, + ); + }, + _ => { + capture_img_bytes::>( + &image_bytes, + &source_size, + &mut curr_img, + frame_id, + extension, + ); + }, + } + } + } +} + +fn capture_img_bytes( + image_bytes: &[P::Subpixel], + source_size: &Extent3d, + curr_img: &mut ResMut, + frame_id: u64, + extension: &str, +) where + P::Subpixel: AnyBitPattern, + [P::Subpixel]: EncodableLayout, +{ + match ImageBuffer::::from_raw(source_size.width, source_size.height, image_bytes) { + Some(image_bytes) => { + curr_img.0.lock().update_data(frame_id, &image_bytes, extension.to_owned()); + + // if let Some(curr_img) = curr_img.as_mut(){ + // curr_img.update_data(frame_id, &image_bytes, extension.to_owned()); + // }else{ + // log::info!("Curr Image Resource still doesn't exist"); + // } + }, + None => { + log::error!("Failed creating image buffer for frame - '{frame_id}'"); + }, + } +} + +/// Plugin enabling the generation of image sequences. +#[derive(Default)] +pub struct ImageExportPlugin; + +#[derive(Debug, Hash, PartialEq, Eq, Clone, SystemSet)] +pub enum ImageExportSystems { + SetupImageExport, + SetupImageExportFlush, +} + +impl Plugin for ImageExportPlugin { + fn build(&self, app: &mut App) { + app.add_plugins( + DefaultPlugins + .set(ImagePlugin::default_nearest()) + .set(WindowPlugin { + primary_window: None, + exit_condition: bevy::window::ExitCondition::DontExit, + close_when_requested: false, + }) + .disable::(), + ); + + // TODO: + let curr_image_container = CurrImageContainer::default(); + + app.insert_resource(curr_image_container.clone()); + + app.configure_sets( + PostUpdate, + (SetupImageExport, SetupImageExportFlush).chain().before(CameraUpdateSystem), + ) + .register_type::() + .init_asset::() + .register_asset_reflect::() + .add_plugins(( + RenderAssetPlugin::::default(), + ExtractComponentPlugin::::default(), + )) + .add_systems( + PostUpdate, + ( + setup_exporters.in_set(SetupImageExport), + apply_deferred.in_set(SetupImageExportFlush), + ), + ); + + let render_app = app.sub_app_mut(RenderApp); + + render_app.insert_resource(curr_image_container); + + render_app.add_systems( + Render, + save_buffer_as_resource.after(RenderSet::Render).before(RenderSet::Cleanup), + ); + + let mut graph = render_app.world.get_resource_mut::().unwrap(); + + graph.add_node(NODE_NAME, ImageExportNode); + graph.add_node_edge(CAMERA_DRIVER, NODE_NAME); + } +} + +#[derive(Clone, Default, Resource)] +pub struct CurrImageContainer(pub Arc>); diff --git a/crates/bevy_headless/src/utils.rs b/crates/bevy_headless/src/utils.rs new file mode 100644 index 0000000..c8cba58 --- /dev/null +++ b/crates/bevy_headless/src/utils.rs @@ -0,0 +1,138 @@ +use bevy::{ + asset::Assets, + ecs::{ + event::Event, + system::{Commands, ResMut, Resource}, + }, + render::{ + camera::RenderTarget, + render_resource::{ + Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages, + }, + texture::Image, + }, +}; +use std::{io::Cursor, ops::Deref}; + +use base64::{engine::general_purpose, Engine}; +use image::{EncodableLayout, ImageBuffer, ImageOutputFormat, Pixel, Rgba, RgbaImage}; + +use crate::{ImageExportBundle, ImageExportSource}; + +#[derive(Default, Resource)] +pub struct CurrImage { + pub img_buffer: ImageBuffer, Vec>, + pub frame_id: u64, + pub extension: String, +} + +impl CurrImage { + pub fn update_data( + &mut self, + frame_id: u64, + image_bytes: &ImageBuffer, + extension: String, + ) where + P: Pixel + image::PixelWithColorType, + [P::Subpixel]: EncodableLayout, + Container: Deref, + { + self.frame_id = frame_id; + + self.extension = extension; + + let (w, h) = image_bytes.dimensions(); + if let Some(rgba_img_buff) = RgbaImage::from_raw(w, h, image_bytes.as_bytes().to_owned()) { + self.img_buffer = rgba_img_buff; + } else { + log::error!("Error updating curr image image buffer"); + }; + } + + pub fn create_path(&self, dir: &str) -> String { + // shouldn't be in loop, remove later + std::fs::create_dir_all(dir).expect("Output path could not be created"); + + format!("{dir}/{:06}.{}", self.frame_id, self.extension) + } + + pub fn to_web_base64(&self) -> anyhow::Result { + base64_browser_img(&self.img_buffer) + } +} + +#[derive(Debug, Default, Resource, Event)] +pub struct SceneInfo { + width: u32, + height: u32, +} + +impl SceneInfo { + pub fn new(width: u32, height: u32) -> SceneInfo { + SceneInfo { width, height } + } + + pub fn dimensions(&self) -> (u32, u32) { + (self.width, self.height) + } +} + +pub fn setup_render_target( + commands: &mut Commands, + images: &mut ResMut>, + scene_controller: &mut ResMut, + mut export_sources: ResMut>, +) -> 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 + // ?? remove ?? + | 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); + + commands.spawn(ImageExportBundle { + source: export_sources.add(render_target_image_handle.clone().into()), + ..Default::default() + }); + + RenderTarget::Image(render_target_image_handle) +} + +fn base64_browser_img(img: &ImageBuffer) -> anyhow::Result +where + P: Pixel + image::PixelWithColorType, + [P::Subpixel]: EncodableLayout, + Container: Deref, +{ + let mut image_data: Vec = Vec::new(); + img.write_to(&mut Cursor::new(&mut image_data), ImageOutputFormat::WebP)?; + let res_base64 = general_purpose::STANDARD.encode(image_data); + Ok(format!("data:image/webp;base64,{}", res_base64)) +} + +fn white_img_placeholder(w: u32, h: u32) -> String { + let img = RgbaImage::new(w, h); + + // img.iter_mut().for_each(|pixel| *pixel = 255); + base64_browser_img(&img).unwrap() +} diff --git a/crates/bevy_ws_server/Cargo.toml b/crates/bevy_ws_server/Cargo.toml index 459b623..250cb2f 100644 --- a/crates/bevy_ws_server/Cargo.toml +++ b/crates/bevy_ws_server/Cargo.toml @@ -18,4 +18,3 @@ 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 index 3f2afd7..644264c 100644 --- a/crates/bevy_ws_server/src/lib.rs +++ b/crates/bevy_ws_server/src/lib.rs @@ -1,11 +1,17 @@ +// Derived from https://github.com/kiwiyou/bevy-ws-server + use async_net::SocketAddr; use bevy::{ - prelude::*, + app::{App, Plugin, Update}, + ecs::{ + component::Component, + system::{Commands, ResMut, Resource}, + }, tasks::{IoTaskPool, Task}, }; use crossbeam_channel::{Receiver, Sender}; -use futures::{pin_mut, select, FutureExt}; +use futures::{pin_mut, select, FutureExt, Stream}; use std::{ net::{TcpListener, TcpStream}, @@ -18,7 +24,7 @@ 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 smol::{stream::StreamExt, Async}; use tungstenite::Message; pub struct WsPlugin; @@ -43,7 +49,7 @@ pub struct WsAcceptQueue { } /// A WebSocket or WebSocket+TLS connection. -enum WsStream { +pub enum WsStream { /// A plain WebSocket connection. Plain(WebSocketStream>), diff --git a/demo/components/WebSocketExample.tsx b/demo/components/WebSocketExample.tsx index 089ab80..1176466 100644 --- a/demo/components/WebSocketExample.tsx +++ b/demo/components/WebSocketExample.tsx @@ -45,7 +45,7 @@ export default function WebSocketExample({port = 8080}:{port?:number}){ return (
{imgUrl ? ( -
+
>, - mut scene_controller: ResMut, - render_device: Res, + mut scene_controller: ResMut, + export_sources: ResMut>, ) { - let render_target = bevy_frame_capture::setup_render_target( + let render_target = bevy_headless::setup_render_target( &mut commands, &mut images, - &render_device, &mut scene_controller, - 15, - String::from("main_scene"), + export_sources, ); commands.spawn(Camera3dBundle { @@ -26,21 +31,38 @@ fn setup( }); } -fn headless_app() { - App::new() - .insert_resource(bevy_frame_capture::SceneInfo::new(1920, 1080)) - .insert_resource(ClearColor(Color::rgb_u8(0, 0, 0))) - .add_plugins(DefaultPlugins.set(ImagePlugin::default_nearest()).set(WindowPlugin { - primary_window: None, - exit_condition: bevy::window::ExitCondition::DontExit, - close_when_requested: false, - })) - .add_plugins(bevy_frame_capture::FrameCapturePlugin) - .add_plugins(ScheduleRunnerPlugin::run_loop(std::time::Duration::from_secs_f64(1.0 / 60.0))) - .add_systems(Startup, setup) - .run() +fn save_img(curr_img: Res) { + let curr_img = curr_img.0.lock(); + if !curr_img.extension.is_empty() { + let path = curr_img.create_path("out"); + log::info!("path is {path}"); + let img = curr_img.img_buffer.clone(); + std::thread::spawn(move || { + if let Err(e) = img.save(path) { + log::error!("Couldn't save image | {e:?}"); + }; + }); + } } pub fn main() { - headless_app(); + pretty_env_logger::formatted_builder() + .filter_module("minimal", log::LevelFilter::Info) + .filter_module("bevy", log::LevelFilter::Info) + .filter_module("bevy_headless", log::LevelFilter::Info) + .init(); + + let (w, h) = (1920, 1080); + + Engine::new() + // .insert_resource(CurrImage::default()) + .insert_resource(bevy_headless::SceneInfo::new(w, h)) + .insert_resource(ClearColor(Color::rgb_u8(0, 0, 0))) + .add_plugins(( + ImageExportPlugin, + ScheduleRunnerPlugin::run_loop(std::time::Duration::from_secs_f64(1.0 / 60.0)), + )) + .add_systems(Startup, setup) + .add_systems(Update, save_img) + .run(); } diff --git a/examples/new_media/Cargo.toml b/examples/new_media/Cargo.toml index a5f523d..ae04407 100644 --- a/examples/new_media/Cargo.toml +++ b/examples/new_media/Cargo.toml @@ -11,7 +11,7 @@ async-net = {workspace = true} base64 = {workspace = true} bevy = { workspace = true } bevy_ws_server = { path = "../../crates/bevy_ws_server" } -bevy_frame_capture = { path = "../../crates/bevy_frame_capture" } +bevy_headless = { path = "../../crates/bevy_headless" } bevy_gaussian_splatting = { version = "2.0.2", default-features = true } bevy_panorbit_camera = { workspace = true } crossbeam-channel = { workspace = true } diff --git a/examples/new_media/src/main.rs b/examples/new_media/src/main.rs index 299691f..83c10c7 100644 --- a/examples/new_media/src/main.rs +++ b/examples/new_media/src/main.rs @@ -1,41 +1,48 @@ mod controls; -// mod server; - -// use bevy_ws_server::WsPlugin; +mod server; use bevy::{ - app::ScheduleRunnerPlugin, core::Name, core_pipeline::tonemapping::Tonemapping, log::LogPlugin, - prelude::*, render::renderer::RenderDevice, + app::{App as Engine, ScheduleRunnerPlugin, Startup, Update}, + asset::{AssetServer, Assets}, + core::Name, + core_pipeline::{clear_color::ClearColor, core_3d::Camera3dBundle, tonemapping::Tonemapping}, + ecs::{ + query::With, + system::{Commands, Query, Res, ResMut}, + }, + math::Vec3, + render::{camera::Camera, color::Color, texture::Image}, + transform::components::Transform, + utils::default, }; +use bevy_headless::{ImageExportPlugin, ImageExportSource}; use bevy_panorbit_camera::{PanOrbitCamera, PanOrbitCameraPlugin}; use bevy_gaussian_splatting::{GaussianCloud, GaussianSplattingBundle, GaussianSplattingPlugin}; - -// use server::{receive_message, start_ws}; +use bevy_ws_server::WsPlugin; +use server::{receive_message, start_ws}; fn setup_gaussian_cloud( mut commands: Commands, - _asset_server: Res, - mut gaussian_assets: ResMut>, - mut scene_controller: ResMut, + asset_server: Res, + _gaussian_assets: ResMut>, + mut scene_controller: ResMut, mut images: ResMut>, - render_device: Res, + export_sources: ResMut>, ) { // let remote_file = Some("https://huggingface.co/datasets/cs50victor/splats/resolve/main/train/point_cloud/iteration_7000/point_cloud.gcloud"); // TODO: figure out how to load remote files later - // let splat_file = "splats/bonsai/point_cloud/iteration_7000/point_cloud.gcloud"; - // log::info!("loading {}", splat_file); - // let cloud = asset_server.load(splat_file.to_string()); + let splat_file = "splats/bonsai/point_cloud/iteration_7000/point_cloud.gcloud"; + log::info!("loading {}", splat_file); + let cloud = asset_server.load(splat_file.to_string()); - let cloud = gaussian_assets.add(GaussianCloud::test_model()); + // let cloud = gaussian_assets.add(GaussianCloud::test_model()); - let render_target = bevy_frame_capture::setup_render_target( + let render_target = bevy_headless::setup_render_target( &mut commands, &mut images, - &render_device, &mut scene_controller, - 15, - String::from("main_scene"), + export_sources, ); let gs = GaussianSplattingBundle { cloud, ..default() }; @@ -70,43 +77,33 @@ fn main() { pretty_env_logger::formatted_builder() .filter_module("new_media", log::LevelFilter::Info) .filter_module("bevy", log::LevelFilter::Info) + .filter_module("bevy_headless", log::LevelFilter::Info) .filter_module("bevy_ws_server", log::LevelFilter::Info) .filter_module("bevy_ws_server", log::LevelFilter::Debug) .init(); let config = AppConfig { width: 1920, height: 1080 }; - App::new() - .insert_resource(bevy_frame_capture::SceneInfo::new(config.width, config.height)) - .insert_resource(bevy_frame_capture::CurrImageBase64(bevy_frame_capture::white_img_placeholder(config.width, config.height))) - .insert_resource(ClearColor(Color::rgb_u8(0, 0, 0))) - .add_plugins(( - DefaultPlugins - .set(ImagePlugin::default_nearest()) - .set(WindowPlugin { - primary_window: None, - exit_condition: bevy::window::ExitCondition::DontExit, - close_when_requested: false, - }).disable::(), - // WsPlugin, - bevy_frame_capture::FrameCapturePlugin, - ScheduleRunnerPlugin::run_loop(std::time::Duration::from_secs_f64(1.0 / 60.0)), - PanOrbitCameraPlugin, - GaussianSplattingPlugin, - )) - .add_systems(Startup, setup_gaussian_cloud) - .add_systems(Update, ( - move_camera, - // receive_message - )) - // .add_systems(OnEnter(AppState::Active), setup_gaussian_cloud) - .run(); + Engine::new() + .insert_resource(bevy_headless::SceneInfo::new(config.width, config.height)) + .insert_resource(ClearColor(Color::rgb_u8(255, 255, 255))) + .add_plugins(( + ImageExportPlugin, + WsPlugin, + ScheduleRunnerPlugin::run_loop(std::time::Duration::from_secs_f64(1.0 / 60.0)), + PanOrbitCameraPlugin, + GaussianSplattingPlugin, + )) + .add_systems(Startup, (start_ws, setup_gaussian_cloud)) + .add_systems(Update, (move_camera, receive_message)) + .run(); } fn move_camera(mut camera: Query<&mut Transform, With>) { + let speed = 0.0005; for mut transform in camera.iter_mut() { - transform.translation.x += 0.0005; - transform.translation.y += 0.0005; - transform.translation.z += 0.0005; + transform.translation.x += speed; + transform.translation.y += speed; + transform.translation.z += speed; } } diff --git a/examples/new_media/src/server.rs b/examples/new_media/src/server.rs index 88a1e99..2aadf07 100644 --- a/examples/new_media/src/server.rs +++ b/examples/new_media/src/server.rs @@ -1,14 +1,16 @@ +use bevy::ecs::{ + entity::Entity, + system::{Commands, Query, Res}, +}; +use bevy_headless::CurrImageContainer; use bevy_ws_server::{ReceiveError, WsConnection, WsListener}; use log::info; use serde_json::json; use std::fmt::Debug; -use bevy::prelude::*; use serde::{Deserialize, Serialize}; -use bevy_frame_capture::CurrImageBase64; - #[derive(Debug, Serialize, Deserialize)] pub struct HttpServerMsg { data: Option, @@ -45,7 +47,7 @@ pub fn start_ws(listener: Res) { pub fn receive_message( mut commands: Commands, - curr_base64_img: Res, + curr_image: Res, connections: Query<(Entity, &WsConnection)>, ) { for (entity, conn) in connections.iter() { @@ -55,7 +57,7 @@ pub fn receive_message( info!("message | {message:?}"); let resp = tungstenite::protocol::Message::Text( json!({ - "image": curr_base64_img.0 + "image": curr_image.0.lock().to_web_base64().unwrap() }) .to_string(), );