Applying a LUT texture with skia

lut-image.js

import {
    Skia,
    TileMode,
    FilterMode,
    MipmapMode,
} from '@shopify/react-native-skia'

export function renderLUTImage({
    baseImage,
    lutImage,
    lutShader,
    width,
    height,
}) {
    const surface = Skia.Surface.MakeOffscreen(width, height)
    if (!surface) return null

    const scaleMatrix = Skia.Matrix()
    scaleMatrix.scale(width / baseImage.width(), height / baseImage.height())

    const baseShader = baseImage.makeShaderOptions(
        TileMode.Clamp,
        TileMode.Clamp,
        FilterMode.Linear,
        MipmapMode.None,
        scaleMatrix
    )

    const lutShaderTex = lutImage.makeShaderOptions(
        TileMode.Clamp,
        TileMode.Clamp,
        FilterMode.Linear,
        MipmapMode.None
    )

    const shader = lutShader.makeShaderWithChildren(
        [],
        [baseShader, lutShaderTex]
    )

    const paint = Skia.Paint()
    paint.setShader(shader)

    const canvas = surface.getCanvas()
    canvas.drawPaint(paint)

    const snapshot = surface.makeImageSnapshot()

    const gpuImage = snapshot.makeNonTextureImage()

    return gpuImage
}

lut-shader-string.js

export const lutShaderString = `
uniform shader image;
uniform shader luts;

vec2 lutCoord(vec3 idx) {
    float slice  = idx.b;
    float tileX  = mod(slice, 8.0);
    float tileY  = floor(slice / 8.0);
    return vec2(tileX * 64.0 + idx.r + 0.5,
                tileY * 64.0 + idx.g + 0.5);
}

vec3 tex3(vec3 idx) { return luts.eval(lutCoord(idx)).rgb; }

vec3 sampleLUT(vec3 rgb) {
    vec3 scaled   = clamp(rgb * 63.0, 0.0, 63.0);
    vec3 floorIdx = floor(scaled);
    vec3 ceilIdx  = min(floorIdx + 1.0, 63.0);
    vec3 frac     = scaled - floorIdx;

    vec3 lerp1 = mix(
        mix(tex3(floorIdx),
            tex3(vec3(ceilIdx.x, floorIdx.y, floorIdx.z)), frac.x),
        mix(tex3(vec3(floorIdx.x, ceilIdx.y, floorIdx.z)),
            tex3(vec3(ceilIdx.x, ceilIdx.y, floorIdx.z)), frac.x),
        frac.y);

    vec3 lerp2 = mix(
        mix(tex3(vec3(floorIdx.x, floorIdx.y, ceilIdx.z)),
            tex3(vec3(ceilIdx.x, floorIdx.y, ceilIdx.z)), frac.x),
        mix(tex3(vec3(floorIdx.x, ceilIdx.y,  ceilIdx.z)),
            tex3(ceilIdx), frac.x),
        frac.y);

    return mix(lerp1, lerp2, frac.z);
}

half4 main(float2 xy) {
    vec4 color = image.eval(xy);

    vec3 lutColor = sampleLUT(rgb);
    return half4(clamp(lutColor, 0.0, 1.0), color.a);
}
`