[WIP] Image Processing

webgl.ts

type FilterChain = {
    func: Function
    args: any[]
}

type FrameBuffer = {
    fbo: WebGLFramebuffer | null
    texture: WebGLTexture | null
}

const SHADER = {
    FRAGMENT_IDENTITY: [
        'precision highp float;',
        'varying vec2 vUv;',
        'uniform sampler2D texture;',
        'void main(void) {',
            'gl_FragColor = texture2D(texture, vUv);',
        '}',
    ].join('\n'),
    VERTEX_IDENTITY: [
        'precision highp float;',
        'attribute vec2 pos;',
        'attribute vec2 uv;',
        'varying vec2 vUv;',
        'uniform float flipY;',

        'void main(void) {',
            'vUv = uv;',
            'gl_Position = vec4(pos.x, pos.y*flipY, 0.0, 1.);',
        '}'
    ].join('\n')
}

const DRAW = { INTERMEDIATE: 1 }

class WebGLProgram {
    gl: WebGLRenderingContext
    vertexSource: string
    fragmentSource: string
    uniform: any = {}
    attribute: any = {}

    constructor(gl: WebGLRenderingContext, vertexSource: string, fragmentSource: string) {
        this.gl = gl
        this.vertexSource = vertexSource
        this.fragmentSource = fragmentSource

        this.init()        
    }

    init() {
        const _vsh = this._compile(this.gl, this.vertexSource, this.gl.VERTEX_SHADER)
        const _fsh = this._compile(this.gl, this.fragmentSource, this.gl.FRAGMENT_SHADER)

        const id = this.gl.createProgram()
        if (!id) {
            throw Error("Couldn't create WebGL program")
        }
        if (!_vsh || !_fsh) {
            throw Error("Couldn't compile WebGLShader")
        }

        this.gl.attachShader(id, _vsh)
        this.gl.attachShader(id, _fsh)
        this.gl.linkProgram(id)

        if (!this.gl.getProgramParameter(id, this.gl.LINK_STATUS)) {
            console.log(this.gl.getProgramInfoLog(id))
            throw Error("Program not linked")
        }
        this.gl.useProgram(id)

        // Collect attributes
        this._collect(this.vertexSource, 'attribute', this.attribute)
        for (const a in this.attribute) {
            // @ts-ignore
            this.attribute[a] = this.gl.getAttribLocation(id, a)
        }

        // collect uniforms
        this._collect(this.vertexSource, 'uniform', this.uniform)
        this._collect(this.fragmentSource, 'uniform', this.uniform)

        for (const u in this.uniform) {
            this.uniform[u] = this.gl.getUniformLocation(id, u)
        }
    }

    
    
    _collect = (source: string, prefix: string, collection: any) => {
        const r = new RegExp('\\b' + prefix + ' \\w+ (\\w+)', 'ig')
        source.replace(r, function(match: any, name: any) {
            collection[name] = 0
            return match
        })
    }
    _compile = (gl: WebGLRenderingContext, source: string, type: number): WebGLShader | null => {
        const shader = gl.createShader(type)
        if (!shader) {
            throw Error("Couldn't create WebGL shader")
        }
        gl.shaderSource(shader, source)
        gl.compileShader(shader)

        if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
            console.log(gl.getShaderInfoLog(shader))
            return null
        }
        return shader
    }
}


export default class WebGLImageFilter {
    private gl: WebGLRenderingContext | null
    private _sourceTexture: WebGLTexture | null
    private _vertexBuffer: WebGLBuffer | null
    private _currentProgram: WebGLProgram | any | null
    private _drawCount: number
    private _lastInChain: boolean
    private _currentFramebufferIndex: number
    private _tempFramebuffers: FrameBuffer[]
    private _filterChain: FilterChain[]
    private _width: number
    private _height: number
    private _canvas: HTMLCanvasElement
    _filter = {
        colorMatrix: (matrix: Iterable<number>) => {
            if (!this.gl) {
                throw Error("Couldn't get WebGL context")
            }
            // Create a Float32 Array and normalize the offset component to 0-1
            var m = new Float32Array(matrix)
            m[4] /= 255
            m[9] /= 255
            m[14] /= 255
            m[19] /= 255

            // Can we ignore the alpha value? Makes things a bit faster.
            var shader = (1 === m[18] && 0 === m[3] && 0 === m[8] && 0 === m[13] && 0 === m[15] && 0 === m[16] && 0 === m[17] && 0 === m[19])
                ? this._filter.colorMatrixSHADERWITHOUT_ALPHA
                : this._filter.colorMatrixSHADERWITH_ALPHA;
            
            var program = this._compileShader(shader)
            this.gl.uniform1fv(program.uniform.m, m)
            this._draw()
        },
        colorMatrixSHADER: {},
        colorMatrixSHADERWITH_ALPHA: [
            'precision highp float;',
            'varying vec2 vUv;',
            'uniform sampler2D texture;',
            'uniform float m[20];',

            'void main(void) {',
                'vec4 c = texture2D(texture, vUv);',
                'gl_FragColor.r = m[0] * c.r + m[1] * c.g + m[2] * c.b + m[3] * c.a + m[4];',
                'gl_FragColor.g = m[5] * c.r + m[6] * c.g + m[7] * c.b + m[8] * c.a + m[9];',
                'gl_FragColor.b = m[10] * c.r + m[11] * c.g + m[12] * c.b + m[13] * c.a + m[14];',
                'gl_FragColor.a = m[15] * c.r + m[16] * c.g + m[17] * c.b + m[18] * c.a + m[19];',
            '}',
        ].join('\n'),
        colorMatrixSHADERWITHOUT_ALPHA: [
            'precision highp float;',
            'varying vec2 vUv;',
            'uniform sampler2D texture;',
            'uniform float m[20];',

            'void main(void) {',
            'vec4 c = texture2D(texture, vUv);',
                'gl_FragColor.r = m[0] * c.r + m[1] * c.g + m[2] * c.b + m[4];',
                'gl_FragColor.g = m[5] * c.r + m[6] * c.g + m[7] * c.b + m[9];',
                'gl_FragColor.b = m[10] * c.r + m[11] * c.g + m[12] * c.b + m[14];',
                'gl_FragColor.a = c.a;',
            '}',
        ].join('\n'),
        brightness: (brightness: number) => {
            const b = (brightness || 0) + 1
            this._filter.colorMatrix([
                b, 0, 0, 0, 0,
                0, b, 0, 0, 0,
                0, 0, b, 0, 0,
                0, 0, 0, 1, 0
            ])
        },
        saturation: (amount: number) => {
            const x = (amount || 0) * 2/3 + 1
            const y = ((x-1) *-0.5)
            this._filter.colorMatrix([
                x, y, y, 0, 0,
                y, x, y, 0, 0,
                y, y, x, 0, 0,
                0, 0, 0, 1, 0
            ])
        },
        desaturate: () => {
            this._filter.saturation(-1)
        },
        contrast: (amount: number) => {
            const v = (amount || 0) + 1
            const o = -128 * (v - 1)

            this._filter.colorMatrix([
                v, 0, 0, 0, o,
                0, v, 0, 0, o,
                0, 0, v, 0, o,
                0, 0, 0, 1, 0 
            ])
        },
        negative: () => {
            this._filter.contrast(-2)
        },
        convolution: (matrix: Iterable<number>) => {
            if (!this.gl) {
                throw Error("Couldn't get WebGL context")
            }

            var m = new Float32Array(matrix)
            var pixelSizeX = 1 / this._width
            var pixelSizeY = 1 / this._height

            var program = this._compileShader(this._filter.convolutionShader)
            this.gl.uniform1fv(program.uniform.m, m);
            this.gl.uniform2f(program.uniform.px, pixelSizeX, pixelSizeY);
            this._draw()
        },
        convolutionShader: [
            'precision highp float;',
            'varying vec2 vUv;',
            'uniform sampler2D texture;',
            'uniform vec2 px;',
            'uniform float m[9];',

            'void main(void) {',
            'vec4 c11 = texture2D(texture, vUv - px);', // top left
            'vec4 c12 = texture2D(texture, vec2(vUv.x, vUv.y - px.y));', // top center
            'vec4 c13 = texture2D(texture, vec2(vUv.x + px.x, vUv.y - px.y));', // top right

            'vec4 c21 = texture2D(texture, vec2(vUv.x - px.x, vUv.y) );', // mid left
            'vec4 c22 = texture2D(texture, vUv);', // mid center
            'vec4 c23 = texture2D(texture, vec2(vUv.x + px.x, vUv.y) );', // mid right

            'vec4 c31 = texture2D(texture, vec2(vUv.x - px.x, vUv.y + px.y) );', // bottom left
            'vec4 c32 = texture2D(texture, vec2(vUv.x, vUv.y + px.y) );', // bottom center
            'vec4 c33 = texture2D(texture, vUv + px );', // bottom right

            'gl_FragColor = ',
            'c11 * m[0] + c12 * m[1] + c22 * m[2] +',
            'c21 * m[3] + c22 * m[4] + c23 * m[5] +',
            'c31 * m[6] + c32 * m[7] + c33 * m[8];',
            'gl_FragColor.a = c22.a;',
            '}',
        ].join('\n')
    }

    constructor() {
        this.gl = null
        this._drawCount = 0
        this._sourceTexture = null
        this._lastInChain = false
        this._currentFramebufferIndex = -1
        this._tempFramebuffers = []
        this._filterChain = []
        this._width = -1
        this._height = -1
        this._vertexBuffer = null
        this._currentProgram = null
        this._canvas = document.createElement('canvas')

        this.gl = this._canvas.getContext('webgl')
        
        if (!this.gl) {
            throw Error("Couldn't get WebGL context")
        }
    }

    public addFilter(name: any, ...args: any[]) {
        var arg = Array.prototype.slice.call(args, 1)
        // @ts-ignore
        var filter = this._filter[name]

        this._filterChain.push({ func: filter, args: arg})
    }

    public reset() {
        this._filterChain = []
    }

    public apply(image: HTMLImageElement) {
        if (!this.gl) return

        console.log(typeof image, image)

        this._resize(image.width, image.height)
        this._drawCount = 0

        // create the texture for the input image

        this._sourceTexture = this.gl.createTexture()
        this.gl.bindTexture(this.gl.TEXTURE_2D, this._sourceTexture)
        this.gl.texParameteri(this.gl.TEXTURE_2D, this.gl.TEXTURE_WRAP_S, this.gl.CLAMP_TO_EDGE)
        this.gl.texParameteri(this.gl.TEXTURE_2D, this.gl.TEXTURE_WRAP_T, this.gl.CLAMP_TO_EDGE)
        this.gl.texParameteri(this.gl.TEXTURE_2D, this.gl.TEXTURE_MIN_FILTER, this.gl.NEAREST)
        this.gl.texParameteri(this.gl.TEXTURE_2D, this.gl.TEXTURE_MAG_FILTER, this.gl.NEAREST)
        this.gl.texImage2D(this.gl.TEXTURE_2D, 0, this.gl.RGBA, this.gl.RGBA, this.gl.UNSIGNED_BYTE, image)        

        // No filters? Just draw
        if (this._filterChain.length === 0) {
            // useless assign
            this._compileShader(SHADER.FRAGMENT_IDENTITY)
            this._draw()
            return this._canvas
        }

        for (let i = 0; i < this._filterChain.length; i++) {
            this._lastInChain = (i === this._filterChain.length-1)
            const f = this._filterChain[i]

            f.func.apply(this, f.args || [])
        }

        return this._canvas
    }

    private _resize(width: number, height: number) {
        if (!this.gl) return

        // Same width/height? Nothing to do here
        if (width === this._width && height === this._height) return

        this._canvas.width = this._width = width
        this._canvas.height = this._height = height

        // create the context if we dont have it yet
        if (!this._vertexBuffer) {
            // create the vertex buffer for the two triangles [x, y, u, v] * 6
            const vertices: ArrayBuffer = new Float32Array([
                -1, -1, 0, 1, 1, -1, 1, 1, -1, 1, 0, 0,
                -1, 1, 0, 0, 1, -1, 1, 1, 1, 1, 1, 0
            ])
            this._vertexBuffer = this.gl.createBuffer()
            this.gl.bindBuffer(this.gl.ARRAY_BUFFER, this._vertexBuffer)
            this.gl.bufferData(this.gl.ARRAY_BUFFER, vertices, this.gl.STATIC_DRAW)

            // Not sure if this is a good idea; at least it makes texture loading
            // in Ejecta instant.
            this.gl.pixelStorei(this.gl.UNPACK_PREMULTIPLY_ALPHA_WEBGL, true)
        }

        this.gl.viewport(0, 0, this._width, this._height)

        // Delete old temp framebuffers
        this._tempFramebuffers = []
    }

    private _getTempFrameBuffer(index: number) {
        this._tempFramebuffers[index] = this._tempFramebuffers[index] || this._createFramebufferTexture(this._width, this._height)
    
        return this._tempFramebuffers[index]
    }

    private _createFramebufferTexture(width: number, height: number): FrameBuffer {
        if (!this.gl) return { fbo: null, texture: null }

        const fbo: WebGLFramebuffer | null = this.gl.createFramebuffer()
        this.gl.bindFramebuffer(this.gl.FRAMEBUFFER, fbo)

        const renderbuffer: WebGLRenderbuffer | null = this.gl.createRenderbuffer()
        this.gl.bindRenderbuffer(this.gl.RENDERBUFFER, renderbuffer)

        const texture: WebGLTexture | null = this.gl.createTexture()
        this.gl.bindTexture(this.gl.TEXTURE_2D, texture)
        this.gl.texImage2D(this.gl.TEXTURE_2D, 0, this.gl.RGBA, width, height, 0, this.gl.RGBA, this.gl.UNSIGNED_BYTE, null)

        this.gl.texParameteri(this.gl.TEXTURE_2D, this.gl.TEXTURE_MAG_FILTER, this.gl.LINEAR)
        this.gl.texParameteri(this.gl.TEXTURE_2D, this.gl.TEXTURE_MIN_FILTER, this.gl.LINEAR)
        this.gl.texParameteri(this.gl.TEXTURE_2D, this.gl.TEXTURE_WRAP_S, this.gl.CLAMP_TO_EDGE)
        this.gl.texParameteri(this.gl.TEXTURE_2D, this.gl.TEXTURE_WRAP_T, this.gl.CLAMP_TO_EDGE)
    
        this.gl.framebufferTexture2D(this.gl.FRAMEBUFFER, this.gl.COLOR_ATTACHMENT0, this.gl.TEXTURE_2D, texture, 0)
    
        this.gl.bindTexture(this.gl.TEXTURE_2D, null)
        this.gl.bindFramebuffer(this.gl.FRAMEBUFFER, null)

        return { fbo: fbo, texture: texture }
    }

    private _draw(flags?: any): void {
        if (!this.gl) {
            throw Error("Couldn't get WebGL context")
        } 

        let source: WebGLTexture | null = null
        let flipY: boolean = false
        let target: WebGLFramebuffer | null = null

        // Set up the source
        if (this._drawCount === 0) {
            // First draw call - u se the source texture
            source = this._sourceTexture
        } else {
            // all following draw calls use the temp buffer lsat drawn to
            source = this._getTempFrameBuffer(this._currentFramebufferIndex).texture
        }
        this._drawCount++

        // Set up the target
        if (this._lastInChain && !(flags && DRAW.INTERMEDIATE)) {
            // last filter in our chain - draw directly to the WebGL Canvas.
            // We may also have to flip the image verticall now
            flipY = this._drawCount % 2 === 0
        } else {
            // Intermediate draw call - get a temp buffer to draw to
            this._currentFramebufferIndex = (this._currentFramebufferIndex+1) % 2
            target = this._getTempFrameBuffer(this._currentFramebufferIndex).fbo
        }

        // bind the source and target and draw the two triangles
        this.gl.bindTexture(this.gl.TEXTURE_2D, source)
        this.gl.bindFramebuffer(this.gl.FRAMEBUFFER, target)

        this.gl.uniform1f(this._currentProgram.uniform.flipY, (flipY ? -1 : 1))
        this.gl.drawArrays(this.gl.TRIANGLES, 0, 6)
    }

    private _compileShader(fragmentSource: string | any) {
        if (!this.gl) {
            throw Error("Couldn't get WebGL context")
        } 


        if (fragmentSource.__program) {
            this._currentProgram = fragmentSource.__program
            this.gl.useProgram(this._currentProgram.id)
        }

        // Compile shaders
        this._currentProgram = new WebGLProgram(this.gl, SHADER.VERTEX_IDENTITY, fragmentSource)

        const floatSize = Float32Array.BYTES_PER_ELEMENT
        const vertSize = 4 * floatSize
        this.gl.enableVertexAttribArray(this._currentProgram.attribute.pos)
        // 0 * x always zero
        // https://github.com/phoboslab/WebGLImageFilter/blob/master/webgl-image-filter.js#L239
        this.gl.vertexAttribPointer(this._currentProgram.attribute.pos, 2, this.gl.FLOAT, false, vertSize, 0 * floatSize)
        this.gl.enableVertexAttribArray(this._currentProgram.attribute.uv)
        this.gl.vertexAttribPointer(this._currentProgram.attribute.uv, 2, this.gl.FLOAT, false, vertSize, 2 * floatSize)
        
        return this._currentProgram
    }
}