PbrTexturePacker.java

package com.jayfella.modular.debug;

import com.jme3.asset.AssetManager;
import com.jme3.math.ColorRGBA;
import com.jme3.texture.Image;
import com.jme3.texture.Texture;
import com.jme3.texture.Texture2D;
import com.jme3.texture.image.ColorSpace;
import com.jme3.texture.image.ImageRaster;
import com.jme3.util.BufferUtils;

import java.nio.ByteBuffer;

public class PbrTexturePacker {

    /**
     * Packs three images into one image.
     * @param aoImage        Ambient Occlusion image - stored in the RED channel.
     * @param metallicImage  Metallic image - stored in the GREEN channel.
     * @param roughnessImage Roughness image - stored in the BLUE channel.
     * @return an Image with the three input images packed together.
     */
    public static Image packAoMetalRoughness(Image aoImage, Image metallicImage, Image roughnessImage) {

        int width = aoImage.getWidth();
        int height = aoImage.getHeight();

        if (metallicImage.getWidth() != width || roughnessImage.getWidth() != width) {
            throw new IllegalArgumentException("Images are not the same width.");
        }

        if (metallicImage.getHeight() != height || roughnessImage.getHeight() != height) {
            throw new IllegalArgumentException("Images are not the same height.");
        }

        ImageRaster aoRaster = ImageRaster.create(aoImage);
        ImageRaster roughnessRaster = ImageRaster.create(roughnessImage);
        ImageRaster metallicRaster = ImageRaster.create(metallicImage);

        ByteBuffer buffer = BufferUtils.createByteBuffer(width * height * 4);
        Image result = new Image(Image.Format.RGB8, width, height, buffer, ColorSpace.sRGB);

        ImageRaster raster = ImageRaster.create(result);

        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {

                ColorRGBA pixelColor = new ColorRGBA();
                pixelColor.r = aoRaster.getPixel(x, y).r; // greyscale, so just grab any of the RGB's
                pixelColor.g = roughnessRaster.getPixel(x, y).r; // greyscale, so just grab any of the RGB's
                pixelColor.b = metallicRaster.getPixel(x, y).r; // greyscale, so just grab any of the RGB's
                pixelColor.a = 1.0f; // unused

                raster.setPixel(x, y, pixelColor);
            }
        }

        return result;

    }

    public static Texture packAoMetalRoughness(Texture aoTex, Texture metallicTex, Texture roughnessTex) {
        Image image = packAoMetalRoughness(aoTex.getImage(), metallicTex.getImage(), roughnessTex.getImage());
        return new Texture2D(image);
    }

    public static Texture packAoMetalRoughness(AssetManager assetManager, String aoPath, String metallicPath, String roughnessPath) {
        return packAoMetalRoughness(
                assetManager.loadTexture(aoPath),
                assetManager.loadTexture(metallicPath),
                assetManager.loadTexture(roughnessPath));
    }


    /**
     * Packs a Normal map and Parallax map into a single image.
     * @param normalImage   Normal Map image - stored in the RGB channels.
     * @param parallaxImage Parallax image - stored in the ALPHA channel.
     * @return an Image with the two input images packed together.
     */
    public static Image packNormalParallax(Image normalImage, Image parallaxImage) {

        if (normalImage.getWidth() != normalImage.getWidth()) {
            throw new IllegalArgumentException("Images are not the same width.");
        }

        if (parallaxImage.getHeight() != parallaxImage.getHeight()) {
            throw new IllegalArgumentException("Images are not the same height.");
        }

        ImageRaster normalRaster = ImageRaster.create(normalImage);
        ImageRaster parallaxRaster = ImageRaster.create(parallaxImage);

        int width = normalImage.getWidth();
        int height = normalImage.getHeight();

        ByteBuffer buffer = BufferUtils.createByteBuffer(width * height * 4);
        Image result = new Image(Image.Format.RGB8, width, height, buffer, ColorSpace.Linear);

        ImageRaster raster = ImageRaster.create(result);

        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {

                ColorRGBA normalColor = normalRaster.getPixel(x, y);
                ColorRGBA parallaxColor = parallaxRaster.getPixel(x, y);

                ColorRGBA pixelColor = new ColorRGBA();

                pixelColor.r = normalColor.r;
                pixelColor.g = normalColor.g;
                pixelColor.b = normalColor.b;
                pixelColor.a = parallaxColor.r; // greyscale, so just grab any of the RGB's

                raster.setPixel(x, y, pixelColor);
            }
        }

        return result;

    }

    public static Texture packNormalParallax(Texture normalTex, Texture parallaxTex) {
        Image image = packNormalParallax(normalTex.getImage(), parallaxTex.getImage());
        return new Texture2D(image);
    }

    public static Texture packNormalParallax(AssetManager assetManager, String normalPath, String parallaxPath) {
        return packNormalParallax(assetManager.loadTexture(normalPath), assetManager.loadTexture(parallaxPath));
    }

}