diff --git a/README.md b/README.md index 457c1b8a..f854cee3 100644 --- a/README.md +++ b/README.md @@ -16,8 +16,9 @@ _NOTE: This only works on desktop Chrome 129+ currently (Oct 2024). Firefox and ## Features The demo can load pretrained ply splats, and can load datasets to train on. Currently only two formats are supported. A .zip file containing: -- A transforms_train.json and images, like the synthetic nerf scene dataset. - An `images` & `sparse` folder with [`COLMAP`](https://github.com/colmap/colmap) data +- A .json and images, like the [nerfstudio format](https://docs.nerf.studio/quickstart/data_conventions.html). + - You can specify a custom transforms_train.json and transforms_eval.json split. While training you can interact with the scene and see the training dynamics live, and compare the current rendering to training / eval views as the training progresses. @@ -72,7 +73,7 @@ Brush is split into various crates. A quick overview of the different responsibi - `brush-viewer` handles the UI and integrating the training loop. - `brush-android` is the binary target for running on android, while `brush-desktop` is for running both on web, and mac/Windows/Linux. - `brush-wgsl` handles some kernel inspection for generating CPU-side structs and interacing with [naga-oil](https://github.com/bevyengine/naga_oil) to handle shader imports. -- `brush-dataset` handles importing different datasets like COLMAP or synthetic nerf data. +- `brush-dataset` handles importing different training data formats. - `brush-prefix-sum` and `brush-sort` are only compute kernels and should be largely independent of Brush (other than `brush-wgsl`). - `rrfd` is a small extension of [`rfd`](https://github.com/PolyMeilex/rfd) diff --git a/crates/brush-dataset/src/formats/mod.rs b/crates/brush-dataset/src/formats/mod.rs index 701b0fbc..2a953da3 100644 --- a/crates/brush-dataset/src/formats/mod.rs +++ b/crates/brush-dataset/src/formats/mod.rs @@ -15,10 +15,8 @@ pub fn load_dataset( load_args: &LoadDatasetArgs, device: &B::Device, ) -> anyhow::Result<(DataStream>, DataStream)> { - let streams = nerfstudio::read_dataset(archive.clone(), load_args, device).or_else(|e| { - log::info!("Not a NeRF synthetic dataset ({e}), trying to load as Colmap."); - colmap::load_dataset::(archive.clone(), load_args, device) - }); + let streams = nerfstudio::read_dataset(archive.clone(), load_args, device) + .or_else(|_| colmap::load_dataset::(archive.clone(), load_args, device)); let Ok(streams) = streams else { anyhow::bail!("Couldn't parse dataset as any format. Only some formats are supported.") diff --git a/crates/brush-dataset/src/formats/nerfstudio.rs b/crates/brush-dataset/src/formats/nerfstudio.rs index f18985d7..f2e2d269 100644 --- a/crates/brush-dataset/src/formats/nerfstudio.rs +++ b/crates/brush-dataset/src/formats/nerfstudio.rs @@ -17,7 +17,7 @@ use tokio_stream::StreamExt; #[derive(serde::Deserialize, Clone)] #[allow(unused)] // not reading camera distortions yet. -struct SyntheticScene { +struct JsonScene { // Simple synthetic nerf camera model. camera_angle_x: Option, // Not really used atm. @@ -92,7 +92,7 @@ struct FrameData { } fn read_transforms_file( - scene: SyntheticScene, + scene: JsonScene, transforms_path: PathBuf, archive: DatasetZip, load_args: &LoadDatasetArgs, @@ -175,11 +175,10 @@ pub fn read_dataset( load_args: &LoadDatasetArgs, device: &B::Device, ) -> Result<(DataStream>, DataStream)> { - log::info!("Loading nerf synthetic dataset"); + log::info!("Loading nerfstudio dataset"); let transforms_path = archive.find_with_extension(".json", "_train")?; - let train_scene: SyntheticScene = - serde_json::from_reader(archive.file_at_path(&transforms_path)?)?; + let train_scene: JsonScene = serde_json::from_reader(archive.file_at_path(&transforms_path)?)?; let train_handles = read_transforms_file( train_scene.clone(), transforms_path.clone(), diff --git a/crates/brush-viewer/src/panels/scene.rs b/crates/brush-viewer/src/panels/scene.rs index b29b29ad..2612642c 100644 --- a/crates/brush-viewer/src/panels/scene.rs +++ b/crates/brush-viewer/src/panels/scene.rs @@ -198,7 +198,7 @@ impl ViewerPanel for ScenePanel { Load a pretrained .ply file to view it Or load a dataset to train on. These are zip files with: - - a transforms_train.json and images, like the synthetic NeRF dataset format. + - a transforms.json and images, like the nerfstudio dataset format. - COLMAP data, containing the `images` & `sparse` folder."#, );