Copy from GPU texture to CPU with using mutex example

main.rs

// License: Apache-2.0 / MIT

use std::sync::{Arc, Mutex};

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, RenderDevice, RenderQueue};
use bevy::render::{Extract, RenderApp, RenderStage};

use bevy::render::render_resource::{
    Buffer, BufferDescriptor, BufferUsages, CommandEncoderDescriptor, Extent3d, ImageCopyBuffer,
    ImageDataLayout, MapMode, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
};

use std::sync::atomic::{AtomicBool, Ordering};

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

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

#[derive(Component, Deref, DerefMut)]
struct ImageToSave(Handle<Image>);

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

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

    // This is the texture that will be rendered to.
    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::TEXTURE_BINDING
                | TextureUsages::COPY_DST
                | TextureUsages::COPY_SRC
                | TextureUsages::RENDER_ATTACHMENT,
        },
        ..Default::default()
    };
    render_target_image.resize(size);
    let render_target_image_handle = images.add(render_target_image);

    // This is the texture that will be copied to.
    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::TEXTURE_BINDING
                | TextureUsages::COPY_DST
                | TextureUsages::COPY_SRC
                | TextureUsages::RENDER_ATTACHMENT,
        },
        ..Default::default()
    };
    cpu_image.resize(size);
    let cpu_image_handle = images.add(cpu_image);

    commands.spawn().insert(ImageCopier::new(
        render_target_image_handle.clone(),
        cpu_image_handle.clone(),
        size,
        &render_device,
    ));

    commands
        .spawn()
        .insert(ImageToSave(cpu_image_handle.clone()));

    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_bundle(PbrBundle {
            mesh: cube_handle,
            material: cube_material_handle,
            transform: Transform::from_translation(Vec3::new(0.0, 0.0, 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()
    });

    commands
        .spawn_bundle(Camera3dBundle {
            transform: Transform::from_xyz(0.7, 0.0, 1.0)
                .looking_at(Vec3::new(0.0, 0.0, 0.0), Vec3::Y),
            ..default()
        })
        .with_children(|parent| {
            parent.spawn_bundle(Camera3dBundle {
                camera: Camera {
                    target: RenderTarget::Image(render_target_image_handle),
                    ..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 get_image_data(images_to_save: Query<&ImageToSave>, images: Res<Assets<Image>>) {
    for image in images_to_save.iter() {
        image::save_buffer(
            "test.png",
            &images.get(image).unwrap().data,
            512,
            512,
            image::ColorType::Rgba8,
        )
        .unwrap();
    }
}

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

pub fn receive_images(
    mut image_copiers: Query<&mut ImageCopier>,
    mut images: ResMut<Assets<Image>>,
    render_device: Res<RenderDevice>,
) {
    for image_copier in image_copiers.iter_mut() {
        if !image_copier.enabled() {
            continue;
        }
        if let Ok(cpu_buffer) = image_copier.buffer.lock() {
            let large_buffer_slice = cpu_buffer.slice(..);
            render_device.map_buffer(&large_buffer_slice, MapMode::Read);
            {
                let large_padded_buffer = large_buffer_slice.get_mapped_range();
                if let Some(mut image) = images.get_mut(&image_copier.dst_image) {
                    image.data = large_padded_buffer.to_vec();
                }
            }
            cpu_buffer.unmap();
        }
    }
}

pub const IMAGE_COPY: &str = "image_copy";

pub struct ImageCopyPlugin;
impl Plugin for ImageCopyPlugin {
    fn build(&self, app: &mut App) {
        let render_app = app.add_system(receive_images).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, Component)]
pub struct ImageCopier {
    buffer: Arc<Mutex<Buffer>>,
    enabled: Arc<AtomicBool>,
    src_image: Handle<Image>,
    dst_image: Handle<Image>,
}

impl ImageCopier {
    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: Arc::new(Mutex::new(cpu_buffer)),
            src_image,
            dst_image,
            enabled: Arc::new(AtomicBool::new(true)),
        }
    }

    fn enable(&self) {
        self.enabled.store(true, Ordering::Relaxed)
    }

    fn disable(&self) {
        self.enabled.store(false, Ordering::Relaxed)
    }

    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(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::<Vec<ImageCopier>>().unwrap();
        let gpu_images = world.get_resource::<RenderAssets<Image>>().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 padded_bytes_per_row =
                RenderDevice::align_copy_bytes_per_row((src_image.size.x) as usize) * 4;

            let texture_extent = Extent3d {
                width: src_image.size.x as u32,
                height: src_image.size.y as u32,
                depth_or_array_layers: 1,
            };

            if let Ok(cpu_buffer) = image_copier.buffer.lock() {
                encoder.copy_texture_to_buffer(
                    src_image.texture.as_image_copy(),
                    ImageCopyBuffer {
                        buffer: &cpu_buffer,
                        layout: ImageDataLayout {
                            offset: 0,
                            bytes_per_row: Some(
                                std::num::NonZeroU32::new(padded_bytes_per_row as u32).unwrap(),
                            ),
                            rows_per_image: None,
                        },
                    },
                    texture_extent,
                );
            }

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

        Ok(())
    }
}