From 6fb325adcc854e4f088b3cd59cc43dfd9e122136 Mon Sep 17 00:00:00 2001 From: ArthurBrussee Date: Tue, 29 Oct 2024 14:52:16 +0000 Subject: [PATCH] Fix exporting ply files on native --- crates/brush-dataset/src/lib.rs | 1 + crates/brush-desktop/src/main.rs | 1 + crates/brush-viewer/src/async_lib.rs | 88 +++++++++++++++++++++++ crates/brush-viewer/src/lib.rs | 2 +- crates/brush-viewer/src/panels/scene.rs | 33 ++++++--- crates/brush-viewer/src/timeout_future.rs | 63 ---------------- crates/brush-viewer/src/viewer.rs | 27 +++---- crates/rrfd/src/android.rs | 9 ++- crates/rrfd/src/lib.rs | 85 +++++++++++++--------- 9 files changed, 182 insertions(+), 127 deletions(-) create mode 100644 crates/brush-viewer/src/async_lib.rs delete mode 100644 crates/brush-viewer/src/timeout_future.rs diff --git a/crates/brush-dataset/src/lib.rs b/crates/brush-dataset/src/lib.rs index 87d6e4e6..19cf2030 100644 --- a/crates/brush-dataset/src/lib.rs +++ b/crates/brush-dataset/src/lib.rs @@ -167,6 +167,7 @@ mod async_helpers { pub(super) fn spawn_future( future: impl Future + 'static, ) -> JoinHandle { + // On wasm, just spawn locally. task::spawn_local(future) } } diff --git a/crates/brush-desktop/src/main.rs b/crates/brush-desktop/src/main.rs index a2df28ad..ea8ab83d 100644 --- a/crates/brush-desktop/src/main.rs +++ b/crates/brush-desktop/src/main.rs @@ -42,6 +42,7 @@ fn main() -> anyhow::Result<()> { .dyn_into::() .unwrap(); + // On wasm, run as a local task. async_std::task::spawn_local(async { eframe::WebRunner::new() .start( diff --git a/crates/brush-viewer/src/async_lib.rs b/crates/brush-viewer/src/async_lib.rs new file mode 100644 index 00000000..70b8558d --- /dev/null +++ b/crates/brush-viewer/src/async_lib.rs @@ -0,0 +1,88 @@ +use std::future::Future; +use web_time::Duration; +use wgpu::WasmNotSend; + +/// Wrap a future to yield back to the browser macrotasks +/// queue each time it yields (using a setTimeout). +/// +/// In the future this could work better with scheduler.yield. +pub fn with_timeout_yield + WasmNotSend + 'static, T: Send + 'static>( + future: F, + min_time_between_yields: Duration, +) -> impl Future + WasmNotSend + 'static { + #[cfg(not(target_family = "wasm"))] + { + let _ = min_time_between_yields; + future + } + + #[cfg(target_family = "wasm")] + { + use gloo_timers::future::TimeoutFuture; + use std::{ + pin::Pin, + task::{Context, Poll}, + }; + use web_time::Instant; + + struct WithTimeoutYield { + future: Pin>, + timeout: Option, + last_yield: Option, + min_duration: Duration, + } + + impl Future for WithTimeoutYield { + type Output = F::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + + if let Some(timeout) = &mut this.timeout { + match std::pin::pin!(timeout).poll(cx) { + Poll::Ready(_) => { + this.timeout = None; + this.last_yield = Some(Instant::now()); + } + Poll::Pending => return Poll::Pending, + } + } + + match this.future.as_mut().poll(cx) { + Poll::Ready(output) => Poll::Ready(output), + Poll::Pending => { + let should_yield = this + .last_yield + .map(|t| t.elapsed() >= this.min_duration) + .unwrap_or(true); + + if should_yield && this.timeout.is_none() { + this.timeout = Some(TimeoutFuture::new(0)); + } + Poll::Pending + } + } + } + } + + WithTimeoutYield { + future: Box::pin(future), + timeout: None, + last_yield: None, + min_duration: min_time_between_yields, + } + } +} + +pub fn spawn_future( + fut: impl Future + WasmNotSend + 'static, +) { + #[cfg(target_family = "wasm")] + { + async_std::task::spawn_local(fut); + } + #[cfg(not(target_family = "wasm"))] + { + async_std::task::spawn(fut); + } +} diff --git a/crates/brush-viewer/src/lib.rs b/crates/brush-viewer/src/lib.rs index ece860f7..3d6905d2 100644 --- a/crates/brush-viewer/src/lib.rs +++ b/crates/brush-viewer/src/lib.rs @@ -5,9 +5,9 @@ use egui_tiles::SimplificationOptions; use viewer::{ViewerContext, ViewerMessage}; +mod async_lib; mod burn_texture; mod orbit_controls; -mod timeout_future; mod panels; mod train_loop; diff --git a/crates/brush-viewer/src/panels/scene.rs b/crates/brush-viewer/src/panels/scene.rs index ebdf15cc..4a55f634 100644 --- a/crates/brush-viewer/src/panels/scene.rs +++ b/crates/brush-viewer/src/panels/scene.rs @@ -1,4 +1,3 @@ -use async_std::task; use brush_dataset::splat_export; use egui::epaint::mutex::RwLock as EguiRwLock; use std::sync::Arc; @@ -243,20 +242,32 @@ runs consider using the native app."#, ui.add_space(15.0); - if ui.button("↑ Export").clicked() { + if ui.button("⬆ Export").clicked() { let splats = splats.clone(); - task::spawn_local(async move { - let data = splat_export::splat_to_ply(*splats).await; - let data = match data { - Ok(data) => data, + + crate::async_lib::spawn_future(async move { + let file = rrfd::save_file("export.ply").await; + + // Not sure where/how to show this error if any. + match file { Err(e) => { log::error!("Failed to save file: {e}"); - return; } - }; - // Not sure where/how to show this error if any. - if let Err(e) = rrfd::save_file("export.ply", data).await { - log::error!("Failed to save file: {e}"); + Ok(file) => { + let data = splat_export::splat_to_ply(*splats).await; + + let data = match data { + Ok(data) => data, + Err(e) => { + log::error!("Failed to serialize file: {e}"); + return; + } + }; + + if let Err(e) = file.write(&data).await { + log::error!("Failed to write file: {e}"); + } + } } }); } diff --git a/crates/brush-viewer/src/timeout_future.rs b/crates/brush-viewer/src/timeout_future.rs deleted file mode 100644 index 988977e0..00000000 --- a/crates/brush-viewer/src/timeout_future.rs +++ /dev/null @@ -1,63 +0,0 @@ -use gloo_timers::future::TimeoutFuture; -use std::{ - future::Future, - pin::Pin, - task::{Context, Poll}, -}; -use web_time::{Duration, Instant}; - -/// Wrap a future to yield back to the browser macrotasks -/// queue each time it yields (using a setTimeout). -/// -/// In the future this could work better with scheduler.yield. -pub fn with_timeout_yield( - future: F, - min_time_between_yields: Duration, -) -> impl Future { - struct WithTimeoutYield { - future: Pin>, - timeout: Option, - last_yield: Option, - min_duration: Duration, - } - - impl Future for WithTimeoutYield { - type Output = F::Output; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.get_mut(); - - if let Some(timeout) = &mut this.timeout { - match std::pin::pin!(timeout).poll(cx) { - Poll::Ready(_) => { - this.timeout = None; - this.last_yield = Some(Instant::now()); - } - Poll::Pending => return Poll::Pending, - } - } - - match this.future.as_mut().poll(cx) { - Poll::Ready(output) => Poll::Ready(output), - Poll::Pending => { - let should_yield = this - .last_yield - .map(|t| t.elapsed() >= this.min_duration) - .unwrap_or(true); - - if should_yield && this.timeout.is_none() { - this.timeout = Some(TimeoutFuture::new(0)); - } - Poll::Pending - } - } - } - } - - WithTimeoutYield { - future: Box::pin(future), - timeout: None, - last_yield: None, - min_duration: min_time_between_yields, - } -} diff --git a/crates/brush-viewer/src/viewer.rs b/crates/brush-viewer/src/viewer.rs index a206a886..85410121 100644 --- a/crates/brush-viewer/src/viewer.rs +++ b/crates/brush-viewer/src/viewer.rs @@ -4,7 +4,6 @@ use async_fn_stream::try_fn_stream; use async_std::{ channel::{Receiver, Sender, TrySendError}, stream::{Stream, StreamExt}, - task, }; use brush_dataset::{self, splat_import, Dataset, LoadDatasetArgs, LoadInitArgs, ZipData}; use brush_render::camera::Camera; @@ -20,6 +19,7 @@ use glam::{Quat, Vec3}; use web_time::Instant; use crate::{ + async_lib, orbit_controls::OrbitControls, panels::{DatasetPanel, LoadDataPanel, PresetsPanel, ScenePanel, StatsPanel}, train_loop::{self, TrainMessage}, @@ -97,12 +97,14 @@ fn process_loop( let stream = try_fn_stream(|emitter| async move { let _ = emitter.emit(ViewerMessage::PickFile).await; let picked = rrfd::pick_file().await?; + let name = picked.file_name(); + + if name.contains(".ply") { + let data = picked.read().await; - if picked.file_name.contains(".ply") { let _ = emitter .emit(ViewerMessage::StartLoading { training: false }) .await; - let data = picked.data; let splat_stream = splat_import::load_splat_from_ply::(data, device.clone()); let mut splat_stream = std::pin::pin!(splat_stream); @@ -114,13 +116,15 @@ fn process_loop( }) .await; } - } else if picked.file_name.contains(".zip") { + } else if name.contains(".zip") { + let data = picked.read().await; + let _ = emitter .emit(ViewerMessage::StartLoading { training: true }) .await; let stream = train_loop::train_loop( - ZipData::from(picked.data), + ZipData::from(data), device, train_receiver, load_data_args, @@ -215,17 +219,8 @@ impl ViewerContext { } }; - #[cfg(target_family = "wasm")] - { - let fut = - crate::timeout_future::with_timeout_yield(fut, web_time::Duration::from_millis(5)); - task::spawn_local(fut); - } - - #[cfg(not(target_family = "wasm"))] - { - task::spawn(fut); - } + let fut = crate::async_lib::with_timeout_yield(fut, web_time::Duration::from_millis(5)); + async_lib::spawn_future(fut); } pub fn send_train_message(&self, message: TrainMessage) { diff --git a/crates/rrfd/src/android.rs b/crates/rrfd/src/android.rs index 548f33a6..5998448b 100644 --- a/crates/rrfd/src/android.rs +++ b/crates/rrfd/src/android.rs @@ -1,4 +1,3 @@ -use super::PickedFile; use anyhow::Result; use async_std::channel::Sender; use jni::objects::{GlobalRef, JByteArray, JClass, JStaticMethodID, JString}; @@ -8,6 +7,12 @@ use lazy_static::lazy_static; use std::sync::Arc; use std::sync::RwLock; +#[derive(Clone, Debug)] +pub struct PickedFile { + pub data: Vec, + pub file_name: String, +} + lazy_static! { static ref VM: RwLock>> = RwLock::new(None); static ref CHANNEL: RwLock>>> = RwLock::new(None); @@ -15,6 +20,7 @@ lazy_static! { static ref FILE_PICKER_CLASS: RwLock> = RwLock::new(None); } +#[allow(unused)] pub fn jni_initialize(vm: Arc) { let mut env = vm.get_env().expect("Cannot get reference to the JNIEnv"); let class = env.find_class("com/splats/app/FilePicker").unwrap(); @@ -30,6 +36,7 @@ pub fn jni_initialize(vm: Arc) { *VM.write().unwrap() = Some(vm); } +#[allow(unused)] pub(crate) async fn pick_file() -> Result { let (sender, receiver) = async_std::channel::bounded(1); { diff --git a/crates/rrfd/src/lib.rs b/crates/rrfd/src/lib.rs index 10a8ee4f..d55fabc1 100644 --- a/crates/rrfd/src/lib.rs +++ b/crates/rrfd/src/lib.rs @@ -1,62 +1,77 @@ -#[cfg(target_os = "android")] pub mod android; +#[allow(unused)] +use anyhow::Context; use anyhow::Result; -#[derive(Clone, Debug)] -pub struct PickedFile { - pub data: Vec, - pub file_name: String, +pub enum FileHandle { + #[cfg(not(target_os = "android"))] + Rfd(rfd::FileHandle), + Android(android::PickedFile), +} + +impl FileHandle { + pub fn file_name(&self) -> String { + match self { + #[cfg(not(target_os = "android"))] + FileHandle::Rfd(file_handle) => file_handle.file_name(), + FileHandle::Android(picked_file) => picked_file.file_name.clone(), + } + } + + pub async fn write(&self, data: &[u8]) -> std::io::Result<()> { + match self { + #[cfg(not(target_os = "android"))] + FileHandle::Rfd(file_handle) => file_handle.write(data).await, + FileHandle::Android(_) => { + let _ = data; + unimplemented!("No saving on Android yet.") + } + } + } + + pub async fn read(&self) -> Vec { + match self { + #[cfg(not(target_os = "android"))] + FileHandle::Rfd(file_handle) => file_handle.read().await, + FileHandle::Android(picked_file) => picked_file.data.clone(), + } + } } /// Pick a file and return the name & bytes of the file. -pub async fn pick_file() -> Result { +pub async fn pick_file() -> Result { #[cfg(not(target_os = "android"))] { - async move { - use anyhow::Context; - let file = rfd::AsyncFileDialog::new() - .pick_file() - .await - .context("No file selected")?; - let file_data = file.read().await; - Ok(PickedFile { - data: file_data, - file_name: file.file_name(), - }) - } - .await - } + let file = rfd::AsyncFileDialog::new() + .pick_file() + .await + .context("No file selected")?; + Ok(FileHandle::Rfd(file)) + } #[cfg(target_os = "android")] { - android::pick_file().await + android::pick_file().await.map(FileHandle::Android) } } /// Saves data to a file and returns the filename the data was saved too. /// /// Nb: Does not work on Android currently. -pub async fn save_file(default_name: &str, data: Vec) -> Result { +pub async fn save_file(default_name: &str) -> Result { #[cfg(not(target_os = "android"))] { - async move { - use anyhow::Context; - let file = rfd::AsyncFileDialog::new() - .set_file_name(default_name) - .save_file() - .await - .context("No file selected")?; - file.write(&data).await?; - Ok(file.file_name()) - } - .await + let file = rfd::AsyncFileDialog::new() + .set_file_name(default_name) + .save_file() + .await + .context("No file selected")?; + Ok(FileHandle::Rfd(file)) } - #[cfg(target_os = "android")] { let _ = default_name; - let _ = data; unimplemented!("No saving on Android yet.") } }