Basic Bevy Mipmap Generation

main.rs

// License: Apache-2.0 / MIT

use std::num::NonZeroU8;

use bevy::{
    asset::HandleId,
    prelude::*,
    render::{
        render_resource::{
            Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
        },
        texture::ImageSampler,
    },
};
use image::{imageops::FilterType, DynamicImage, RgbaImage};

fn main() {
    let mut app = App::new();
    app.add_plugins(DefaultPlugins)
        .insert_resource(MipmapHandleSets(Vec::new()))
        .add_startup_system(setup)
        .add_system(create_mipmaps);

    app.run();
}

fn create_test_image(size: u32, cx: f32, cy: f32) -> Vec<u8> {
    use std::iter;

    (0..size * size)
        .flat_map(|id| {
            // get high five for recognizing this ;)
            let mut x = 4.0 * (id % size) as f32 / (size - 1) as f32 - 2.0;
            let mut y = 2.0 * (id / size) as f32 / (size - 1) as f32 - 1.0;
            let mut count = 0;
            while count < 0xFF && x * x + y * y < 4.0 {
                let old_x = x;
                x = x * x - y * y + cx;
                y = 2.0 * old_x * y + cy;
                count += 1;
            }
            iter::once(0xFF - (count * 2) as u8)
                .chain(iter::once(0xFF - (count * 5) as u8))
                .chain(iter::once(0xFF - (count * 13) as u8))
                .chain(iter::once(std::u8::MAX))
        })
        .collect()
}

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

    let texture_descriptor = TextureDescriptor {
        label: None,
        size,
        dimension: TextureDimension::D2,
        format: TextureFormat::Bgra8UnormSrgb,
        mip_level_count: 1,
        sample_count: 1,
        usage: TextureUsages::TEXTURE_BINDING
            | TextureUsages::COPY_DST
            | TextureUsages::RENDER_ATTACHMENT,
    };

    let mut image = Image {
        texture_descriptor,
        ..Default::default()
    };

    image.data = create_test_image(size.width, -0.8, 0.156);

    let mipmaphandle = MipmapHandleSet::from(images.add(image));

    let image_handle_a = mipmaphandle.dst.clone();

    mipmaphandles.push(mipmaphandle);

    // plane
    commands.spawn_bundle(PbrBundle {
        mesh: meshes.add(Mesh::from(shape::Plane { size: 20.0 })),
        material: materials.add(StandardMaterial::from(image_handle_a)),
        ..default()
    });
    // light
    commands.spawn_bundle(PointLightBundle {
        point_light: PointLight {
            intensity: 1500.0,
            shadows_enabled: true,
            ..default()
        },
        transform: Transform::from_xyz(4.0, 8.0, 4.0),
        ..default()
    });
    // camera
    commands.spawn_bundle(Camera3dBundle {
        transform: Transform::from_xyz(0.0, 0.5, 10.0).looking_at(Vec3::ZERO, Vec3::Y),
        ..default()
    });
}

#[derive(Deref, DerefMut)]
struct MipmapHandleSets(Vec<MipmapHandleSet>);

struct MipmapHandleSet {
    src: Handle<Image>,
    dst: Handle<Image>,
}

impl MipmapHandleSet {
    fn from(src_image: Handle<Image>) -> MipmapHandleSet {
        MipmapHandleSet {
            src: src_image,
            dst: Handle::weak(HandleId::random::<Image>()),
        }
    }
}

fn create_mipmaps(mut mipmaphandles: ResMut<MipmapHandleSets>, mut images: ResMut<Assets<Image>>) {
    mipmaphandles.retain(|mipmaphandle| {
        let (image_data, texture_descriptor) = if let Some(image) = images.get(&mipmaphandle.src) {
            let mut image_data = image.data.clone();
            let mut mip_level_count = 1;
            let mut texture_descriptor = image.texture_descriptor.clone();
            let mut width = texture_descriptor.size.width;
            let mut height = texture_descriptor.size.height;
            let mut img = DynamicImage::ImageRgba8(
                RgbaImage::from_raw(width, height, image.data.clone()).unwrap(),
            );
            while width / 2 >= 2 && height / 2 >= 2 {
                mip_level_count += 1;
                width /= 2;
                height /= 2;
                img = img.resize_exact(width, height, FilterType::CatmullRom);
                image_data.append(&mut img.as_rgba8().unwrap().as_raw().clone());
            }

            texture_descriptor.mip_level_count = mip_level_count;
            (image_data, texture_descriptor)
        } else {
            return true;
        };

        let mut descriptor = ImageSampler::linear_descriptor();
        descriptor.anisotropy_clamp = NonZeroU8::new(16); // Use 16x anisotropic filtering
        let _ = images.set(
            mipmaphandle.dst.clone(),
            Image {
                texture_descriptor,
                sampler_descriptor: ImageSampler::Descriptor(descriptor),
                data: image_data,
            },
        );
        false
    })
}