From 6f5a488937469b850507799d52ace4c4986b0a7c Mon Sep 17 00:00:00 2001
From: mosure <mitchell@mosure.me>
Date: Tue, 30 Jul 2024 18:27:09 -0500
Subject: [PATCH] feat: scaffold for ffi io

---
 .github/workflows/clippy.yml            |   2 +-
 Cargo.toml                              |  18 +-
 ffi/bevy_zeroverse/Cargo.toml           |   2 +
 ffi/bevy_zeroverse/python/dataloader.py |  21 +-
 ffi/bevy_zeroverse/src/lib.rs           | 127 ++++++++-
 src/app.rs                              |  74 +++--
 src/io.rs                               | 342 ++++++++++++++++++++++++
 src/lib.rs                              |   1 +
 src/render/mod.rs                       |   2 +
 src/scene/mod.rs                        |   2 +
 10 files changed, 545 insertions(+), 46 deletions(-)
 create mode 100644 src/io.rs

diff --git a/.github/workflows/clippy.yml b/.github/workflows/clippy.yml
index 0c07027..7c18160 100644
--- a/.github/workflows/clippy.yml
+++ b/.github/workflows/clippy.yml
@@ -34,4 +34,4 @@ jobs:
         enable-sccache: "true"
 
     - name: lint
-      run: cargo clippy -- -Dwarnings
+      run: cargo clippy --all -- -Dwarnings
diff --git a/Cargo.toml b/Cargo.toml
index 111bac0..7f1e8f8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -31,6 +31,13 @@ exclude = [
 default-run = "viewer"
 
 
+[workspace]
+members = [
+  ".",
+  "ffi/bevy_zeroverse",
+]
+
+
 [features]
 default = [
   "asset_pipeline",
@@ -71,14 +78,18 @@ bevy_args = "1.6"
 bevy-inspector-egui = { version = "0.25", optional = true }
 bevy_panorbit_camera = { version = "0.19", optional = true, features = ["bevy_egui"] }
 clap = { version = "4.4", features = ["derive"] }
+futures-intrusive = "0.5"
 glob = "0.3"
 itertools = "0.13"
 noise = { version = "0.9" }
+pollster = "0.3"
+pyo3 = { version = "0.22", features = ["extension-module"] }
 rand = "0.8"
 rayon = { version = "1.10", optional = true }
 serde = "1.0"
 strum = "0.26"
 strum_macros = "0.26"
+wgpu = "0.20.0"
 
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
@@ -145,10 +156,3 @@ path = "src/lib.rs"
 name = "viewer"
 path = "src/viewer.rs"
 required-features = ["viewer"]
-
-
-[workspace]
-members = [
-  "ffi/bevy_zeroverse",
-  ".",
-]
diff --git a/ffi/bevy_zeroverse/Cargo.toml b/ffi/bevy_zeroverse/Cargo.toml
index 0c6dda5..afb7422 100644
--- a/ffi/bevy_zeroverse/Cargo.toml
+++ b/ffi/bevy_zeroverse/Cargo.toml
@@ -6,6 +6,8 @@ edition = "2021"
 
 [dependencies]
 bevy_zeroverse = { path = "../../" }
+ndarray = { version = "0.15", features = ["blas"] }
+once_cell = "1.19"
 pyo3 = { version = "0.22", features = ["extension-module"] }
 pyo3-log = "0.11"
 
diff --git a/ffi/bevy_zeroverse/python/dataloader.py b/ffi/bevy_zeroverse/python/dataloader.py
index dd3b0ef..9389d80 100644
--- a/ffi/bevy_zeroverse/python/dataloader.py
+++ b/ffi/bevy_zeroverse/python/dataloader.py
@@ -1,20 +1,13 @@
 import bevy_zeroverse
 
 
-print('initializing bevy zeroverse...')
-bevy_zeroverse.initialize()
-print('bevy zeroverse initialized!')
+config = bevy_zeroverse.BevyZeroverseConfig()
 
+config.headless = True
+config.num_cameras = 4
 
-# dataloader = bevy_zeroverse.ZeroverseDataloader(
-#     width=256,
-#     height=144,
-#     num_cameras=4,
-#     render_modes=['color', 'depth', 'normal'],
-#     seed=0,
-#     scene_type='room',
-# )
+bevy_zeroverse.initialize(config)
 
-# for batch in dataloader:
-#     print(batch)
-#     break
+
+sample = bevy_zeroverse.next()
+print(sample)
diff --git a/ffi/bevy_zeroverse/src/lib.rs b/ffi/bevy_zeroverse/src/lib.rs
index 8d79c2f..9f704c5 100644
--- a/ffi/bevy_zeroverse/src/lib.rs
+++ b/ffi/bevy_zeroverse/src/lib.rs
@@ -1,15 +1,93 @@
 use std::{
     sync::{
-        atomic::{AtomicBool, Ordering},
         Arc,
+        Mutex,
+        atomic::{
+            AtomicBool,
+            Ordering,
+        },
+        mpsc::{
+            Sender,
+            Receiver,
+            RecvTimeoutError,
+        },
     },
     thread,
+    time::Duration,
 };
 
-// use bevy::prelude::*;
-use pyo3::prelude::*;
+use bevy::prelude::*;
+use ndarray::{Array2, Array3};
+use once_cell::sync::OnceCell;
+use pyo3::{
+    prelude::*,
+    exceptions::PyTimeoutError,
+    types::PyList,
+};
+
+use ::bevy_zeroverse::app::{
+    viewer_app,
+    BevyZeroverseConfig,
+};
+
+
+type ColorImage = Array3<u8>;
+type DepthImage = Array2<f32>;
+type NormalImage = Array2<[f32; 3]>;
+
+
+#[derive(Clone)]
+#[pyclass]
+struct View {
+    color: ColorImage,
+    depth: DepthImage,
+    normal: NormalImage,
+
+    #[pyo3(get, set)]
+    view_from_world: [[f32; 4]; 4],
+
+    #[pyo3(get, set)]
+    fovy: f32,
+}
+
+#[pymethods]
+impl View {
+    // TODO: upgrade rust-numpy to latest once pyo3 0.22 support is available
+    #[getter]
+    fn color<'py>(&self, py: Python<'py>) -> Bound<'py, PyList> {
+        let color_list: Vec<_> = self.color.iter().map(|&v| v.into_py(py)).collect();
+        PyList::new_bound(py, color_list)
+    }
+
+    #[getter]
+    fn depth<'py>(&self, py: Python<'py>) -> Bound<'py, PyList> {
+        let depth_list: Vec<_> = self.depth.iter().map(|&v| v.into_py(py)).collect();
+        PyList::new_bound(py, depth_list)
+    }
+
+    #[getter]
+    fn normal<'py>(&self, py: Python<'py>) -> Bound<'py, PyList> {
+        let normal_list: Vec<_> = self.normal.iter().map(|&v| v.into_py(py)).collect();
+        PyList::new_bound(py, normal_list)
+    }
+}
 
-use ::bevy_zeroverse::app::viewer_app;
+#[pyclass]
+struct Sample {
+    views: Vec<View>,
+}
+
+#[pymethods]
+impl Sample {
+    #[getter]
+    fn views<'py>(&self, py: Python<'py>) -> Bound<'py, PyList> {
+        let views_list: Vec<_> = self.views.iter().map(|v| v.clone().into_py(py)).collect();
+        PyList::new_bound(py, views_list)
+    }
+}
+
+static SAMPLE_RECEIVER: OnceCell<Arc<Mutex<Receiver<Sample>>>> = OnceCell::new();
+// static SAMPLE_SENDER: OnceCell<Sender<Sample>> = OnceCell::new();
 
 
 // TODO: create Dataloader torch class (or a render `n` frames and return capture fn, used within a python wrapper dataloader, wrapper requires setup.py to include the python module)
@@ -17,6 +95,7 @@ use ::bevy_zeroverse::app::viewer_app;
 
 pub fn setup_and_run_app(
     new_thread: bool,
+    override_args: Option<BevyZeroverseConfig>,
 ) {
     let ready = Arc::new(AtomicBool::new(false));
 
@@ -24,7 +103,7 @@ pub fn setup_and_run_app(
         let ready = Arc::clone(&ready);
 
         move || {
-            let mut app = viewer_app(None);
+            let mut app = viewer_app(override_args);
 
             ready.store(true, Ordering::Release);
 
@@ -33,6 +112,8 @@ pub fn setup_and_run_app(
     };
 
     if new_thread {
+        info!("starting bevy_zeroverse in a new thread");
+
         thread::spawn(startup);
 
         while !ready.load(Ordering::Acquire) {
@@ -44,21 +125,53 @@ pub fn setup_and_run_app(
 }
 
 
-// TODO: add BevyZeroverseViewer struct parameter
 #[pyfunction]
+#[pyo3(signature = (override_args=None))]
 fn initialize(
     py: Python<'_>,
+    override_args: Option<BevyZeroverseConfig>,
 ) {
     py.allow_threads(|| {
-        setup_and_run_app(true);
+        setup_and_run_app(true, override_args);
     });
 }
 
 
+// TODO: support batch dimension (e.g. single array allocation for multiple samples)
+// TODO: add systems for pushing and receiving camera outputs + metadata to python
+// TODO: add options to bevy_zeroverse.next (e.g. render_modes, scene parameters, etc.)
+#[pyfunction]
+fn next() -> PyResult<Sample> {
+    // TODO: advance 'n' frames - requires app reference
+    // TODO: capture frame - requires app system registration to write to textures and readback, triggered by app event after 'n' frames
+    // TODO: send frame
+
+    let receiver = SAMPLE_RECEIVER.get().unwrap();
+    let receiver = receiver.lock().unwrap();
+
+    let timeout = Duration::from_secs(5);
+
+    match receiver.recv_timeout(timeout) {
+        Ok(sample) => Ok(sample),
+        Err(RecvTimeoutError::Timeout) => {
+            Err(PyTimeoutError::new_err("receive operation timed out"))
+        }
+        Err(RecvTimeoutError::Disconnected) => {
+            Err(PyTimeoutError::new_err("channel disconnected"))
+        }
+    }
+}
+
+
 #[pymodule]
 fn bevy_zeroverse(m: &Bound<'_, PyModule>) -> PyResult<()> {
     pyo3_log::init();
 
+    m.add_class::<BevyZeroverseConfig>()?;
+    m.add_class::<Sample>()?;
+    m.add_class::<View>()?;
+
     m.add_function(wrap_pyfunction!(initialize, m)?)?;
+    m.add_function(wrap_pyfunction!(next, m)?)?;
     Ok(())
 }
diff --git a/src/app.rs b/src/app.rs
index 6b86237..fa8e286 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -3,6 +3,7 @@ use bevy::{
     app::AppExit,
     time::Stopwatch,
     render::camera::RenderTarget,
+    winit::WinitPlugin,
 };
 use bevy_args::{
     parse_args,
@@ -15,6 +16,7 @@ use bevy_panorbit_camera::{
     PanOrbitCamera,
     PanOrbitCameraPlugin,
 };
+use pyo3::prelude::*;
 
 use crate::{
     BevyZeroversePlugin,
@@ -54,62 +56,93 @@ use crate::{
 )]
 #[command(about = "bevy_zeroverse viewer", version, long_about = None)]
 #[reflect(Resource)]
-pub struct BevyZeroverseViewer {
+#[pyclass]
+pub struct BevyZeroverseConfig {
     /// enable the bevy inspector
+    #[pyo3(get, set)]
     #[arg(long, default_value = "true")]
     pub editor: bool,
 
     /// no window will be shown
-    #[arg(long, default_value = "true")]
+    #[pyo3(get, set)]
+    #[arg(long, default_value = "false")]
     pub headless: bool,
 
     /// view available material basecolor textures in a grid
+    #[pyo3(get, set)]
     #[arg(long, default_value = "false")]
     pub material_grid: bool,
 
     /// view plücker embeddings
+    #[pyo3(get, set)]
     #[arg(long, default_value = "false")]
     pub plucker_visualization: bool,
 
     /// enable closing the window with the escape key (doesn't work in web)
+    #[pyo3(get, set)]
     #[arg(long, default_value = "true")]
     pub press_esc_close: bool,
 
+    #[pyo3(get, set)]
     #[arg(long, default_value = "1920.0")]
     pub width: f32,
 
+    #[pyo3(get, set)]
     #[arg(long, default_value = "1080.0")]
     pub height: f32,
 
+    #[pyo3(get, set)]
     #[arg(long, default_value = "0")]
     pub num_cameras: usize,
 
     /// display a grid of Zeroverse cameras
+    #[pyo3(get, set)]
     #[arg(long, default_value = "false")]
     pub camera_grid: bool,
 
     /// window title
+    #[pyo3(get, set)]
     #[arg(long, default_value = "bevy_zeroverse")]
     pub name: String,
 
     /// move to the next scene after `regenerate_ms` milliseconds
+    #[pyo3(get, set)]
     #[arg(long, default_value = "0")]
     pub regenerate_ms: u32,
 
     /// automatically rotate the root scene object in the y axis
+    #[pyo3(get, set)]
     #[arg(long, default_value = "0.0")]
     pub yaw_speed: f32,
 
+    #[pyo3(get, set)]
     #[arg(long, value_enum, default_value_t = RenderMode::Color)]
     pub render_mode: RenderMode,
 
+    #[pyo3(get, set)]
     #[arg(long, value_enum, default_value_t = ZeroverseSceneType::Object)]
     pub scene_type: ZeroverseSceneType,
 }
 
-impl Default for BevyZeroverseViewer {
-    fn default() -> BevyZeroverseViewer {
-        BevyZeroverseViewer {
+#[pymethods]
+impl BevyZeroverseConfig {
+    #[new]
+    pub fn new() -> Self {
+        Default::default()
+    }
+
+    fn __str__(&self) -> PyResult<String> {
+        Ok(format!("{:?}", self))
+    }
+
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!("{:?}", self))
+    }
+}
+
+impl Default for BevyZeroverseConfig {
+    fn default() -> BevyZeroverseConfig {
+        BevyZeroverseConfig {
             editor: true,
             headless: false,
             material_grid: false,
@@ -130,11 +163,11 @@ impl Default for BevyZeroverseViewer {
 
 
 pub fn viewer_app(
-    override_args: Option<BevyZeroverseViewer>,
+    override_args: Option<BevyZeroverseConfig>,
 ) -> App {
     let args = match override_args {
         Some(args) => args,
-        None => parse_args::<BevyZeroverseViewer>(),
+        None => parse_args::<BevyZeroverseConfig>(),
     };
 
     let mut app = App::new();
@@ -175,12 +208,19 @@ pub fn viewer_app(
     });
 
     app.insert_resource(ClearColor(Color::srgba(0.0, 0.0, 0.0, 0.0)));
+
     let default_plugins = DefaultPlugins
-        .set(ImagePlugin::default_nearest())
-        .set(WindowPlugin {
+        .set(ImagePlugin::default_nearest());
+
+    let default_plugins = if args.headless {
+        default_plugins
+            .disable::<WinitPlugin>()
+    } else {
+        default_plugins.set(WindowPlugin {
             primary_window,
             ..default()
-        });
+        })
+    };
 
     app.add_plugins(default_plugins);
 
@@ -189,7 +229,7 @@ pub fn viewer_app(
     app.insert_resource(Msaa::Sample8);
 
     if args.editor {
-        app.register_type::<BevyZeroverseViewer>();
+        app.register_type::<BevyZeroverseConfig>();
         app.add_plugins(WorldInspectorPlugin::new());
     }
 
@@ -232,7 +272,7 @@ pub fn viewer_app(
 struct MaterialGridCameraMarker;
 
 fn setup_camera(
-    args: Res<BevyZeroverseViewer>,
+    args: Res<BevyZeroverseConfig>,
     mut commands: Commands,
     material_grid_cameras: Query<Entity, With<MaterialGridCameraMarker>>,
     editor_cameras: Query<Entity, With<EditorCameraMarker>>,
@@ -290,7 +330,7 @@ struct CameraGrid;
 
 fn setup_camera_grid(
     mut commands: Commands,
-    args: Res<BevyZeroverseViewer>,
+    args: Res<BevyZeroverseConfig>,
     camera_grids: Query<Entity, With<CameraGrid>>,
     zeroverse_cameras: Query<
         (Entity, &Camera),
@@ -360,7 +400,7 @@ struct MaterialGrid;
 
 fn setup_material_grid(
     mut commands: Commands,
-    args: Res<BevyZeroverseViewer>,
+    args: Res<BevyZeroverseConfig>,
     standard_materials: Res<Assets<StandardMaterial>>,
     zeroverse_materials: Res<ZeroverseMaterials>,
     material_grids: Query<Entity, With<MaterialGrid>>,
@@ -429,7 +469,7 @@ fn setup_scene(
 
 #[allow(clippy::too_many_arguments)]
 fn regenerate_scene_system(
-    args: Res<BevyZeroverseViewer>,
+    args: Res<BevyZeroverseConfig>,
     keys: Res<ButtonInput<KeyCode>>,
     time: Res<Time>,
     mut regenerate_stopwatch: Local<Stopwatch>,
@@ -451,7 +491,7 @@ fn regenerate_scene_system(
 
 fn rotate_scene(
     time: Res<Time>,
-    args: Res<BevyZeroverseViewer>,
+    args: Res<BevyZeroverseConfig>,
     mut scene_roots: Query<&mut Transform, With<ZeroverseSceneRoot>>,
 ) {
     for mut transform in scene_roots.iter_mut() {
@@ -462,7 +502,7 @@ fn rotate_scene(
 
 
 fn propagate_cli_settings(
-    args: Res<BevyZeroverseViewer>,
+    args: Res<BevyZeroverseConfig>,
     mut plucker_settings: ResMut<ZeroversePluckerSettings>,
     mut scene_settings: ResMut<ZeroverseSceneSettings>,
 ) {
diff --git a/src/io.rs b/src/io.rs
new file mode 100644
index 0000000..8d155cb
--- /dev/null
+++ b/src/io.rs
@@ -0,0 +1,342 @@
+/// Derived from: https://github.com/bevyengine/bevy/pull/5550
+pub mod image_copy {
+    use std::sync::Arc;
+
+    use bevy::prelude::*;
+    use bevy::render::render_asset::RenderAssets;
+    use bevy::render::render_graph::{self, NodeRunError, RenderGraph, RenderGraphContext, RenderLabel};
+    use bevy::render::renderer::{RenderContext, RenderDevice, RenderQueue};
+    use bevy::render::{Extract, RenderApp};
+    use bevy::render::texture::GpuImage;
+
+    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<Assets<Image>>,
+        render_device: Res<RenderDevice>,
+    ) {
+        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();
+        }
+    }
+
+    #[derive(Debug, Hash, PartialEq, Eq, Clone, RenderLabel)]
+    pub struct ImageCopyLabel;
+
+    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_mut().get_resource_mut::<RenderGraph>().unwrap();
+
+            graph.add_node(ImageCopyLabel, ImageCopyDriver);
+            graph.add_node_edge(ImageCopyLabel, bevy::render::graph::CameraDriverLabel);
+        }
+    }
+
+    #[derive(Clone, Default, Resource, Deref, DerefMut)]
+    pub struct ImageCopiers(pub Vec<ImageCopier>);
+
+    #[derive(Clone, Component)]
+    pub struct ImageCopier {
+        buffer: Buffer,
+        enabled: Arc<AtomicBool>,
+        src_image: Handle<Image>,
+        dst_image: Handle<Image>,
+    }
+
+    impl ImageCopier {
+        pub fn new(
+            src_image: Handle<Image>,
+            dst_image: Handle<Image>,
+            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<Query<&ImageCopier>>,
+    ) {
+        commands.insert_resource(ImageCopiers(
+            image_copiers.iter().cloned().collect::<Vec<ImageCopier>>(),
+        ));
+    }
+
+    #[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::<ImageCopiers>().unwrap();
+            let gpu_images = world.get_resource::<RenderAssets<GpuImage>>().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_copy_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,
+                    height: src_image.size.y,
+                    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::<RenderQueue>().unwrap();
+                render_queue.submit(std::iter::once(encoder.finish()));
+            }
+
+            Ok(())
+        }
+    }
+}
+
+pub mod scene {
+    use std::path::PathBuf;
+
+    use bevy::{
+        app::AppExit,
+        prelude::*,
+        render::{camera::RenderTarget, renderer::RenderDevice},
+    };
+    use wgpu::{Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages};
+
+    use super::image_copy::ImageCopier;
+
+
+    #[derive(Component, Default)]
+    pub struct CaptureCamera;
+
+    #[derive(Component, Deref, DerefMut)]
+    struct ImageToSave(Handle<Image>);
+
+    pub struct CaptureFramePlugin;
+    impl Plugin for CaptureFramePlugin {
+        fn build(&self, app: &mut App) {
+            app.add_systems(PostUpdate, update);
+        }
+    }
+
+    #[derive(Debug, Default, Resource, Event)]
+    pub struct SceneController {
+        state: SceneState,
+        name: String,
+        width: u32,
+        height: u32,
+        single_image: bool,
+        output_directory: String,
+    }
+
+    impl SceneController {
+        pub fn new(
+            width:u32,
+            height:u32,
+            single_image: bool,
+            name: String,
+            output_directory: String,
+        ) -> SceneController {
+            SceneController {
+                state: SceneState::BuildScene,
+                name,
+                width,
+                height,
+                single_image,
+                output_directory,
+            }
+        }
+    }
+
+    #[derive(Debug, Default)]
+    pub enum SceneState {
+        #[default]
+        BuildScene,
+        Render(u32),
+    }
+
+    pub fn setup_render_target(
+        commands: &mut Commands,
+        images: &mut ResMut<Assets<Image>>,
+        render_device: &Res<RenderDevice>,
+        scene_controller: &mut ResMut<SceneController>,
+        pre_roll_frames: u32,
+    ) -> RenderTarget {
+        let size = Extent3d {
+            width: scene_controller.width,
+            height: scene_controller.height,
+            ..Default::default()
+        };
+
+        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);
+
+        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);
+        RenderTarget::Image(render_target_image_handle)
+    }
+
+    fn update(
+        images_to_save: Query<&ImageToSave>,
+        mut images: ResMut<Assets<Image>>,
+        mut scene_controller: ResMut<SceneController>,
+        mut app_exit_writer: EventWriter<AppExit>,
+    ) {
+        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 img = match img_bytes.clone().try_into_dynamic() {
+                        Ok(img) => img.to_rgba8(),
+                        Err(e) => panic!("failed to create image buffer {e:?}"),
+                    };
+
+                    let images_dir = PathBuf::from(&scene_controller.output_directory);
+                    std::fs::create_dir_all(&images_dir).unwrap();
+
+                    let image_path = images_dir.join(format!("{}.png", scene_controller.name));
+                    if let Err(e) = img.save(image_path){
+                        panic!("failed to save image: {}", e);
+                    };
+                }
+                if scene_controller.single_image {
+                    app_exit_writer.send(AppExit::Success);
+                }
+            } else {
+                scene_controller.state = SceneState::Render(n - 1);
+            }
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 842cec4..5158407 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@ use bevy::prelude::*;
 
 pub mod app;
 pub mod camera;
+pub mod io;
 pub mod manifold;
 pub mod material;
 pub mod mesh;
diff --git a/src/render/mod.rs b/src/render/mod.rs
index 4e60571..3ae3bdc 100644
--- a/src/render/mod.rs
+++ b/src/render/mod.rs
@@ -7,6 +7,7 @@ use bevy_args::{
     Serialize,
     ValueEnum,
 };
+use pyo3::prelude::*;
 
 use crate::primitive::process_primitives;
 
@@ -26,6 +27,7 @@ pub mod normal;
     ValueEnum,
 )]
 #[reflect(Resource)]
+#[pyclass(eq, eq_int)]
 pub enum RenderMode {
     #[default]
     Color,
diff --git a/src/scene/mod.rs b/src/scene/mod.rs
index 9795af9..8ac9cb5 100644
--- a/src/scene/mod.rs
+++ b/src/scene/mod.rs
@@ -4,6 +4,7 @@ use bevy_args::{
     Serialize,
     ValueEnum,
 };
+use pyo3::prelude::*;
 
 pub mod lighting;
 pub mod object;
@@ -48,6 +49,7 @@ pub struct ZeroverseSceneRoot;
     Reflect,
     ValueEnum,
 )]
+#[pyclass(eq, eq_int)]
 pub enum ZeroverseSceneType {
     #[default]
     Object,