Repository: mosure/bevy_zeroverse (language: rust)

Issue:
add projection near/far
https://github.com/mosure/bevy_zeroverse/blob/ba3a142b7c653b6579edb7b856916a74894260cb/ffi/src/lib.rs#L60

```rust


    #[pyo3(get, set)]
    fovy: f32,

    // TODO: add projection near/far
}

#[pymethods]

```


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 4674dbd..ee7573a 100644
--- a/README.md
+++ b/README.md
@@ -36,7 +36,7 @@ from torch.utils.data import DataLoader
 
 dataset = BevyZeroverseDataset(
     editor=False, headless=True, num_cameras=6,
-    width=640, height=360, num_samples=1e6,
+    width=640, height=480, num_samples=1e6,
 )
 dataloader = DataLoader(
     dataset, batch_size=4, shuffle=True, num_workers=1,
diff --git a/ffi/Cargo.toml b/ffi/Cargo.toml
index e75dfa1..043cf65 100644
--- a/ffi/Cargo.toml
+++ b/ffi/Cargo.toml
@@ -1,7 +1,8 @@
 [package]
 name = "bevy_zeroverse_ffi"
-version = "0.1.0"
+version = "0.3.0"
 edition = "2021"
+default-run = "generate"
 
 
 [features]
@@ -14,12 +15,17 @@ extension-module = [
 
 
 [dependencies]
+bevy_args = "1.6"
 bevy_zeroverse = { path = "../", default-features = false, features = ["python"] }
+bytemuck = "1.17"
+clap = { version = "4.4", features = ["derive"] }
 image = "0.25"
 ndarray = { version = "0.15", features = ["blas"] }
 once_cell = "1.19"
 pyo3 = { version = "0.22", features = ["macros"] }
 pyo3-log = "0.11"
+safetensors = "0.4"
+serde = "1.0"
 
 
 [dependencies.bevy]
@@ -34,5 +40,10 @@ features = [
 
 
 [lib]
-name = "bevy_zeroverse"
+name = "bevy_zeroverse_ffi"
 path = "src/lib.rs"
+
+
+[[bin]]
+name = "generate"
+path = "src/generate.rs"
diff --git a/ffi/README.md b/ffi/README.md
index d3841a3..ab87687 100644
--- a/ffi/README.md
+++ b/ffi/README.md
@@ -23,7 +23,7 @@ from torch.utils.data import DataLoader
 
 dataset = BevyZeroverseDataset(
     editor=False, headless=True, num_cameras=6,
-    width=640, height=360, num_samples=1e6,
+    width=640, height=480, num_samples=1e6,
 )
 dataloader = DataLoader(dataset, batch_size=2, shuffle=True, num_workers=1)
 
@@ -33,7 +33,13 @@ for batch in dataloader:
 ```
 
 
-<!-- ### macos setup -->
+### macos setup
+
+macos does not support running the generator off main thread. right now, the only way to generate on mac is from rust. e.g.
+
+```bash
+cargo run -p bevy_zeroverse_ffi --bin generate -- --help
+```
 
 <!-- ```bash
 LIBTORCH_PATH=$(python3 -c "import site; print(site.getsitepackages()[0] + '/torch/lib')")
diff --git a/ffi/python/bevy_zeroverse_dataloader/__init__.py b/ffi/python/bevy_zeroverse_dataloader/__init__.py
index 540e31e..3be19b6 100644
--- a/ffi/python/bevy_zeroverse_dataloader/__init__.py
+++ b/ffi/python/bevy_zeroverse_dataloader/__init__.py
@@ -5,17 +5,19 @@
 import torch
 from torch.utils.data import DataLoader, Dataset
 
-import bevy_zeroverse
+import bevy_zeroverse_ffi
 
 
 # TODO: add sample-level world rotation augment
 class View:
-    def __init__(self, color, depth, normal, world_from_view, fovy, width, height):
+    def __init__(self, color, depth, normal, world_from_view, fovy, near, far, width, height):
         self.color = color
         self.depth = depth
         self.normal = normal
         self.world_from_view = world_from_view
         self.fovy = fovy
+        self.near = near
+        self.far = far
         self.width = width
         self.height = height
 
@@ -42,7 +44,9 @@ def reshape_data(data, dtype):
 
         world_from_view = np.array(rust_view.world_from_view)
         fovy = rust_view.fovy
-        return cls(color, depth, normal, world_from_view, fovy, width, height)
+        near = rust_view.near
+        far = rust_view.far
+        return cls(color, depth, normal, world_from_view, fovy, near, far, width, height)
 
     def to_tensors(self):
         color_tensor = torch.tensor(self.color, dtype=torch.float32)
@@ -50,27 +54,35 @@ def to_tensors(self):
         normal_tensor = torch.tensor(self.normal, dtype=torch.float32)
 
         color_tensor = color_tensor[..., :3]
-        depth_tensor = depth_tensor[..., 0]
+        depth_tensor = depth_tensor[..., 0:1]
         normal_tensor = normal_tensor[..., :3]
 
         world_from_view_tensor = torch.tensor(self.world_from_view, dtype=torch.float32)
-        fovy_tensor = torch.tensor(self.fovy, dtype=torch.float32)
+
+        fovy_tensor = torch.tensor(self.fovy, dtype=torch.float32).unsqueeze(-1)
+        near_tensor = torch.tensor(self.near, dtype=torch.float32).unsqueeze(-1)
+        far_tensor = torch.tensor(self.far, dtype=torch.float32).unsqueeze(-1)
+
         return {
             'color': color_tensor,
             'depth': depth_tensor,
             'normal': normal_tensor,
             'world_from_view': world_from_view_tensor,
-            'fovy': fovy_tensor
+            'fovy': fovy_tensor,
+            'near': near_tensor,
+            'far': far_tensor,
         }
 
 class Sample:
-    def __init__(self, views):
+    def __init__(self, views, aabb):
         self.views = views
+        self.aabb = aabb
 
     @classmethod
     def from_rust(cls, rust_sample, width, height):
         views = [View.from_rust(view, width, height) for view in rust_sample.views]
-        return cls(views)
+        aabb = np.array(rust_sample.aabb)
+        return cls(views, aabb)
 
     def to_tensors(self):
         tensor_dict = {
@@ -78,7 +90,9 @@ def to_tensors(self):
             'depth': [],
             'normal': [],
             'world_from_view': [],
-            'fovy': []
+            'fovy': [],
+            'near': [],
+            'far': [],
         }
 
         if len(self.views) == 0:
@@ -93,14 +107,16 @@ def to_tensors(self):
         for key in tensor_dict:
             tensor_dict[key] = torch.stack(tensor_dict[key], dim=0)
 
+        tensor_dict['aabb'] = torch.tensor(self.aabb, dtype=torch.float32)
+
         return tensor_dict
 
 
 # TODO: add dataset seed parameter to config
 class BevyZeroverseDataset(Dataset):
     scene_map = {
-        'object': bevy_zeroverse.ZeroverseSceneType.Object,
-        'room': bevy_zeroverse.ZeroverseSceneType.Room,
+        'object': bevy_zeroverse_ffi.ZeroverseSceneType.Object,
+        'room': bevy_zeroverse_ffi.ZeroverseSceneType.Room,
     }
 
     def __init__(
@@ -125,7 +141,7 @@ def __init__(
         self.scene_type = scene_type
 
     def initialize(self):
-        config = bevy_zeroverse.BevyZeroverseConfig()
+        config = bevy_zeroverse_ffi.BevyZeroverseConfig()
         config.editor = self.editor
         config.headless = self.headless
         config.num_cameras = self.num_cameras
@@ -133,7 +149,7 @@ def initialize(self):
         config.height = self.height
         config.scene_type = BevyZeroverseDataset.scene_map[self.scene_type]
         config.regenerate_scene_material_shuffle_period = 256
-        bevy_zeroverse.initialize(
+        bevy_zeroverse_ffi.initialize(
             config,
             self.root_asset_folder,
         )
@@ -146,7 +162,7 @@ def __getitem__(self, idx):
         if not self.initialized:
             self.initialize()
 
-        rust_sample = bevy_zeroverse.next()
+        rust_sample = bevy_zeroverse_ffi.next()
         sample = Sample.from_rust(rust_sample, self.width, self.height)
         return sample.to_tensors()
 
@@ -157,18 +173,24 @@ def chunk_and_save(
     bytes_per_chunk: int = int(256 * 1024 * 1024),
     n_workers: int = 1,
 ):
+    output_dir.mkdir(exist_ok=True, parents=True)
+    existing_chunks = sorted(output_dir.glob("*.safetensors"))
+    if existing_chunks:
+        latest_chunk = existing_chunks[-1]
+        chunk_index = int(latest_chunk.stem)
+        print(f"resuming from chunk {chunk_index}.")
+    else:
+        chunk_index = 0
+
     chunk_size = 0
-    chunk_index = 0
     chunk = []
-    original_samples = []
-    chunk_file_paths = []
+    chunk_file_paths = [output_dir / f"{int(chunk.stem):0>6}.safetensors" for chunk in existing_chunks]
 
     def save_chunk():
-        nonlocal chunk_size, chunk_index, chunk, original_samples, chunk_file_paths
+        nonlocal chunk_size, chunk_index, chunk, chunk_file_paths
 
         chunk_key = f"{chunk_index:0>6}"
         print(f"saving chunk {chunk_key} of {len(dataset)} ({chunk_size / 1e6:.2f} MB).")
-        output_dir.mkdir(exist_ok=True, parents=True)
         file_path = output_dir / f"{chunk_key}.safetensors"
 
         batch = {}
@@ -187,13 +209,17 @@ def save_chunk():
         chunk_index += 1
         chunk = []
 
+        del batch
+        torch.cuda.empty_cache()
+        import gc
+        gc.collect()
+
     dataloader = DataLoader(dataset, batch_size=1, num_workers=n_workers, shuffle=False)
 
     for idx, sample in enumerate(dataloader):
         sample = {k: v.squeeze(0) for k, v in sample.items()}
         sample_size = sum(tensor.numel() * tensor.element_size() for tensor in sample.values())
         chunk.append(sample)
-        original_samples.append(sample)
         chunk_size += sample_size
 
         print(f"    added sample {idx} to chunk ({sample_size / 1e6:.2f} MB).")
@@ -203,13 +229,13 @@ def save_chunk():
     if chunk_size > 0:
         save_chunk()
 
-    return original_samples, chunk_file_paths
+    return chunk_file_paths
 
 
 class ChunkedDataset(Dataset):
     def __init__(self, output_dir: Path):
-        self.output_dir = output_dir
-        self.chunk_files = sorted(output_dir.glob("*.safetensors"))
+        self.output_dir = Path(output_dir)
+        self.chunk_files = sorted(self.output_dir.glob("*.safetensors"))
 
     def load_chunk(self, file_path: Path):
         with safe_open(str(file_path), framework="pt", device="cpu") as f:
diff --git a/ffi/python/generate.py b/ffi/python/generate.py
index cea9248..3b14a24 100644
--- a/ffi/python/generate.py
+++ b/ffi/python/generate.py
@@ -12,13 +12,10 @@ def generate_chunked_dataset(
         headless=True,
         num_cameras=4,
         width=640,
-        height=360,
+        height=480,
         num_samples=100,
     )
 ) -> list:
-    if output_dir.exists():
-        shutil.rmtree(output_dir)
-
     return chunk_and_save(
         dataset,
         output_dir,
diff --git a/ffi/python/view.py b/ffi/python/view.py
new file mode 100644
index 0000000..d2c5e4e
--- /dev/null
+++ b/ffi/python/view.py
@@ -0,0 +1,17 @@
+from torch.utils.data import DataLoader
+
+from bevy_zeroverse_dataloader import ChunkedDataset
+
+from test import visualize
+
+
+def main():
+    chunked_dataset = ChunkedDataset("data/zeroverse/rust")
+    dataloader = DataLoader(chunked_dataset, batch_size=1, shuffle=False)
+
+    for batch in dataloader:
+        visualize(batch)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/ffi/setup.py b/ffi/setup.py
index 4aedb02..1518d36 100644
--- a/ffi/setup.py
+++ b/ffi/setup.py
@@ -3,12 +3,12 @@
 
 
 setup(
-    name="bevy_zeroverse",
-    version="0.1",
+    name="bevy_zeroverse_ffi",
+    version="0.3",
     packages=find_packages(where="python"),
     package_dir={"": "python"},
     rust_extensions=[
-        RustExtension("bevy_zeroverse")
+        RustExtension("bevy_zeroverse_ffi")
     ],
     include_package_data=True,
     zip_safe=False,
diff --git a/ffi/src/generate.rs b/ffi/src/generate.rs
new file mode 100644
index 0000000..7da279b
--- /dev/null
+++ b/ffi/src/generate.rs
@@ -0,0 +1,410 @@
+use std::{
+    borrow::Cow,
+    path::Path,
+    sync::mpsc::RecvTimeoutError,
+    time::Duration,
+};
+
+use bevy::prelude::*;
+use bevy_args::{
+    parse_args,
+    Deserialize,
+    Parser,
+    Serialize,
+};
+use bevy_zeroverse::{
+    app::BevyZeroverseConfig,
+    scene::ZeroverseSceneType,
+};
+use bevy_zeroverse_ffi::{
+    create_app,
+    setup_globals,
+    SAMPLE_RECEIVER,
+    APP_FRAME_SENDER,
+    Sample,
+};
+use bytemuck::cast_slice;
+use ndarray::{Array5, Array4, Array3, Array2, Array1, ArrayBase, OwnedRepr, Axis, Dimension, s};
+use safetensors::{
+    serialize_to_file,
+    View,
+    Dtype,
+};
+
+
+
+pub struct StackedViews {
+    pub color: Array5<f32>,             // Shape: (batch_size, num_views, height, width, channels)
+    pub depth: Array5<f32>,             // Shape: (batch_size, num_views, height, width, 1)
+    pub normal: Array5<f32>,            // Shape: (batch_size, num_views, height, width, channels)
+    pub world_from_view: Array4<f32>,   // Shape: (batch_size, num_views, 4, 4)
+    pub fovy: Array3<f32>,              // Shape: (batch_size, num_views, 1)
+    pub near: Array3<f32>,              // Shape: (batch_size, num_views, 1)
+    pub far: Array3<f32>,               // Shape: (batch_size, num_views, 1)
+    pub aabb: Array3<f32>,              // Shape: (batch_size, 2, 3) - min and max
+}
+
+struct Wrapper<A, D>(ArrayBase<OwnedRepr<A>, D>);
+
+impl<D: Dimension> Wrapper<f32, D> {
+    fn buffer(&self) -> &[u8] {
+        let slice = self.0.as_slice().expect("Non-contiguous tensors");
+        let new_slice: &[u8] = unsafe {
+            std::slice::from_raw_parts(slice.as_ptr() as *const u8, std::mem::size_of_val(slice))
+        };
+        new_slice
+    }
+}
+
+impl<D: Dimension> View for Wrapper<f32, D> {
+    fn dtype(&self) -> Dtype {
+        Dtype::F32
+    }
+    fn shape(&self) -> &[usize] {
+        self.0.shape()
+    }
+    fn data(&self) -> Cow<[u8]> {
+        self.buffer().into()
+    }
+    fn data_len(&self) -> usize {
+        self.buffer().len()
+    }
+}
+
+fn stack_samples(
+    samples: Vec<Sample>,
+    zeroverse_config: &BevyZeroverseConfig,
+) -> StackedViews {
+    let _batch_size = samples.len();
+    let _num_views = samples.first().map_or(0, |sample| sample.views.len());
+    let height = zeroverse_config.height as usize;
+    let width = zeroverse_config.width as usize;
+
+    let mut color_stacks = Vec::new();
+    let mut depth_stacks = Vec::new();
+    let mut normal_stacks = Vec::new();
+    let mut world_from_view_stacks = Vec::new();
+    let mut fovy_stacks = Vec::new();
+    let mut near_stacks = Vec::new();
+    let mut far_stacks = Vec::new();
+    let mut aabb_stacks = Vec::new();
+
+    for sample in samples {
+        let mut color_views = Vec::new();
+        let mut depth_views = Vec::new();
+        let mut normal_views = Vec::new();
+        let mut world_from_view_views = Vec::new();
+        let mut fovy_views = Vec::new();
+        let mut near_views = Vec::new();
+        let mut far_views = Vec::new();
+
+        for view in sample.views {
+            let color_f32: &[f32] = cast_slice(&view.color);
+            let depth_f32: &[f32] = cast_slice(&view.depth);
+            let normal_f32: &[f32] = cast_slice(&view.normal);
+
+            let color = Array3::from_shape_vec((height, width, 4), color_f32.to_vec()).unwrap();
+            let depth = Array3::from_shape_vec((height, width, 4), depth_f32.to_vec()).unwrap();
+            let normal = Array3::from_shape_vec((height, width, 4), normal_f32.to_vec()).unwrap();
+
+            let world_from_view = Array2::from_shape_vec((4, 4), view.world_from_view.concat()).unwrap();
+            let fovy = Array1::from_elem(1, view.fovy);
+            let near = Array1::from_elem(1, view.near);
+            let far = Array1::from_elem(1, view.far);
+
+            let color = color.slice(s![.., .., 0..3]).to_owned();       // rgb
+            let depth = depth.slice(s![.., .., 0..1]).to_owned();       // x
+            let normal = normal.slice(s![.., .., 0..3]).to_owned();     // xyz
+
+            color_views.push(color);
+            depth_views.push(depth);
+            normal_views.push(normal);
+            world_from_view_views.push(world_from_view);
+            fovy_views.push(fovy);
+            near_views.push(near);
+            far_views.push(far);
+        }
+
+        let color_views_stacked = ndarray::stack(Axis(0), &color_views.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+        let depth_views_stacked = ndarray::stack(Axis(0), &depth_views.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+        let normal_views_stacked = ndarray::stack(Axis(0), &normal_views.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+        let world_from_view_views_stacked = ndarray::stack(Axis(0), &world_from_view_views.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+        let fovy_views_stacked = ndarray::stack(Axis(0), &fovy_views.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+        let near_views_stacked = ndarray::stack(Axis(0), &near_views.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+        let far_views_stacked = ndarray::stack(Axis(0), &far_views.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+
+        color_stacks.push(color_views_stacked);
+        depth_stacks.push(depth_views_stacked);
+        normal_stacks.push(normal_views_stacked);
+        world_from_view_stacks.push(world_from_view_views_stacked);
+        fovy_stacks.push(fovy_views_stacked);
+        near_stacks.push(near_views_stacked);
+        far_stacks.push(far_views_stacked);
+
+        let sample_aabb = Array2::from_shape_vec((2, 3), sample.aabb.concat()).unwrap();
+        aabb_stacks.push(sample_aabb);
+    }
+
+    let color_stacked = ndarray::stack(Axis(0), &color_stacks.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+    let depth_stacked = ndarray::stack(Axis(0), &depth_stacks.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+    let normal_stacked = ndarray::stack(Axis(0), &normal_stacks.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+    let world_from_view_stacked = ndarray::stack(Axis(0), &world_from_view_stacks.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+    let fovy_stacked = ndarray::stack(Axis(0), &fovy_stacks.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+    let near_stacked = ndarray::stack(Axis(0), &near_stacks.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+    let far_stacked = ndarray::stack(Axis(0), &far_stacks.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+    let aabb_stacked = ndarray::stack(Axis(0), &aabb_stacks.iter().map(|a| a.view()).collect::<Vec<_>>()).unwrap();
+
+    StackedViews {
+        color: color_stacked,                       // Shape: (batch_size, num_views, height, width, 3)
+        depth: depth_stacked,                       // Shape: (batch_size, num_views, height, width, 1)
+        normal: normal_stacked,                     // Shape: (batch_size, num_views, height, width, 3)
+        world_from_view: world_from_view_stacked,   // Shape: (batch_size, num_views, 4, 4)
+        fovy: fovy_stacked,                         // Shape: (batch_size, num_views)
+        near: near_stacked,                         // Shape: (batch_size, num_views)
+        far: far_stacked,                           // Shape: (batch_size, num_views)
+        aabb: aabb_stacked,                         // Shape: (batch_size, 2, 3)
+    }
+}
+
+
+enum TensorView {
+    Color(Wrapper<f32, ndarray::Ix5>),
+    Depth(Wrapper<f32, ndarray::Ix5>),
+    Normal(Wrapper<f32, ndarray::Ix5>),
+    WorldFromView(Wrapper<f32, ndarray::Ix4>),
+    Aabb(Wrapper<f32, ndarray::Ix3>),
+    Singular(Wrapper<f32, ndarray::Ix3>),
+}
+
+impl View for TensorView {
+    fn dtype(&self) -> Dtype {
+        match self {
+            TensorView::Color(t) => t.dtype(),
+            TensorView::Depth(t) => t.dtype(),
+            TensorView::Normal(t) => t.dtype(),
+            TensorView::WorldFromView(t) => t.dtype(),
+            TensorView::Aabb(t) => t.dtype(),
+            TensorView::Singular(t) => t.dtype(),
+        }
+    }
+
+    fn shape(&self) -> &[usize] {
+        match self {
+            TensorView::Color(t) => t.shape(),
+            TensorView::Depth(t) => t.shape(),
+            TensorView::Normal(t) => t.shape(),
+            TensorView::WorldFromView(t) => t.shape(),
+            TensorView::Aabb(t) => t.shape(),
+            TensorView::Singular(t) => t.shape(),
+        }
+    }
+
+    fn data(&self) -> Cow<[u8]> {
+        match self {
+            TensorView::Color(t) => t.data(),
+            TensorView::Depth(t) => t.data(),
+            TensorView::Normal(t) => t.data(),
+            TensorView::WorldFromView(t) => t.data(),
+            TensorView::Aabb(t) => t.data(),
+            TensorView::Singular(t) => t.data(),
+        }
+    }
+
+    fn data_len(&self) -> usize {
+        match self {
+            TensorView::Color(t) => t.data_len(),
+            TensorView::Depth(t) => t.data_len(),
+            TensorView::Normal(t) => t.data_len(),
+            TensorView::WorldFromView(t) => t.data_len(),
+            TensorView::Aabb(t) => t.data_len(),
+            TensorView::Singular(t) => t.data_len(),
+        }
+    }
+}
+
+
+fn save_stacked_views_to_safetensors(stacked_views: StackedViews, output_path: &Path) -> Result<(), safetensors::SafeTensorError> {
+    let data: Vec<(&str, TensorView)> = vec![
+        ("color", TensorView::Color(Wrapper(stacked_views.color))),
+        ("depth", TensorView::Depth(Wrapper(stacked_views.depth))),
+        ("normal", TensorView::Normal(Wrapper(stacked_views.normal))),
+        ("world_from_view", TensorView::WorldFromView(Wrapper(stacked_views.world_from_view))),
+        ("fovy", TensorView::Singular(Wrapper(stacked_views.fovy))),
+        ("near", TensorView::Singular(Wrapper(stacked_views.near))),
+        ("far", TensorView::Singular(Wrapper(stacked_views.far))),
+        ("aabb", TensorView::Aabb(Wrapper(stacked_views.aabb))),
+    ];
+
+    serialize_to_file(data, &None, output_path)
+}
+
+
+
+#[derive(
+    Clone,
+    Debug,
+    Resource,
+    Serialize,
+    Deserialize,
+    Parser,
+    Reflect,
+)]
+#[command(about = "bevy_zeroverse viewer", version, long_about = None)]
+#[reflect(Resource)]
+pub struct GeneratorConfig {
+    #[arg(long, default_value = "10")]
+    pub num_samples: usize,
+
+    #[arg(long, default_value = "268435456")]  // 256 MB
+    pub bytes_per_chunk: usize,
+
+    #[arg(long, default_value = "data/zeroverse/rust")]
+    pub output_dir: String,
+}
+
+impl Default for GeneratorConfig {
+    fn default() -> GeneratorConfig {
+        GeneratorConfig {
+            num_samples: 10,
+            bytes_per_chunk: 268435456,
+            output_dir: "data/zeroverse/rust".to_string(),
+        }
+    }
+}
+
+
+fn receive_samples(
+    generator_config: &GeneratorConfig,
+    zeroverse_config: &BevyZeroverseConfig,
+) {
+    let receiver = SAMPLE_RECEIVER.get().unwrap();
+    let receiver = receiver.lock().unwrap();
+
+    let mut chunk_size = 0;
+    let mut chunk_index = 0;
+    let mut chunk_samples = Vec::new();
+
+    for sample_index in 0..generator_config.num_samples {
+        {
+            let app_frame_sender = APP_FRAME_SENDER.get().unwrap();
+            app_frame_sender.send(()).unwrap();
+        }
+
+        let timeout = Duration::from_secs(30);
+        match receiver.recv_timeout(timeout) {
+            Ok(sample) => {
+                chunk_samples.push(sample);
+
+                let sample_size = estimate_sample_size(chunk_samples.last().unwrap());
+                chunk_size += sample_size;
+
+                info!(
+                    "    added sample {} to chunk ({:.2} MB).",
+                    sample_index,
+                    sample_size as f64 / 1e6
+                );
+
+                if chunk_size >= generator_config.bytes_per_chunk {
+                    save_chunk(&chunk_samples, chunk_index, generator_config, zeroverse_config);
+
+                    chunk_samples.clear();
+                    chunk_size = 0;
+                    chunk_index += 1;
+                }
+            },
+            Err(RecvTimeoutError::Timeout) => {
+                error!("receive operation timed out");
+                std::process::exit(1);
+            }
+            Err(RecvTimeoutError::Disconnected) => {
+                error!("channel disconnected");
+                std::process::exit(1);
+            }
+        }
+    }
+
+    if !chunk_samples.is_empty() {
+        save_chunk(&chunk_samples, chunk_index, generator_config, zeroverse_config);
+    }
+
+    info!("finished generating samples");
+    std::process::exit(0);
+}
+
+
+fn estimate_sample_size(sample: &Sample) -> usize {
+    let mut size = 0;
+
+    for view in &sample.views {
+        size += view.color.len() * 3 / 4;
+        size += view.depth.len() / 4;
+        size += view.normal.len() * 3 / 4;
+        size += view.world_from_view.len();
+        size += 1;  // fovy
+        size += 1;  // near
+        size += 1;  // far
+    }
+
+    size += sample.aabb.len();
+
+    size
+}
+
+fn save_chunk(
+    chunk_samples: &[Sample],
+    chunk_index: usize,
+    generator_config: &GeneratorConfig,
+    zeroverse_config: &BevyZeroverseConfig,
+) {
+    let stacked_views = stack_samples(chunk_samples.to_vec(), zeroverse_config);
+
+    let file_name = format!("{:06}.safetensors", chunk_index);
+    let output_dir = Path::new(generator_config.output_dir.as_str());
+
+    if !output_dir.exists() {
+        std::fs::create_dir_all(output_dir).unwrap();
+    }
+
+    let output_path = output_dir.join(file_name);
+
+    let chunk_size = chunk_samples.iter().map(estimate_sample_size).sum::<usize>();
+    info!(
+        "saving chunk {} of {} ({:.2} MB).",
+        chunk_index,
+        generator_config.num_samples,
+        chunk_size as f64 / 1e6
+    );
+
+    match save_stacked_views_to_safetensors(stacked_views, &output_path) {
+        Ok(_) => info!("successfully saved chunk {}", chunk_index),
+        Err(e) => warn!("failed to save chunk {}: {:?}", chunk_index, e),
+    }
+}
+
+
+fn main() {
+    let generator_args = parse_args::<GeneratorConfig>();
+    let mut zeroverse_args = parse_args::<BevyZeroverseConfig>();
+
+    zeroverse_args.editor = false;
+    zeroverse_args.headless = true;
+    zeroverse_args.num_cameras = 4;
+    zeroverse_args.width = 640.0;
+    zeroverse_args.height = 480.0;
+    zeroverse_args.regenerate_scene_material_shuffle_period = 256;
+    zeroverse_args.scene_type = ZeroverseSceneType::Object;
+
+    setup_globals(None);
+
+    let mut app = create_app(zeroverse_args.clone().into());
+
+    std::thread::spawn(move || {
+        receive_samples(&generator_args, &zeroverse_args);
+    });
+
+    // TODO: set to headless mode
+    app.run();
+}
+
+
+// TODO: add test comparing safetensors from python vs rust
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs
index 2c5bc1d..b286c76 100644
--- a/ffi/src/lib.rs
+++ b/ffi/src/lib.rs
@@ -37,6 +37,8 @@ use ::bevy_zeroverse::{
     render::RenderMode,
     scene::{
         RegenerateSceneEvent,
+        SceneAabb,
+        ZeroverseSceneRoot,
         ZeroverseSceneType,
     },
 };
@@ -46,18 +48,22 @@ use ::bevy_zeroverse::{
 
 #[derive(Clone, Debug, Default)]
 #[pyclass]
-struct View {
-    color: Vec<u8>,
-    depth: Vec<u8>,
-    normal: Vec<u8>,
+pub struct View {
+    pub color: Vec<u8>,
+    pub depth: Vec<u8>,
+    pub normal: Vec<u8>,
 
     #[pyo3(get, set)]
-    world_from_view: [[f32; 4]; 4],
+    pub world_from_view: [[f32; 4]; 4],
 
     #[pyo3(get, set)]
-    fovy: f32,
+    pub fovy: f32,
 
-    // TODO: add projection near/far
+    #[pyo3(get, set)]
+    pub near: f32,
+
+    #[pyo3(get, set)]
+    pub far: f32,
 }
 
 #[pymethods]
@@ -86,10 +92,14 @@ impl View {
     }
 }
 
-#[derive(Debug, Default, Resource)]
+#[derive(Clone, Debug, Default, Resource)]
 #[pyclass]
-struct Sample {
-    views: Vec<View>,
+pub struct Sample {
+    pub views: Vec<View>,
+
+    /// min and max corners of the axis-aligned bounding box
+    #[pyo3(get, set)]
+    pub aabb: [[f32; 3]; 2],
 }
 
 #[pymethods]
@@ -160,11 +170,11 @@ impl SamplerState {
 }
 
 
-static APP_FRAME_RECEIVER: OnceCell<Arc<Mutex<Receiver<()>>>> = OnceCell::new();
-static APP_FRAME_SENDER: OnceCell<Sender<()>> = OnceCell::new();
+pub static APP_FRAME_RECEIVER: OnceCell<Arc<Mutex<Receiver<()>>>> = OnceCell::new();
+pub static APP_FRAME_SENDER: OnceCell<Sender<()>> = OnceCell::new();
 
-static SAMPLE_RECEIVER: OnceCell<Arc<Mutex<Receiver<Sample>>>> = OnceCell::new();
-static SAMPLE_SENDER: OnceCell<Sender<Sample>> = OnceCell::new();
+pub static SAMPLE_RECEIVER: OnceCell<Arc<Mutex<Receiver<Sample>>>> = OnceCell::new();
+pub static SAMPLE_SENDER: OnceCell<Sender<Sample>> = OnceCell::new();
 
 
 fn signaled_runner(mut app: App) -> AppExit {
@@ -202,6 +212,26 @@ fn signaled_runner(mut app: App) -> AppExit {
 }
 
 
+pub fn create_app(
+    override_args: Option<BevyZeroverseConfig>,
+) -> App {
+    let mut app = viewer_app(override_args);
+
+    app.init_resource::<Sample>();
+
+    // initialize to disabled state
+    app.insert_resource(SamplerState {
+        enabled: false,
+        ..Default::default()
+    });
+
+    app.add_systems(PreUpdate, sample_stream);
+    app.set_runner(signaled_runner);
+
+    app
+}
+
+
 pub fn setup_and_run_app(
     new_thread: bool,
     override_args: Option<BevyZeroverseConfig>,
@@ -212,21 +242,8 @@ pub fn setup_and_run_app(
         let ready = Arc::clone(&ready);
 
         move || {
-            let mut app = viewer_app(override_args);
-
-            app.init_resource::<Sample>();
-
-            // initialize to disabled state
-            app.insert_resource(SamplerState {
-                enabled: false,
-                ..Default::default()
-            });
-
-            app.add_systems(PreUpdate, sample_stream);
-
+            let mut app = create_app(override_args);
             ready.store(true, Ordering::Release);
-
-            app.set_runner(signaled_runner);
             app.run();
         }
     };
@@ -251,6 +268,7 @@ fn sample_stream(
         &Projection,
         &ImageCopier,
     )>,
+    scene: Query<(&ZeroverseSceneRoot, &SceneAabb)>,
     images: Res<Assets<Image>>,
     render_mode: ResMut<RenderMode>,
     mut regenerate_event: EventWriter<RegenerateSceneEvent>,
@@ -281,11 +299,21 @@ fn sample_stream(
         }
     }
 
+    let scene_aabb = scene.single().1;
+    buffered_sample.aabb = [
+        scene_aabb.min.into(),
+        scene_aabb.max.into(),
+    ];
+
     for (i, (camera_transform, projection, image_copier)) in cameras.iter().enumerate() {
         let view = &mut buffered_sample.views[i];
 
-        view.fovy = match projection {
-            Projection::Perspective(perspective) => perspective.fov,
+        match projection {
+            Projection::Perspective(perspective) => {
+                view.fovy = perspective.fov;
+                view.near = perspective.near;
+                view.far = perspective.far;
+            }
             Projection::Orthographic(_) => panic!("orthographic projection not supported"),
         };
 
@@ -311,6 +339,7 @@ fn sample_stream(
         let views = std::mem::take(&mut buffered_sample.views);
         let sample = Sample {
             views,
+            aabb: buffered_sample.aabb,
         };
 
         let sender = SAMPLE_SENDER.get().unwrap();
@@ -321,12 +350,7 @@ fn sample_stream(
 }
 
 
-// TODO: add process idx, disable logging on worker idx > 0
-#[pyfunction]
-#[pyo3(signature = (override_args=None, asset_root=None))]
-fn initialize(
-    py: Python<'_>,
-    override_args: Option<BevyZeroverseConfig>,
+pub fn setup_globals(
     asset_root: Option<String>,
 ) {
     if APP_FRAME_RECEIVER.get().is_some() {
@@ -354,6 +378,18 @@ fn initialize(
 
     SAMPLE_RECEIVER.set(Arc::new(Mutex::new(sample_receiver))).unwrap();
     SAMPLE_SENDER.set(sample_sender).unwrap();
+}
+
+
+// TODO: add process idx, disable logging on worker idx > 0
+#[pyfunction]
+#[pyo3(signature = (override_args=None, asset_root=None))]
+fn initialize(
+    py: Python<'_>,
+    override_args: Option<BevyZeroverseConfig>,
+    asset_root: Option<String>,
+) {
+    setup_globals(asset_root);
 
     py.allow_threads(|| {
         setup_and_run_app(true, override_args);
@@ -391,7 +427,7 @@ fn next(
 
 
 #[pymodule]
-fn bevy_zeroverse(m: &Bound<'_, PyModule>) -> PyResult<()> {
+fn bevy_zeroverse_ffi(m: &Bound<'_, PyModule>) -> PyResult<()> {
     pyo3_log::init();
 
     m.add_class::<BevyZeroverseConfig>()?;
diff --git a/src/scene/mod.rs b/src/scene/mod.rs
index aafecea..dd3375f 100644
--- a/src/scene/mod.rs
+++ b/src/scene/mod.rs
@@ -1,4 +1,7 @@
-use bevy::prelude::*;
+use bevy::{
+    prelude::*,
+    render::primitives::Aabb,
+};
 use bevy_args::{
     Deserialize,
     Serialize,
@@ -8,6 +11,11 @@ use bevy_args::{
 #[cfg(feature = "python")]
 use pyo3::prelude::*;
 
+use crate::{
+    app::BevyZeroverseConfig,
+    camera::EditorCameraGizmoConfigGroup,
+};
+
 pub mod lighting;
 pub mod object;
 pub mod room;
@@ -22,12 +30,16 @@ impl Plugin for ZeroverseScenePlugin {
 
         app.init_resource::<ZeroverseSceneSettings>();
         app.register_type::<ZeroverseSceneSettings>();
+        app.register_type::<SceneAabb>();
 
         app.add_plugins((
             lighting::ZeroverseLightingPlugin,
             object::ZeroverseObjectPlugin,
             room::ZeroverseRoomPlugin,
         ));
+
+        app.add_systems(PreUpdate, create_scene_aabb);
+        app.add_systems(Update, draw_scene_aabb);
     }
 }
 
@@ -70,3 +82,117 @@ pub struct RegenerateSceneEvent;
 
 #[derive(Event)]
 pub struct SceneLoadedEvent;
+
+
+
+#[derive(Component, Debug, Reflect)]
+pub struct SceneAabb {
+    pub min: Vec3,
+    pub max: Vec3,
+}
+
+impl SceneAabb {
+    fn new(center: Vec3) -> Self {
+        Self {
+            min: center,
+            max: center,
+        }
+    }
+
+    fn merge_aabb(&mut self, aabb: &Aabb, global_transform: &GlobalTransform) {
+        let min = aabb.min();
+        let max = aabb.max();
+        let corners = [
+            global_transform.transform_point(Vec3::new(max.x, max.y, max.z)),
+            global_transform.transform_point(Vec3::new(min.x, max.y, max.z)),
+            global_transform.transform_point(Vec3::new(min.x, max.y, min.z)),
+            global_transform.transform_point(Vec3::new(max.x, max.y, min.z)),
+            global_transform.transform_point(Vec3::new(max.x, min.y, max.z)),
+            global_transform.transform_point(Vec3::new(min.x, min.y, max.z)),
+            global_transform.transform_point(Vec3::new(min.x, min.y, min.z)),
+            global_transform.transform_point(Vec3::new(max.x, min.y, min.z)),
+        ];
+
+        for corner in corners {
+            let gt = corner.cmpgt(self.max);
+            let lt = corner.cmplt(self.min);
+
+            if gt.x {
+                self.max.x = corner.x;
+            } else if lt.x {
+                self.min.x = corner.x;
+            }
+
+            if gt.y {
+                self.max.y = corner.y;
+            } else if lt.y {
+                self.min.y = corner.y;
+            }
+
+            if gt.z {
+                self.max.z = corner.z;
+            } else if lt.z {
+                self.min.z = corner.z;
+            }
+        }
+    }
+}
+
+impl From<&SceneAabb> for Transform {
+    fn from(scene_aabb: &SceneAabb) -> Transform {
+        let min = scene_aabb.min;
+        let max = scene_aabb.max;
+
+        let center = (min + max) / 2.0;
+        let size = max - min;
+
+        Transform::from_translation(center)
+            * Transform::from_scale(size)
+    }
+}
+
+
+fn create_scene_aabb(
+    mut commands: Commands,
+    scene_instances: Query<
+        (
+            Entity,
+            &ZeroverseSceneRoot,
+            &GlobalTransform,
+        ),
+        Without<SceneAabb>,
+    >,
+    children: Query<&Children>,
+    bounding_boxes: Query<(&Aabb, &GlobalTransform)>,
+) {
+    for (entity, _instance, global_transform) in scene_instances.iter() {
+        let mut scene_aabb = SceneAabb::new(global_transform.translation());
+
+        if children.iter_descendants(entity).count() == 0 {
+            continue;
+        }
+
+        for child in children.iter_descendants(entity) {
+            let Ok((bb, transform)) = bounding_boxes.get(child) else { continue };
+            scene_aabb.merge_aabb(bb, transform);
+        }
+
+        commands.entity(entity).insert(scene_aabb);
+    }
+}
+
+
+fn draw_scene_aabb(
+    args: Res<BevyZeroverseConfig>,
+    scene_instances: Query<&SceneAabb>,
+    mut gizmos: Gizmos<EditorCameraGizmoConfigGroup>,
+) {
+    if !args.editor {
+        return;
+    }
+
+    for aabb in &scene_instances {
+        let color = Color::srgb(0.0, 1.0, 1.0);
+        gizmos.cuboid(Transform::from(aabb), color);
+    }
+}


Test patch:
diff --git a/ffi/python/test.py b/ffi/python/test.py
index b830bde..a6e3a63 100644
--- a/ffi/python/test.py
+++ b/ffi/python/test.py
@@ -12,10 +12,17 @@
 def visualize(batch):
     print(batch['color'].shape)
 
+    is_chunked = len(batch['color'].shape) == 6
+
     color_images = batch['color'].numpy()
     depth_images = batch['depth'].numpy()
     normal_images = batch['normal'].numpy()
 
+    if is_chunked:
+        color_images = color_images.squeeze(0)
+        depth_images = depth_images.squeeze(0)
+        normal_images = normal_images.squeeze(0)
+
     batch_size = color_images.shape[0]
     num_cameras = color_images.shape[1]
     num_image_types = 3  # color, depth, normal
@@ -82,7 +89,7 @@ def benchmark(dataloader):
 def test():
     dataset = BevyZeroverseDataset(
         editor=False, headless=True, num_cameras=6,
-        width=640, height=360, num_samples=1e6,
+        width=640, height=480, num_samples=1e6,
     )
     dataloader = DataLoader(dataset, batch_size=5, shuffle=True, num_workers=2)
 
@@ -99,7 +106,7 @@ def setUp(self):
         self.headless = True
         self.num_cameras = 4
         self.width = 640
-        self.height = 360
+        self.height = 480
         self.num_samples = 10
         self.bytes_per_chunk = int(256 * 1024 * 1024)
         self.stage = "test"
@@ -109,31 +116,7 @@ def setUp(self):
             shutil.rmtree(self.output_dir)
 
         self.dataset = BevyZeroverseDataset(self.editor, self.headless, self.num_cameras, self.width, self.height, self.num_samples)
-        self.original_samples = chunk_and_save(self.dataset, self.output_dir, self.bytes_per_chunk)
-
-    def test_chunked_dataset_loading(self):
-        chunked_dataset = ChunkedDataset(self.output_dir)
-        dataloader = DataLoader(chunked_dataset, batch_size=1, shuffle=False)
-
-        num_chunks = 0
-        total_loaded_samples = 0
-
-        expected_shapes = {key: tensor.shape for key, tensor in self.original_samples[0][0].items()}
-
-        for batch in dataloader:
-            num_chunks += 1
-
-            for key, tensor in batch.items():
-                tensor = tensor.squeeze(0)
-                expected_shape = (tensor.shape[0],) + expected_shapes[key]
-                self.assertEqual(tensor.shape, expected_shape, f"Mismatch in tensor shape for key {key}")
-
-            total_loaded_samples += batch['color'].squeeze(0).shape[0]
-
-        expected_num_chunks = len(chunked_dataset)
-        self.assertEqual(num_chunks, expected_num_chunks, "Mismatch in number of chunks")
-
-        self.assertEqual(total_loaded_samples, len(self.original_samples[0]))
+        self.chunk_paths = chunk_and_save(self.dataset, self.output_dir, self.bytes_per_chunk)
 
     def test_benchmark_chunked_dataloader(self):
         chunked_dataset = ChunkedDataset(self.output_dir)