Copy from GPU texture to CPU with using channel example

main.rs

// License: Apache-2.0 / MIT

use std::sync::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::{RenderApp, RenderStage};

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

#[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>>,
) {
    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);

    setup_image_copy(
        &mut commands,
        render_target_image_handle.clone(),
        cpu_image_handle.clone(),
    );

    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 receivers: Query<&mut ImageReceiver>, mut images: ResMut<Assets<Image>>) {
    for image_receiver in receivers.iter_mut() {
        if let Ok(receiver) = image_receiver.receiver.try_lock() {
            if !receiver.is_empty() {
                if let Some(mut image) = images.get_mut(&image_receiver.image) {
                    if let Ok(received_image) = receiver.recv() {
                        image.data = received_image;
                    }
                }
            }
        }
    }
}

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 ImageSender {
    image: Handle<Image>,
    sender: Sender<Vec<u8>>,
}

#[derive(Component)]
pub struct ImageReceiver {
    image: Handle<Image>,
    receiver: Mutex<Receiver<Vec<u8>>>,
}

fn setup_image_copy(commands: &mut Commands, src_image: Handle<Image>, dst_image: Handle<Image>) {
    let (sender, receiver) = bounded(1);
    commands
        .spawn()
        .insert(ImageSender {
            image: src_image,
            sender,
        })
        .insert(ImageReceiver {
            image: dst_image,
            receiver: Mutex::new(receiver),
        });
}

pub fn image_copy_extract(mut commands: Commands, senders: Query<&mut ImageSender>) {
    commands.insert_resource(senders.iter().cloned().collect::<Vec<ImageSender>>());
}

#[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_senders = world.get_resource::<Vec<ImageSender>>().unwrap();
        let gpu_images = world.get_resource::<RenderAssets<Image>>().unwrap();

        for image_sender in image_senders.iter() {
            let src_image = gpu_images.get(&image_sender.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,
            };

            let cpu_buffer = render_context
                .render_device
                .create_buffer(&BufferDescriptor {
                    label: None,
                    size: padded_bytes_per_row as u64 * src_image.size.y as u64,
                    usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                });

            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()));

            let large_buffer_slice = cpu_buffer.slice(..);
            render_context
                .render_device
                .map_buffer(&large_buffer_slice, MapMode::Read);
            {
                let large_padded_buffer = large_buffer_slice.get_mapped_range();
                if !image_sender.sender.is_full() {
                    image_sender
                        .sender
                        .send(large_padded_buffer.to_vec())
                        .unwrap();
                }
            }
            cpu_buffer.unmap();
        }

        Ok(())
    }
}