Convert 32bit, 16bit, or RGB9E5 images to either 16bit or RGB9E5 KTX2 with ZSTD supercompression in 3D

main.rs

// Convert 32bit, 16bit, or RGB9E5 images in 3D or vertical strip format
// to either 16bit or RGB9E5 KTX2 with ZSTD supercompression in 3D

// Example cargo run -- --inputs tony_mc_mapface.dds,blender_-11_12.exr,AgX-default_contrast.exr --outputs tony_mc_mapface.ktx2,blender_-11_12.ktx2,AgX-default_contrast.ktx2

// [dependencies]
// Should be able to use bevy 0.10 instead when it is release. Needed for exr support. (TODO rewrite without bevy)
// bevy = { git = "https://github.com/DGriffin91/bevy", branch = "tonemap_options", features = ["zstd", "ktx2", "exr", "dds"] }
// half = { version = "2.1" }
// ktx2 = { git = "https://github.com/BVE-Reborn/ktx2", rev = "4a7cc48ffa4deb3aa1ef5b453292220489908fa1" }
// zstd = "0.12"
// clap = { version = "4.1", features = ["derive"] }

use std::path::Path;

use bevy::{prelude::Image, render::render_resource::TextureFormat};
use ktx2::SupercompressionScheme;

use std::{path::PathBuf, time::Duration};

use bevy::{
    app::{AppExit, ScheduleRunnerSettings},
    prelude::*,
    render::render_resource::{Extent3d, TextureDescriptor, TextureDimension},
};
use clap::Parser;

/// Encode images in ktx2 files.
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Input file paths
    #[arg(short, long, value_delimiter = ',')]
    inputs: Vec<PathBuf>,

    /// Output file paths
    #[arg(short, long, value_delimiter = ',')]
    outputs: Vec<PathBuf>,
}

fn main() {
    let args = Args::parse();

    if args.inputs.is_empty() {
        panic!("No input paths provided");
    }

    if args.outputs.is_empty() {
        panic!("No output paths provided");
    }

    if args.inputs.len() != args.outputs.len() {
        panic!("Input and output path lengths don't match");
    }

    let mut app = App::new();
    // TODO don't be ridiculous
    app.insert_resource(ScheduleRunnerSettings::run_loop(Duration::from_secs_f64(
        1.0 / 100.0,
    )))
    .add_plugins(
        MinimalPlugins
            .build()
            .add(AssetPlugin::default())
            .add(ImagePlugin::default()),
    )
    .add_system(convert);

    for (input, output) in args.inputs.iter().zip(args.outputs.iter()) {
        let asset_server = app.world.resource_mut::<AssetServer>();
        // using canonicalize to avoid being relative to the asset folder
        let image_h = asset_server.load(std::fs::canonicalize(input).unwrap());
        app.world.spawn(ImageToConvert {
            image_h,
            output_path: PathBuf::from(output),
        });
    }

    app.run();
}

#[derive(Component)]
struct Converted;

#[derive(Component)]
struct ImageToConvert {
    image_h: Handle<Image>,
    output_path: PathBuf,
}

fn convert(
    mut commands: Commands,
    query: Query<(Entity, &ImageToConvert), Without<Converted>>,
    mut images: ResMut<Assets<Image>>,
    mut app_exit_events: EventWriter<AppExit>,
) {
    if query.is_empty() {
        app_exit_events.send(AppExit);
    }
    for (entity, conv) in &query {
        if let Some(mut image) = images.get_mut(&conv.image_h) {
            println!(
                "Converting {}, {:?}, format:{:?}",
                &conv.output_path.display(),
                image.texture_descriptor.size,
                image.texture_descriptor.format,
            );
            let block_size = image.size().x as u32;

            image.texture_descriptor = TextureDescriptor {
                label: Some("tonemapping lut"),
                size: Extent3d {
                    width: block_size,
                    height: block_size,
                    depth_or_array_layers: block_size,
                },
                mip_level_count: 1,
                sample_count: 1,
                dimension: TextureDimension::D3,
                format: image.texture_descriptor.format,
                usage: image.texture_descriptor.usage,
                view_formats: image.texture_descriptor.view_formats,
            };
            write_ktx2(image, &conv.output_path);
            commands.entity(entity).insert(Converted);
        }
    }
}

pub fn write_ktx2(image: &Image, output_path: &Path) {
    if image.is_compressed() {
        panic!("Only uncompressed images supported");
    }

    let mut image = image.clone();

    if let TextureFormat::Rgba32Float = image.texture_descriptor.format {
        let mut f16data = Vec::new();
        for n in to_vec_f32_from_byte_slice(&image.data).iter() {
            f16data.push(half::f16::from_f32(*n));
        }
        image.data = f16_to_bytes(&f16data).to_vec();
        image.texture_descriptor.format = TextureFormat::Rgba16Float;
    }

    let data = WriterLevel {
        uncompressed_length: image.data.len(),
        bytes: zstd::bulk::compress(&image.data, 0).unwrap(),
    };

    let format = match image.texture_descriptor.format {
        TextureFormat::Rgba16Float => ktx2::Format::R16G16B16A16_SFLOAT,
        TextureFormat::Rgb9e5Ufloat => ktx2::Format::E5B9G9R9_UFLOAT_PACK32,
        _ => panic!("format not setup"),
    };

    let type_size = match image.texture_descriptor.format {
        TextureFormat::Rgba16Float => 2,
        TextureFormat::Rgb9e5Ufloat => 4,
        _ => panic!("format not setup"),
    };

    dbg!(&image.texture_descriptor);

    let header = Header {
        format: Some(format),
        type_size,
        pixel_width: image.texture_descriptor.size.width,
        pixel_height: image.texture_descriptor.size.height,
        pixel_depth: image.texture_descriptor.size.depth_or_array_layers,
        layer_count: 0,
        face_count: 1,
        supercompression_scheme: Some(SupercompressionScheme::Zstandard),
    };

    dbg!(&header);

    // https://github.khronos.org/KTX-Specification/
    let writer = KTX2Writer {
        header,
        dfd_bytes: u32_to_bytes(&[0u32, 0, 2]),
        levels_descending: vec![data],
    };

    writer
        .write(&mut std::fs::File::create(output_path).unwrap())
        .unwrap();
}

pub fn to_vec_f32_from_byte_slice(vecs: &[u8]) -> &[f32] {
    unsafe { std::slice::from_raw_parts(vecs.as_ptr() as *const _, vecs.len() / 4) }
}

pub fn u32_to_bytes(vecs: &[u32]) -> &[u8] {
    unsafe { std::slice::from_raw_parts(vecs.as_ptr() as *const _, vecs.len() * 4) }
}

pub fn f16_to_bytes(vecs: &[half::f16]) -> &[u8] {
    unsafe { std::slice::from_raw_parts(vecs.as_ptr() as *const _, vecs.len() * 2) }
}

pub struct KTX2Writer<'a> {
    pub header: Header,
    pub dfd_bytes: &'a [u8],
    pub levels_descending: Vec<WriterLevel>,
}

impl<'a> KTX2Writer<'a> {
    pub fn write<T: std::io::Write>(&self, writer: &mut T) -> std::io::Result<()> {
        let dfd_offset =
            ktx2::Header::LENGTH + self.levels_descending.len() * ktx2::LevelIndex::LENGTH;

        writer.write_all(
            &ktx2::Header {
                format: self.header.format,
                type_size: if self.header.supercompression_scheme.is_some() {
                    1
                } else {
                    self.header.type_size
                },
                pixel_width: self.header.pixel_width,
                pixel_height: self.header.pixel_height,
                pixel_depth: self.header.pixel_depth,
                layer_count: self.header.layer_count,
                face_count: self.header.face_count,
                supercompression_scheme: self.header.supercompression_scheme,
                level_count: self.levels_descending.len() as u32,
                index: ktx2::Index {
                    dfd_byte_length: self.dfd_bytes.len() as u32,
                    kvd_byte_length: 0,
                    sgd_byte_length: 0,
                    dfd_byte_offset: dfd_offset as u32,
                    kvd_byte_offset: 0,
                    sgd_byte_offset: 0,
                },
            }
            .as_bytes()[..],
        )?;

        let mut offset = dfd_offset + self.dfd_bytes.len();

        let mut levels = self
            .levels_descending
            .iter()
            .rev()
            .map(|level| {
                let index = ktx2::LevelIndex {
                    byte_offset: offset as u64,
                    byte_length: level.bytes.len() as u64,
                    uncompressed_byte_length: level.uncompressed_length as u64,
                };

                offset += level.bytes.len();

                index
            })
            .collect::<Vec<_>>();

        levels.reverse();

        for level in levels {
            writer.write_all(&level.as_bytes())?;
        }

        writer.write_all(self.dfd_bytes)?;

        for level in self.levels_descending.iter().rev() {
            writer.write_all(&level.bytes)?;
        }

        Ok(())
    }
}

pub struct WriterLevel {
    pub uncompressed_length: usize,
    pub bytes: Vec<u8>,
}

#[derive(Debug)]
pub struct Header {
    pub format: Option<ktx2::Format>,
    pub type_size: u32,
    pub pixel_width: u32,
    pub pixel_height: u32,
    pub pixel_depth: u32,
    pub layer_count: u32,
    pub face_count: u32,
    pub supercompression_scheme: Option<ktx2::SupercompressionScheme>,
}