taisei/scripts/gen-blur-shader.py

#!/usr/bin/env python3

from taiseilib.common import (
    add_common_args,
    run_main,
    update_text_file,
)

from math import (
    exp,
)

import argparse


def gen_comment(args):
    return f'Generated by gen-blur-shader.py. kernel size = {args.kernel_size}; sigma = {args.sigma}'


def gen_shader_func(args):
    s = ""
    kernel = gen_kernel(args.kernel_size, args.sigma)
    kernel_half_size = args.kernel_size // 2

    s += f'vec4 sample_{args.name}(sampler2D tex, vec2 uv, vec2 dir) {"{"}\n'
    s += f'    return\n'

    for i in range(-kernel_half_size, kernel_half_size+1):
        prefix = ' ' * 8
        term = ';' if i == kernel_half_size else ' +'

        uv_ofs_mul = f' * {float(abs(i))}' if abs(i) > 1 else ''
        uv_ofs = ((' + dir' if i >= 0 else ' - dir') + uv_ofs_mul) if i != 0 else ""
        s += f'{prefix}texture(tex, uv{uv_ofs}) * {kernel[kernel_half_size + i]}{term}\n'

    s += '}\n'

    return s


def gen_lib_shader(args):
    macro = f'BLUR_{args.name}_H'

    return (
        '\n'
        f'// {gen_comment(args)}\n'
        '\n'
        f'#ifndef {macro}\n'
        f'#define {macro}\n'
        '\n'
        f'{gen_shader_func(args)}'
        '\n'
        '#endif\n'
    )


def gen_fragment_shader(args):
    return (
        '#version 330 core\n'
        '\n'
        f'// {gen_comment(args)}\n'
        '\n'
        '#include "interface/standard.glslh"\n'
        f'#include "lib/blur/{args.name}.glslh"\n'
        '\n'
        'UNIFORM(1) vec2 blur_resolution;\n'
        'UNIFORM(2) vec2 blur_direction;\n'
        '\n'
        'void main(void) {\n'
        f'    fragColor = sample_{args.name}(tex, texCoord, blur_direction / blur_resolution);\n'
        '}\n'
    )


def gen_meta(args):
    return (
        '\n'
        f'# {gen_comment(args)}\n'
        '\n'
        f'objects = standardnotex.vert {args.name}.frag\n'
    )


def normpdf(x, sigma):
    return 0.39894 * exp(-0.5 * x ** 2 / sigma ** 2) / sigma


def gen_kernel(size, sigma):
    assert size & 1

    kernel = [0] * size
    halfsize = size // 2
    ksum = 0.0

    kernel[halfsize] = normpdf(0, sigma);
    ksum += kernel[halfsize];

    for i in range(1, halfsize + 1):
        k = normpdf(i, sigma)
        kernel[halfsize + i] = kernel[halfsize - i] = k
        ksum += k * 2

    for i in range(0, size):
        kernel[i] /= ksum

    return kernel


def main(args):
    parser = argparse.ArgumentParser(description='Generate a reusable gaussian blur shader.', prog=args[0])

    def kernsize(v):
        v = int(v)
        if not v & 1:
            raise ValueError('Kernel size must be an odd integer')
        return v

    parser.add_argument('kernel_size',
        help='size of the 1D convolution kernel, must be odd',
        type=kernsize,
    )

    parser.add_argument('sigma',
        help='standard deviation value for the probability density function',
        type=float,
    )

    parser.add_argument('name',
        help='name of the shader, defaults to blurX, where X is the kernel size',
        type=str,
        default=None,
        nargs='?',
    )

    add_common_args(parser)
    args = parser.parse_args()

    if args.name is None:
        args.name = f'blur{args.kernel_size}'

    shaders = args.rootdir / 'resources' / 'shader'
    blurs = shaders / 'lib' / 'blur'
    blurs.mkdir(parents=True, exist_ok=True)

    update_text_file(blurs / f'{args.name}.glslh', gen_lib_shader(args))
    update_text_file(shaders / f'{args.name}.frag.glsl', gen_fragment_shader(args))
    update_text_file(shaders / f'{args.name}.prog', gen_meta(args))


if __name__ == '__main__':
    run_main(main)