rbf_interp_cli_vectors.py

"""
Try with vector function values instead of scalars
(triplets instead of single values).
Should be a little cleaner.
"""
from argparse import ArgumentParser
import numpy as np
import xalglib

def rbf_interp(source, target, cube_size):
    """
    Takes a 2D array of RGB triplet points as both args.
    Returns an array of interpolated grid values.
    """
    data = np.hstack((source, target))

    model = xalglib.rbfcreate(3, 3)
    xalglib.rbfsetpoints(model, data.tolist())

    xalglib.rbfsetalgohierarchical(model, 5.0, 5, 0.0)
    xalglib.rbfbuildmodel(model)

    grid = np.linspace(0, 1, cube_size)
    values = []

    # Haven't figured out how to get triplets from the
    # official gridcalc3v function, so doing it manually, so it's slower.
    for b in grid:
        for g in grid:
            for r in grid:
                values.append(xalglib.rbfcalc(model, [r, g, b]))

    return values

def write_lut(grid_values, size, name):
    file = open(name, "w")

    file.write(f"LUT_3D_SIZE {size}\n\n")

    for i in range(0, size ** 3):
        file.write(
            "{:.12f} {:.12f} {:.12f}\n"
                .format(grid_values[i][0], grid_values[i][1], grid_values[i][2])
        )

    file.close()

def collect_args():
    parser = ArgumentParser(description="Generate a 3D LUT to match datasets")

    parser.add_argument("source", type=str, help="source dataset")
    parser.add_argument("target", type=str, help="target dataset")
    parser.add_argument("-n", "--name", default="output", type=str, help="LUT filename")
    parser.add_argument("-s", "--size", default=33, type=int, help="LUT size (defaults to 33)")
    
    return parser.parse_args()

def main():
    args = collect_args()

    source = np.loadtxt(args.source, delimiter=" ")
    target = np.loadtxt(args.target, delimiter=" ")

    lut = rbf_interp(source, target, args.size)
    write_lut(lut, args.size, args.name + ".cube")

main()