Skip to content

Commit

Permalink
use a rendering pipeline instead of compiling program for every mask
Browse files Browse the repository at this point in the history
  • Loading branch information
@sashankaryal
sashankaryal committed Jan 3, 2025
1 parent 4091a48 commit fb23c02
Show file tree
Hide file tree
Showing 2 changed files with 172 additions and 135 deletions.
297 changes: 167 additions & 130 deletions app/packages/looker/src/worker/painter-gpu.ts
View file
Original file line number Diff line number Diff line change
@@ -1,193 +1,230 @@
import { OverlayMask } from "../numpy";

const canvas = new OffscreenCanvas(1, 1);
const gl = canvas.getContext("webgl2");
// todo: have a fallback strategy?
if (!gl) throw new Error("webgl2 not supported in this browser");
// note: for POC only
export function getColorForCategoryTesting(cat: number): number {
if (cat === 1) return 0xff_00_00_ff;
if (cat === 2) return 0x00_ff_00_ff;
if (cat === 3) return 0x00_00_ff_ff;
return 0x00_00_00_00; // transparent
}

const createShader = (
gl: WebGL2RenderingContext,
source: string,
type: number
): WebGLShader => {
const shader = gl.createShader(type)!;
gl.shaderSource(shader, source);
gl.compileShader(shader);

if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
const msg = gl.getShaderInfoLog(shader);
gl.deleteShader(shader);
throw new Error("shader compile error: " + msg);
}
return shader;
};

export const renderSegmentationMask = (
mask: OverlayMask,
const createProgram = (
gl: WebGL2RenderingContext,
vsSource: string,
fsSource: string
): WebGLProgram => {
const vs = createShader(gl, vsSource, gl.VERTEX_SHADER);
const fs = createShader(gl, fsSource, gl.FRAGMENT_SHADER);
const program = gl.createProgram()!;

gl.attachShader(program, vs);
gl.attachShader(program, fs);
gl.linkProgram(program);

if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
const msg = gl.getProgramInfoLog(program);
gl.deleteProgram(program);
throw new Error("program link error: " + msg);
}

// can detach and delete now that it’s linked
gl.detachShader(program, vs);
gl.detachShader(program, fs);
gl.deleteShader(vs);
gl.deleteShader(fs);

return program;
};

/** create a 256x1 look up texture for category -> RGBA */
const createLUTTexture = (
gl: WebGL2RenderingContext,
getColorForCategory: (cat: number) => number
) => {
const [height, width] = mask.shape;
canvas.width = width;
canvas.height = height;
): WebGLTexture => {
const lutSize = 256;
const lutData = new Uint8Array(lutSize * 4);

const maskData = new Uint8Array(mask.buffer);
for (let i = 0; i < lutSize; i++) {
const rgba32 = getColorForCategory(i);
lutData[i * 4 + 0] = (rgba32 >>> 24) & 0xff; // R
lutData[i * 4 + 1] = (rgba32 >>> 16) & 0xff; // G
lutData[i * 4 + 2] = (rgba32 >>> 8) & 0xff; // B
lutData[i * 4 + 3] = (rgba32 >>> 0) & 0xff; // A
}

// vertex shader:
// render full-screen quad. we'll compute textre coords from aPosition
const vsSource = `#version 300 es
layout(location = 0) in vec2 aPosition;
out vec2 vTexCoord;
void main() {
// Map aPosition -1..1 to 0..1 texture coords
vTexCoord = (aPosition + 1.0) * 0.5;
gl_Position = vec4(aPosition, 0.0, 1.0);
}
`;

// fragment shader:
// - sample the mask texture to get an 8-bit ID in [0..1].
// - convert that to [0..255], then sample the LUT texture.
const fsSource = `#version 300 es
precision mediump float;
in vec2 vTexCoord;
out vec4 outColor;
// R8
uniform sampler2D uMaskTex;
// RGBA8 (256 x 1)
uniform sampler2D uLutTex;
void main() {
float catId = texture(uMaskTex, vTexCoord).r; // catId in [0..1]
float index = catId * 255.0; // [0..255]
// sample the LUT at x = (index + 0.5)/256.0, y=0.5 (a single row)
outColor = texture(uLutTex, vec2((index + 0.5)/256.0, 0.5));
}
`;

function createShader(
glCtx: WebGL2RenderingContext,
source: string,
type: number
) {
const shader = glCtx.createShader(type)!;
glCtx.shaderSource(shader, source);
glCtx.compileShader(shader);

if (!glCtx.getShaderParameter(shader, glCtx.COMPILE_STATUS)) {
const msg = glCtx.getShaderInfoLog(shader);
glCtx.deleteShader(shader);
throw new Error("Shader compile error: " + msg);
}
const lutTex = gl.createTexture()!;
gl.bindTexture(gl.TEXTURE_2D, lutTex);
gl.texImage2D(
gl.TEXTURE_2D,
0,
gl.RGBA8,
lutSize,
1,
0,
gl.RGBA,
gl.UNSIGNED_BYTE,
lutData
);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);

return lutTex;
};

// pipeline to reuse for every render
export interface SegmentationPipeline {
canvas: OffscreenCanvas;
gl: WebGL2RenderingContext;
program: WebGLProgram;
vao: WebGLVertexArrayObject;
lutTex: WebGLTexture;
}

return shader;
export function initSegmentationPipeline(
getColorForCategory: (cat: number) => number
): SegmentationPipeline {
const canvas = new OffscreenCanvas(1, 1);

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/app/src/utils.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/overlays/heatmap.ts:4:31

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/looker/src/overlays/index.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/overlays/heatmap.ts:4:31

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/looker/src/worker/painter.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/worker/painter.test.ts:2:31

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/playback/src/lib/utils.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/overlays/heatmap.ts:4:31

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/state/src/recoil/aggregations.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/overlays/heatmap.ts:4:31

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/state/src/recoil/dynamicGroups.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/overlays/heatmap.ts:4:31

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/state/src/recoil/filters.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/overlays/heatmap.ts:4:31

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/state/src/recoil/groups.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/overlays/heatmap.ts:4:31

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/state/src/recoil/options.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/overlays/heatmap.ts:4:31

Check failure on line 105 in app/packages/looker/src/worker/painter-gpu.ts

View workflow job for this annotation

GitHub Actions / test / test-app

packages/state/src/recoil/queryPerformance.test.ts 

ReferenceError: OffscreenCanvas is not defined
 ❯ Module.initSegmentationPipeline packages/looker/src/worker/painter-gpu.ts:105:18
 ❯ packages/looker/src/worker/painter.ts:24:18
 ❯ packages/looker/src/overlays/heatmap.ts:4:31
const gl = canvas.getContext("webgl2");

// todo: fallback strategy...?
if (!gl) {
throw new Error("WebGL2 not supported in this browser/environment");
}

function createProgram(
glCtx: WebGL2RenderingContext,
vs: string,
fs: string
) {
const vShader = createShader(glCtx, vs, glCtx.VERTEX_SHADER);
const fShader = createShader(glCtx, fs, glCtx.FRAGMENT_SHADER);
const prog = glCtx.createProgram()!;
glCtx.attachShader(prog, vShader);
glCtx.attachShader(prog, fShader);
glCtx.linkProgram(prog);

if (!glCtx.getProgramParameter(prog, glCtx.LINK_STATUS)) {
const msg = glCtx.getProgramInfoLog(prog);
glCtx.deleteProgram(prog);
throw new Error("Program link error: " + msg);
// vertex shader
const vsSource = `#version 300 es
layout(location = 0) in vec2 aPosition;
out vec2 vTexCoord;
void main() {
// map aPosition from [-1..1] to [0..1] texture coords
vTexCoord = (aPosition + 1.0) * 0.5;
gl_Position = vec4(aPosition, 0.0, 1.0);
}
`;

return prog;
}
// fragment shader
const fsSource = `#version 300 es
precision mediump float;
in vec2 vTexCoord;
out vec4 outColor;
uniform sampler2D uMaskTex;
uniform sampler2D uLutTex;
void main() {
// mask.r is in [0..1], multiply by 255 to get category
float catId = texture(uMaskTex, vTexCoord).r;
float index = catId * 255.0;
// sample LUT at (index + 0.5)/256.0, single row at y=0.5
outColor = texture(uLutTex, vec2((index + 0.5)/256.0, 0.5));
}
`;

// compile + link program once
const program = createProgram(gl, vsSource, fsSource);

// create full screen quad
const vao = gl.createVertexArray();
// VAO for the fullscreen quad
const vao = gl.createVertexArray()!;
gl.bindVertexArray(vao);

// fullscreen quad positions
const quadVerts = new Float32Array([-1, -1, 1, -1, -1, 1, 1, 1]);
const vbo = gl.createBuffer();
const vbo = gl.createBuffer()!;
gl.bindBuffer(gl.ARRAY_BUFFER, vbo);
gl.bufferData(gl.ARRAY_BUFFER, quadVerts, gl.STATIC_DRAW);

// enable the attribute location 0 for aPosition
gl.enableVertexAttribArray(0);
gl.vertexAttribPointer(0, 2, gl.FLOAT, false, 0, 0);

// unbind
gl.bindVertexArray(null);
gl.bindBuffer(gl.ARRAY_BUFFER, null);

// create a LUT texture for category->color
const lutTex = createLUTTexture(gl, getColorForCategory);

return { canvas, gl, program, vao, lutTex };
}

// mask as r8 texture (todo: different for rgb masks)
const maskTex = gl.createTexture();
export const renderSegmentationMask = (
pipeline: SegmentationPipeline,
mask: OverlayMask
): Uint8Array => {
const { canvas, gl, program, vao, lutTex } = pipeline;

const [height, width] = mask.shape;
canvas.width = width;
canvas.height = height;

// upload the mask data as a R8 texture
const maskData = new Uint8Array(mask.buffer);
const maskTex = gl.createTexture()!;
gl.activeTexture(gl.TEXTURE0);
gl.bindTexture(gl.TEXTURE_2D, maskTex);
gl.texImage2D(
gl.TEXTURE_2D,
0,
// internl format
gl.R8,
width,
height,
0,
// src format
gl.RED,
gl.UNSIGNED_BYTE,
maskData
);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);

// build a look up table (lut) texture (256 x 1) for category-color
// assume maskData fits in 0..255
const lutSize = 256;
const lutData = new Uint8Array(lutSize * 4);
for (let i = 0; i < lutSize; i++) {
const rgba32 = getColorForCategory(i);
// red (bits 24-31)
lutData[i * 4 + 0] = (rgba32 >>> 24) & 0xff;
// green (bits 16-23)
lutData[i * 4 + 1] = (rgba32 >>> 16) & 0xff;
// blue (bits 8-15)
lutData[i * 4 + 2] = (rgba32 >>> 8) & 0xff;
// alpha (bits 0-7)
lutData[i * 4 + 3] = (rgba32 >>> 0) & 0xff;
}

const lutTex = gl.createTexture();
gl.activeTexture(gl.TEXTURE1);
gl.bindTexture(gl.TEXTURE_2D, lutTex);
gl.texImage2D(
gl.TEXTURE_2D,
0,
gl.RGBA8,
lutSize,
1,
0,
gl.RGBA,
gl.UNSIGNED_BYTE,
lutData
);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);

// draw
gl.viewport(0, 0, width, height);
gl.useProgram(program);

gl.activeTexture(gl.TEXTURE1);
gl.bindTexture(gl.TEXTURE_2D, lutTex);

// set uniforms (maskTex -> unit 0, lutTex -> unit 1)
const maskTexLoc = gl.getUniformLocation(program, "uMaskTex");
const lutTexLoc = gl.getUniformLocation(program, "uLutTex");
// texture unit 0
gl.uniform1i(maskTexLoc, 0);
// texture unit 1
gl.uniform1i(lutTexLoc, 1);

// draw fullscreen quad
gl.bindVertexArray(vao);
gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);

// read back pixels
// read back the painted pixels
const paintedPixels = new Uint8Array(width * height * 4);
gl.readPixels(0, 0, width, height, gl.RGBA, gl.UNSIGNED_BYTE, paintedPixels);

// cleanup
gl.bindVertexArray(null);
gl.useProgram(null);
gl.deleteTexture(maskTex);

return paintedPixels;
};

// note: for POC only
export const getColorForCategoryTESTING = (cat: number) => {
// example: 1=red, 2=green, 3=blue, else transparent
if (cat === 1) return 0xff0000ff;
if (cat === 2) return 0x00ff00ff;
if (cat === 3) return 0x0000ffff;
return 0x00000000;
};
10 changes: 5 additions & 5 deletions app/packages/looker/src/worker/painter.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,13 @@ import {
RgbMaskTargets,
} from "../state";
import {
getColorForCategoryTESTING,
getColorForCategoryTesting,
initSegmentationPipeline,
renderSegmentationMask,
} from "./painter-gpu";

const pipeline = initSegmentationPipeline(getColorForCategoryTesting);

export const PainterFactory = (requestColor) => ({
Detection: async (
field,
Expand Down Expand Up @@ -323,10 +326,7 @@ export const PainterFactory = (requestColor) => ({
// discard the buffer values of other channels
maskData.buffer = maskData.buffer.slice(0, overlay.length);
} else {
const gpuResult = renderSegmentationMask(
maskData,
getColorForCategoryTESTING
);
const gpuResult = renderSegmentationMask(pipeline, maskData);
overlay.set(new Uint32Array(gpuResult.buffer));

// const cache = {};
Expand Down

0 comments on commit fb23c02

Please sign in to comment.