GPU image copy plugin example

main.rs

// License: Apache-2.0 / MIT
// Shows how to render to a buffered texture that is copied and can then be used on the same render layer.

use bevy::prelude::*;
use bevy::render::camera::RenderTarget;
use bevy::render::render_asset::RenderAssets;
use bevy::render::render_graph::{self, NodeRunError, RenderGraph, RenderGraphContext};
use bevy::render::renderer::{RenderContext, RenderQueue};
use bevy::render::{RenderApp, RenderStage};

use bevy::render::render_resource::{
    CommandEncoderDescriptor, Extent3d, TextureDescriptor, TextureDimension, TextureFormat,
    TextureUsages,
};

#[derive(Component, Default)]
pub struct CaptureCamera;

#[derive(Component)]
struct Cube {
    rotate_speed: f32,
}

#[derive(Component, Deref, DerefMut)]
struct ImagesToBeCopied(ImageCopyCommand);

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_plugin(ImageCopyPlugin)
        .add_startup_system(setup)
        .add_system(cube_rotator_system)
        .add_system(copy_images)
        .run();
}

fn setup(
    mut commands: Commands,
    mut images: ResMut<Assets<Image>>,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
) {
    let size = Extent3d {
        width: 512,
        height: 512,
        ..Default::default()
    };

    // This is the texture that will be rendered to.
    let mut image = Image {
        texture_descriptor: TextureDescriptor {
            label: None,
            size,
            dimension: TextureDimension::D2,
            format: TextureFormat::Rgba8UnormSrgb,
            mip_level_count: 1,
            sample_count: 1,
            usage: TextureUsages::TEXTURE_BINDING
                | TextureUsages::COPY_DST
                | TextureUsages::COPY_SRC
                | TextureUsages::RENDER_ATTACHMENT,
        },
        ..Default::default()
    };
    image.resize(size);
    let src_image = images.add(image);

    // This is the buffered texture that copied to.
    let mut image = Image {
        texture_descriptor: TextureDescriptor {
            label: None,
            size,
            dimension: TextureDimension::D2,
            format: TextureFormat::Rgba8UnormSrgb,
            mip_level_count: 1,
            sample_count: 1,
            usage: TextureUsages::TEXTURE_BINDING
                | TextureUsages::COPY_DST
                | TextureUsages::COPY_SRC
                | TextureUsages::RENDER_ATTACHMENT,
        },
        ..Default::default()
    };
    image.resize(size);
    let dst_image = images.add(image);

    let cube_handle = meshes.add(Mesh::from(shape::Cube { size: 0.25 }));
    let cube_material_handle = materials.add(StandardMaterial {
        base_color: Color::rgb(0.8, 0.7, 0.6),
        reflectance: 0.02,
        unlit: false,
        ..default()
    });

    commands.spawn().insert(ImagesToBeCopied(ImageCopyCommand {
        src_image: src_image.clone(),
        dst_image: dst_image.clone(),
    }));

    // The cube that will be rendered to the texture.
    commands
        .spawn_bundle(PbrBundle {
            mesh: cube_handle,
            material: cube_material_handle,
            transform: Transform::from_translation(Vec3::new(0.0, 0.25, 0.0)),
            ..default()
        })
        .insert(Cube { rotate_speed: 1.0 });

    commands.spawn_bundle(PointLightBundle {
        transform: Transform::from_translation(Vec3::new(0.0, 0.0, 10.0)),
        ..default()
    });

    let cube_size = 0.25;
    let cube_handle = meshes.add(Mesh::from(shape::Box::new(cube_size, cube_size, cube_size)));

    // This material has the texture that has been rendered.
    let material_handle = materials.add(StandardMaterial {
        base_color_texture: Some(dst_image.clone()),
        reflectance: 0.02,
        unlit: false,
        ..default()
    });
    // Main pass cube, with material containing the rendered first pass texture.
    commands
        .spawn_bundle(PbrBundle {
            mesh: cube_handle,
            material: material_handle,
            transform: Transform {
                translation: Vec3::new(0.0, 0.5, 0.0),
                rotation: Quat::from_rotation_x(-std::f32::consts::PI / 5.0),
                ..default()
            },
            ..default()
        })
        .insert(Cube { rotate_speed: 1.3 });

    commands
        .spawn_bundle(Camera3dBundle {
            transform: Transform::from_xyz(0.7, 0.7, 1.0)
                .looking_at(Vec3::new(0.0, 0.3, 0.0), Vec3::Y),
            ..default()
        })
        .with_children(|parent| {
            let render_target = RenderTarget::Image(src_image);
            parent.spawn_bundle(Camera3dBundle {
                camera: Camera {
                    target: render_target,
                    ..default()
                },
                ..default()
            });
        });
}

/// Rotates the cubes
fn cube_rotator_system(time: Res<Time>, mut query: Query<(&mut Transform, &Cube)>) {
    for (mut transform, cube) in query.iter_mut() {
        transform.rotation *= Quat::from_rotation_x(cube.rotate_speed * time.delta_seconds());
        transform.rotation *= Quat::from_rotation_y(cube.rotate_speed * time.delta_seconds());
    }
}

fn copy_images(
    mut copy_commands: ResMut<ImageCopyCommands>,
    to_be_copied: Query<&ImagesToBeCopied>,
) {
    for copy_command in to_be_copied.iter() {
        copy_commands.push(copy_command.0.clone());
    }
}

//------------------------------------------------------
// ImageCopyPlugin -------------------------------------
//------------------------------------------------------

pub const IMAGE_COPY: &str = "image_copy";

pub struct ImageCopyPlugin;
impl Plugin for ImageCopyPlugin {
    fn build(&self, app: &mut App) {
        let render_app = app
            .insert_resource(ImageCopyCommands::default())
            .sub_app_mut(RenderApp);

        render_app.add_system_to_stage(RenderStage::Extract, image_copy_extract);

        let mut graph = render_app.world.get_resource_mut::<RenderGraph>().unwrap();

        graph.add_node(IMAGE_COPY, ImageCopyDriver::default());
        graph
            .add_node_edge(IMAGE_COPY, bevy::render::main_graph::node::CAMERA_DRIVER)
            .unwrap();
    }
}

#[derive(Clone)]
pub struct ImageCopyCommand {
    pub src_image: Handle<Image>,
    pub dst_image: Handle<Image>,
}

#[derive(Clone, Default, DerefMut, Deref)]
pub struct ImageCopyCommands(Vec<ImageCopyCommand>);

pub fn image_copy_extract(mut commands: Commands, mut copy_commands: ResMut<ImageCopyCommands>) {
    commands.insert_resource(copy_commands.clone());
    copy_commands.clear();
}

#[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 copy_commands = world.get_resource::<ImageCopyCommands>().unwrap();
        for image_copy in copy_commands.iter() {
            let gpu_images = world.get_resource::<RenderAssets<Image>>().unwrap();
            let src_image = gpu_images.get(&image_copy.src_image).unwrap();
            let dst_image = gpu_images.get(&image_copy.dst_image).unwrap();

            let mut encoder = render_context
                .render_device
                .create_command_encoder(&CommandEncoderDescriptor::default());

            encoder.copy_texture_to_texture(
                src_image.texture.as_image_copy(),
                dst_image.texture.as_image_copy(),
                Extent3d {
                    width: src_image.size.x as u32,
                    height: src_image.size.y as u32,
                    depth_or_array_layers: 1,
                },
            );

            let render_queue = world.get_resource::<RenderQueue>().unwrap();
            render_queue.submit(std::iter::once(encoder.finish()));
        }

        Ok(())
    }
}