Source: ui/color/ColorizedImage.js

import iro from '@jaames/iro';

import { UiElement } from "../UiElement.js";
import { createProgram, createTexture } from "../webgl/Utils.js";
import { validateColorScale } from "./Utils.js";

export { ColorizedImage };

/**
 * Module for applying a different color scale to an image
 * 
 * @author rhess <robin.hess@awi.de>
 * @memberof vef.ui.color
 */
class ColorizedImage extends UiElement {

    /**
     * options = {
     *   src: "example.png"
     *   scale: {object[]} ColorScale array,
     *   maxScaleSize: 10 //default
     * }
     * 
     * @param {HTMLElement | string} target 
     * @param {object} options 
     */
    constructor(target, options) {
        super(target);

        // initialize options
        this.options = Object.assign({
            src: "",
            scale: [],
            maxScaleSize: 10
        }, options);

        if (this.options.src.length == 0) throw ("no image: options.src is empty");

        // initialize properties
        this.scale = options.scale || [];
        this.canvas_ = document.createElement("canvas");
        this.glProgram_ = null;
        this.image_ = new Image();

        this.image_.onload = () => {
            this.setSize_();
            this.initWebGlProgram_();
            this.initWebGlData_();
            this.setScale(this.options.scale);
        };

        // load image
        this.image_.src = this.options.src;

        // append canvas to container
        this.getElement().appendChild(this.canvas_);
    }

    /**
     * @param {object[]} image 
     */
    setScale(scale) {
        validateColorScale(scale);
        this.options.scale = scale;
        this.applyScale_();
    }

    /**
     * @private
     * @param {object[]} image 
     */
    applyScale_() {
        // abort if the webgl program is not initialized
        if (!this.glProgram_) return;

        const gl = this.glProgram_.gl;
        const scale = this.convertScale_();

        // set color scale uniforms
        gl.uniform1i(this.glProgram_.uniforms.u_colorScaleSize, scale.size);
        gl.uniform1fv(this.glProgram_.uniforms.u_colorScaleValues, scale.values);
        gl.uniform4fv(this.glProgram_.uniforms.u_colorScale, scale.scale);

        // render the result
        gl.drawArrays(gl.TRIANGLES, 0, 6);
    }

    /**
     * @private
     * @returns color scale formatted for passing it to the shader
     */ 
    convertScale_() {
        const scale = this.options.scale;
        const maxSize = this.options.maxScaleSize;

        const rgbaScale = []
        const values = [];

        for (let i = 0; i < scale.length; ++i) {
            const color = new iro.Color(scale[i].color).rgb;
            rgbaScale.push(color.r / 255);
            rgbaScale.push(color.g / 255);
            rgbaScale.push(color.b / 255);
            rgbaScale.push((Number.isFinite(scale.opacity) ? scale.opacity : 1.0));
            values.push(scale[i].value);
        }

        while (rgbaScale.length < (maxSize * 4)) rgbaScale.push(0.0);
        while (values.length < maxSize) values.push(0.0);

        return {
            size: scale.length,
            scale: new Float32Array(rgbaScale),
            values: new Float32Array(values)
        }
    }

    /**
     * @private
     */
    initWebGlProgram_() {
        const gl = this.canvas_.getContext('webgl')

        // create vertex shader
        const vShader = `
            attribute vec4 a_position;
            attribute vec2 a_textureCoordinates;
            
            varying vec2 v_textureCoordinates;
            
            void main() {
                gl_Position = a_position;
                v_textureCoordinates = a_textureCoordinates;
            }
        `;

        // create fragment shader
        const fShader = `
            precision highp float;

            uniform int u_colorScaleSize;
            uniform float u_colorScaleValues[${this.options.maxScaleSize}];
            uniform vec4 u_colorScale[${this.options.maxScaleSize}];

            uniform sampler2D u_texture;

            varying vec2 v_textureCoordinates;

            vec4 getColor(float intensity) {
            
                float minValue = u_colorScaleValues[0];
                vec4 minColor = u_colorScale[0];
                float maxValue = u_colorScaleValues[0];
                vec4 maxColor = u_colorScale[0];

                for (int i = 0; i < ${this.options.maxScaleSize}; ++i) {
                    
                    // only run loop for array size (loop header must be hardcoded)
                    if (i < u_colorScaleSize) {
                        if (u_colorScaleValues[i] <= intensity) {
                            minValue = u_colorScaleValues[i];
                            minColor = u_colorScale[i];
                        }

                        if (u_colorScaleValues[i] >= intensity) {
                            maxValue = u_colorScaleValues[i];
                            maxColor = u_colorScale[i];
                        }

                        // finish condition -> size cannot be used in loop header
                        if (i == (u_colorScaleSize - 1)) {

                            // Has to be done in loop, because indices cannot be accessed using vaiables.
                            // Accessing using loop indices works, because the loop is unwrapped before execution

                            if (intensity < u_colorScaleValues[0]) {
                                return u_colorScale[0];
                            } else if (intensity > u_colorScaleValues[i]) {
                                return u_colorScale[i];
                            } else {
                                float ratio = (intensity - minValue) / (maxValue - minValue);
                                return minColor + ((maxColor - minColor) * ratio);
                            }
                        }
                    }
                }
            
                // fallback: use original color
                return texture2D(u_texture, v_textureCoordinates);
            }

            void main() {
                vec4 color = texture2D(u_texture, v_textureCoordinates);
                float intensity = (color.r + color.g + color.b) / 3.0;

                gl_FragColor = getColor(intensity);
            }

        `;

        // initialize the webgl shader program
        this.glProgram_ = createProgram(gl, vShader, fShader, ["a_position", "a_textureCoordinates"], ["u_texture", "u_colorScaleSize", "u_colorScaleValues", "u_colorScale"]);
        gl.useProgram(this.glProgram_.program);
    }

    /**
     * @private
     */
    initWebGlData_() {
        const program = this.glProgram_;
        const gl = program.gl;

        // texture coordinates
        const textureCoordinates = new Float32Array([
            0.0, 1.0,
            0.0, 0.0,
            1.0, 0.0,
            1.0, 0.0,
            1.0, 1.0,
            0.0, 1.0
        ]);
        this.assignAttribute_(gl, textureCoordinates, program.attributes.a_textureCoordinates);

        // create vertices for a quad that fills the whole canvas
        const vertices = new Float32Array([
            -1, -1,
            -1, 1,
            1, 1,
            1, 1,
            1, -1,
            -1, -1,
        ]);
        this.assignAttribute_(gl, vertices, program.attributes.a_position);

        // texture uniform
        const texture = createTexture(gl, this.image_);
        gl.uniform1i(program.uniforms.u_texture, 0);
    }

    /**
     * Helper method for assigning attribute values to a buffer and to the
     * shader attribute
     * 
     * @private
     * @param {WebGLContext} gl 
     * @param {Float32Array} data 
     * @param {AttributePointer} attributePointer 
     */
    assignAttribute_(gl, data, attributePointer) {
        const buffer = gl.createBuffer();
        gl.bindBuffer(gl.ARRAY_BUFFER, buffer);
        gl.bufferData(gl.ARRAY_BUFFER, data, gl.STATIC_DRAW);
        gl.enableVertexAttribArray(attributePointer);
        gl.vertexAttribPointer(attributePointer, 2, gl.FLOAT, false, 0, 0);
    }

    /**
     * @private
     */
    setSize_() {
        const element = this.getElement();
        this.canvas_.width = this.image_.width;
        this.canvas_.height = this.image_.height;
        element.style.width = this.image_.width + "px";
        element.style.height = this.image_.height + "px";
    }
}