Add initial brush VFS

crates/brush-dataset/src/brush_vfs.rs

@@ -0,0 +1,132 @@
+//
+// This class helps working with an archive as a somewhat more regular filesystem.
+//
+// [1] really we want to just read directories.
+// The reason is that picking directories isn't supported on
+// rfd on wasm, nor is drag-and-dropping folders in egui.
+use std::{
+    collections::HashMap,
+    io::{Cursor, Read},
+    path::{Component, Path, PathBuf},
+    sync::Arc,
+};
+
+use anyhow::Context;
+use tokio::{io::AsyncRead, io::AsyncReadExt, sync::Mutex};
+
+use zip::{
+    result::{ZipError, ZipResult},
+    ZipArchive,
+};
+
+type DynRead = Box<dyn AsyncRead + Send + Unpin>;
+
+#[derive(Clone)]
+pub struct ZipData {
+    data: Arc<Vec<u8>>,
+}
+
+impl AsRef<[u8]> for ZipData {
+    fn as_ref(&self) -> &[u8] {
+        &self.data
+    }
+}
+
+pub(crate) fn normalized_path(path: &Path) -> PathBuf {
+    path.components()
+        .filter(|c| !matches!(c, Component::CurDir | Component::ParentDir))
+        .collect()
+}
+
+#[derive(Clone, Default)]
+pub struct PathReader {
+    paths: HashMap<PathBuf, Arc<Mutex<Option<DynRead>>>>,
+}
+
+impl PathReader {
+    fn paths(&self) -> impl Iterator<Item = &PathBuf> {
+        self.paths.keys()
+    }
+
+    pub fn add(&mut self, path: &Path, reader: impl AsyncRead + Send + Unpin + 'static) {
+        self.paths.insert(
+            path.to_path_buf(),
+            Arc::new(Mutex::new(Some(Box::new(reader)))),
+        );
+    }
+
+    async fn open(&mut self, path: &Path) -> anyhow::Result<Box<dyn AsyncRead + Send + Unpin>> {
+        let entry = self.paths.remove(path).context("File not found")?;
+        let reader = entry.lock().await.take();
+        reader.context("Missing reader")
+    }
+}
+
+#[derive(Clone)]
+pub enum BrushVfs {
+    Zip(ZipArchive<Cursor<ZipData>>),
+    Manual(PathReader),
+    #[cfg(not(target_family = "wasm"))]
+    Directory(PathBuf, Vec<PathBuf>),
+}
+
+// TODO: This is all awfully ad-hoc.
+impl BrushVfs {
+    pub async fn from_zip_reader(reader: impl AsyncRead + Unpin) -> ZipResult<Self> {
+        let mut bytes = vec![];
+        let mut reader = reader;
+        reader.read_to_end(&mut bytes).await?;
+
+        let zip_data = ZipData {
+            data: Arc::new(bytes),
+        };
+        let archive = ZipArchive::new(Cursor::new(zip_data))?;
+        Ok(BrushVfs::Zip(archive))
+    }
+
+    pub fn from_paths(paths: PathReader) -> Self {
+        BrushVfs::Manual(paths)
+    }
+
+    #[cfg(not(target_family = "wasm"))]
+    pub async fn from_directory(dir: &Path) -> anyhow::Result<Self> {
+        let mut read = ::tokio::fs::read_dir(dir).await?;
+        let mut paths = vec![];
+        while let Some(entry) = read.next_entry().await? {
+            paths.push(entry.path());
+        }
+        Ok(BrushVfs::Directory(dir.to_path_buf(), paths))
+    }
+
+    pub fn file_names(&self) -> impl Iterator<Item = &Path> + '_ {
+        let iterator: Box<dyn Iterator<Item = &Path>> = match self {
+            BrushVfs::Zip(archive) => Box::new(archive.file_names().map(Path::new)),
+            BrushVfs::Manual(map) => Box::new(map.paths().map(|p| p.as_path())),
+            #[cfg(not(target_family = "wasm"))]
+            BrushVfs::Directory(_, paths) => Box::new(paths.iter().map(|p| p.as_path())),
+        };
+        // stupic macOS.
+        iterator.filter(|p| !p.starts_with("__MACOSX"))
+    }
+
+    pub async fn open_path(&mut self, path: &Path) -> anyhow::Result<DynRead> {
+        match self {
+            BrushVfs::Zip(archive) => {
+                let name = archive
+                    .file_names()
+                    .find(|name| path == Path::new(name))
+                    .ok_or(ZipError::FileNotFound)?;
+                let name = name.to_owned();
+                let mut buffer = vec![];
+                archive.by_name(&name)?.read_to_end(&mut buffer)?;
+                Ok(Box::new(Cursor::new(buffer)))
+            }
+            BrushVfs::Manual(map) => map.open(path).await,
+            #[cfg(not(target_family = "wasm"))]
+            BrushVfs::Directory(path_buf, _) => {
+                let file = tokio::fs::File::open(path_buf).await?;
+                Ok(Box::new(file))
+            }
+        }
+    }
+}

crates/brush-dataset/src/formats/colmap.rs

@@ -1,8 +1,16 @@
-use std::{future::Future, sync::Arc};
+use std::{
+    future::Future,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 
-use super::{DataStream, DatasetZip, LoadDatasetArgs};
-use crate::{splat_import::SplatMessage, stream_fut_parallel, Dataset};
-use anyhow::{Context, Result};
+use super::{DataStream, LoadDatasetArgs};
+use crate::{
+    brush_vfs::{normalized_path, BrushVfs},
+    splat_import::SplatMessage,
+    stream_fut_parallel, Dataset,
+};
+use anyhow::Result;
 use async_fn_stream::try_fn_stream;
 use brush_render::{
     camera::{self, Camera},
@@ -12,22 +20,37 @@ use brush_render::{
 };
 use brush_train::scene::SceneView;
 use glam::Vec3;
+use tokio::io::AsyncReadExt;
 use tokio_stream::StreamExt;
 
-fn read_views(
-    mut archive: DatasetZip,
+fn find_base_path(archive: &BrushVfs, search_path: &str) -> Option<PathBuf> {
+    for file in archive.file_names() {
+        let path = normalized_path(Path::new(file));
+        if path.ends_with(search_path) {
+            return path
+                .ancestors()
+                .nth(Path::new(search_path).components().count())
+                .map(|x| x.to_owned());
+        }
+    }
+    None
+}
+
+async fn read_views(
+    archive: BrushVfs,
     load_args: &LoadDatasetArgs,
 ) -> Result<Vec<impl Future<Output = Result<SceneView>>>> {
     log::info!("Loading colmap dataset");
+    let mut archive = archive;
 
-    let (is_binary, base_path) = if let Some(path) = archive.find_base_path("sparse/0/cameras.bin")
-    {
-        (true, path)
-    } else if let Some(path) = archive.find_base_path("sparse/0/cameras.txt") {
-        (false, path)
-    } else {
-        anyhow::bail!("No COLMAP data found (either text or binary.")
-    };
+    let (is_binary, base_path) =
+        if let Some(path) = find_base_path(&archive, "sparse/0/cameras.bin") {
+            (true, path)
+        } else if let Some(path) = find_base_path(&archive, "sparse/0/cameras.txt") {
+            (false, path)
+        } else {
+            anyhow::bail!("No COLMAP data found (either text or binary.)")
+        };
 
     let (cam_path, img_path) = if is_binary {
         (
@@ -42,14 +65,14 @@ fn read_views(
     };
 
     let cam_model_data = {
-        let mut cam_file = archive.file_at_path(&cam_path)?;
-        colmap_reader::read_cameras(&mut cam_file, is_binary)?
+        let mut cam_file = archive.open_path(&cam_path).await?;
+        colmap_reader::read_cameras(&mut cam_file, is_binary).await?
     };
 
     let img_infos = {
-        let img_file = archive.file_at_path(&img_path)?;
-        let mut buf_reader = std::io::BufReader::new(img_file);
-        colmap_reader::read_images(&mut buf_reader, is_binary)?
+        let img_file = archive.open_path(&img_path).await?;
+        let mut buf_reader = tokio::io::BufReader::new(img_file);
+        colmap_reader::read_images(&mut buf_reader, is_binary).await?
     };
 
     let mut img_info_list = img_infos.into_iter().collect::<Vec<_>>();
@@ -82,7 +105,12 @@ fn read_views(
 
                 let img_path = base_path.join(format!("images/{}", img_info.name));
 
-                let img_bytes = archive.read_bytes_at_path(&img_path)?;
+                let mut img_bytes = vec![];
+                archive
+                    .open_path(&img_path)
+                    .await?
+                    .read_to_end(&mut img_bytes)
+                    .await?;
                 let mut img = image::load_from_memory(&img_bytes)?;
 
                 if let Some(max) = load_args.max_resolution {
@@ -98,7 +126,7 @@ fn read_views(
                 let camera = Camera::new(translation, quat, fovx, fovy, center_uv);
 
                 let view = SceneView {
-                    name: img_path.to_str().context("Invalid file name")?.to_owned(),
+                    name: img_path.to_string_lossy().to_string(),
                     camera,
                     image: Arc::new(img),
                 };
@@ -110,12 +138,12 @@ fn read_views(
     Ok(handles)
 }
 
-pub(crate) fn load_dataset<B: Backend>(
-    mut archive: DatasetZip,
+pub(crate) async fn load_dataset<B: Backend>(
+    mut archive: BrushVfs,
     load_args: &LoadDatasetArgs,
     device: &B::Device,
 ) -> Result<(DataStream<SplatMessage<B>>, DataStream<Dataset>)> {
-    let mut handles = read_views(archive.clone(), load_args)?;
+    let mut handles = read_views(archive.clone(), load_args).await?;
 
     if let Some(subsample) = load_args.subsample_frames {
         handles = handles.into_iter().step_by(subsample as usize).collect();
@@ -149,9 +177,9 @@ pub(crate) fn load_dataset<B: Backend>(
 
     let init_stream = try_fn_stream(|emitter| async move {
         let (is_binary, base_path) =
-            if let Some(path) = archive.find_base_path("sparse/0/cameras.bin") {
+            if let Some(path) = find_base_path(&archive, "sparse/0/cameras.bin") {
                 (true, path)
-            } else if let Some(path) = archive.find_base_path("sparse/0/cameras.txt") {
+            } else if let Some(path) = find_base_path(&archive, "sparse/0/cameras.txt") {
                 (false, path)
             } else {
                 anyhow::bail!("No COLMAP data found (either text or binary.")
@@ -165,8 +193,8 @@ pub(crate) fn load_dataset<B: Backend>(
 
         // Extract COLMAP sfm points.
         let points_data = {
-            let mut points_file = archive.file_at_path(&points_path)?;
-            colmap_reader::read_points3d(&mut points_file, is_binary)
+            let mut points_file = archive.open_path(&points_path).await?;
+            colmap_reader::read_points3d(&mut points_file, is_binary).await
         };
 
         // Ignore empty points data.

crates/brush-dataset/src/formats/mod.rs

@@ -1,11 +1,11 @@
 use crate::{
+    brush_vfs::BrushVfs,
     splat_import::{load_splat_from_ply, SplatMessage},
-    zip::DatasetZip,
     Dataset, LoadDatasetArgs,
 };
 use anyhow::Result;
 use brush_render::Backend;
-use std::{io::Cursor, pin::Pin};
+use std::{path::Path, pin::Pin};
 use tokio_stream::Stream;
 
 pub mod colmap;
@@ -14,32 +14,36 @@ pub mod nerfstudio;
 // A dynamic stream of datasets
 type DataStream<T> = Pin<Box<dyn Stream<Item = Result<T>> + Send + 'static>>;
 
-pub fn load_dataset<B: Backend>(
-    mut archive: DatasetZip,
+pub async fn load_dataset<B: Backend>(
+    mut vfs: BrushVfs,
     load_args: &LoadDatasetArgs,
     device: &B::Device,
 ) -> anyhow::Result<(DataStream<SplatMessage<B>>, DataStream<Dataset>)> {
-    let streams = nerfstudio::read_dataset(archive.clone(), load_args, device)
-        .or_else(|_| colmap::load_dataset::<B>(archive.clone(), load_args, device));
+    let stream = nerfstudio::read_dataset(vfs.clone(), load_args, device).await;
 
-    let Ok(streams) = streams else {
-        anyhow::bail!("Couldn't parse dataset as any format. Only some formats are supported.")
+    let stream = match stream {
+        Ok(s) => Ok(s),
+        Err(_) => colmap::load_dataset::<B>(vfs.clone(), load_args, device).await,
     };
 
-    // If there's an init.ply definitey override the init stream with that.
-    let init_path = archive.find_with_extension(".ply", "init");
+    let stream = match stream {
+        Ok(stream) => stream,
+        Err(e) => anyhow::bail!(
+            "Couldn't parse dataset as any format. Only some formats are supported. {e}"
+        ),
+    };
 
-    let init_stream = if let Ok(path) = init_path {
-        let ply_data = archive.read_bytes_at_path(&path)?;
-        log::info!("Using {path:?} as initial point cloud.");
+    // If there's an init.ply definitey override the init stream with that.
+    let init_stream = if let Ok(reader) = vfs.open_path(Path::new("init.ply")).await {
+        log::info!("Using init.ply as initial point cloud.");
         Box::pin(load_splat_from_ply(
-            Cursor::new(ply_data),
+            reader,
             load_args.subsample_points,
             device.clone(),
         ))
     } else {
-        streams.0
+        stream.0
     };
 
-    Ok((init_stream, streams.1))
+    Ok((init_stream, stream.1))
 }