xxxxxxxxxx
235
// TODO:
// 1. Pass resolution of image down to compute shader for indexing
// 2. Write descriptor code (figure out casting bool to int)
// 3. Write descriptor matching code (maybe use another compute shader if too slow?)
// 4. Write line drawing code (features id to coordinate )
let cam;
let camg1;
const MAX_FEATURES = 5000;
const MAX_IMAGE_SIZE = 1000000
const N_PAIRS = 128;
const S = 10; // BRIEF diameter
let n_features;
let keypoints = new Int32Array(MAX_FEATURES);
let descriptor = new Int32Array(MAX_FEATURES*4);
let pairs = new Int32Array(N_PAIRS*2);
let img = new Float32Array(MAX_IMAGE_SIZE);
let device, commandBuffer, keypointBuffer, imageBuffer, resultBuffer, descriptorBuffer, descriptorStagingBuffer, pairsBuffer, pipeline, bindGroup;
function preload(){
camShader1 = loadShader('effect.vert', 'fast.frag');
}
async function runGPU() {
submitCommand();
}
async function initWebGPU() {
const adapter = await navigator.gpu?.requestAdapter();
device = await adapter?.requestDevice();
if (!device) {
fail('need a browser that supports WebGPU');
return;
}
const module = device.createShaderModule({
label: 'doubling compute module',
code: `
@group(0) @binding(0) var<storage, read_write> keypoints: array<i32>;
@group(0) @binding(1) var<storage, read_write> img: array<f32>;
@group(0) @binding(2) var<storage, read_write> pairs: array<i32>;
@group(0) @binding(3) var<storage, read_write> descriptors: array<i32>;
@compute @workgroup_size(1) fn brief(
@builtin(global_invocation_id) id: vec3<u32>
) {
let i = id.x; // keypoint feature id
keypoints[i] = keypoints[i];
img[i] = img[i];
descriptors[i] = descriptors[i];
pairs[i] = pairs[i];
if(img[keypoints[i]] > img[keypoints[i] +1]) {
descriptors[i] = 1;
} else {
descriptors[i] = 0;
}
}
`,
});
pipeline = device.createComputePipeline({
label: 'brief pipeline',
layout: 'auto',
compute: {
module,
entryPoint: 'brief',
},
});
// create a buffer on the GPU to hold our computation
// input and output
keypointBuffer = device.createBuffer({
label: 'keypoint buffer',
size: keypoints.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
});
pairsBuffer = device.createBuffer({
label: 'pairs buffer',
size: keypoints.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
});
imageBuffer = device.createBuffer({
label: 'image buffer',
size: img.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
});
descriptorBuffer = device.createBuffer({
label: 'descriptor buffer',
size: descriptor.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
});
descriptorStagingBuffer = device.createBuffer({
label: 'descriptor staging buffer',
size: descriptor.byteLength,
usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});
resultBuffer = device.createBuffer({
label: 'result buffer',
size: keypoints.byteLength,
usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});
bindGroup = device.createBindGroup({
label: 'bindGroup for work buffer',
layout: pipeline.getBindGroupLayout(0),
entries: [
{ binding: 0, resource: { buffer: keypointBuffer } },
{ binding: 1, resource: { buffer: imageBuffer } },
{ binding: 2, resource: { buffer: pairsBuffer } },
{ binding: 3, resource: { buffer: descriptorBuffer } },
],
});
}
async function submitCommand() {
// Encode commands to do the computation
const encoder = device.createCommandEncoder({
label: 'doubling encoder',
});
const pass = encoder.beginComputePass({
label: 'doubling compute pass',
});
pass.setPipeline(pipeline);
pass.setBindGroup(0, bindGroup);
// We will iterature through the features
pass.dispatchWorkgroups(n_features);
pass.end();
// Encode a command to copy the results to a mappable buffer.
encoder.copyBufferToBuffer(keypointBuffer, 0, resultBuffer, 0, resultBuffer.size);
encoder.copyBufferToBuffer(descriptorBuffer, 0, descriptorStagingBuffer, 0, descriptorStagingBuffer.size);
// Finish encoding and submit the commands
commandBuffer = encoder.finish();
device.queue.writeBuffer(keypointBuffer, 0, keypoints);
device.queue.writeBuffer(imageBuffer, 0, img);
device.queue.writeBuffer(pairsBuffer, 0, pairs);
device.queue.submit([commandBuffer]);
// Read the results
await resultBuffer.mapAsync(GPUMapMode.READ);
const result = new Int32Array(resultBuffer.getMappedRange().slice());
resultBuffer.unmap();
await descriptorStagingBuffer.mapAsync(GPUMapMode.READ);
const descriptorResults = new Int32Array(descriptorStagingBuffer.getMappedRange().slice());
descriptorStagingBuffer.unmap();
}
function featureId_to_coordinate(keypoint_id) {
let id = keypoints[keypoint_id]
return {x: id % width, y: id / height};
}
function setup() {
initWebGPU();
// Generate random pairs
for(let i = 0; i < 2*N_PAIRS; i += 1) {
pairs[i] = floor(random() * S);
}
createCanvas(windowWidth, windowHeight);
camg1 = createGraphics(windowWidth, windowHeight, WEBGL);
camg1.pixelDensity(1);
camg1.noStroke();
cam = createCapture(VIDEO);
cam.size(windowWidth, windowHeight);
cam.hide();
}
function draw() {
camg1.shader(camShader1);
camShader1.setUniform('tex0', cam);
camShader1.setUniform('iResolutionX', windowWidth);
camShader1.setUniform('iResolutionY', windowHeight);
camg1.rect(0, 0, width, height);
image(camg1, 0, 0);
getKeypointsFromTexture(camg1);
if(frameCount % 30 == 0) runGPU();
stroke(0, 255, 0)
noFill();
for(let i = 0; i < 10; ++i) {
let p1 = featureId_to_coordinate(0);
let p2 = featureId_to_coordinate(floor(random()*n_features));
line(p1.x, p1.y, p2.x, p2.y);
circle(p1.x, p1.y, 5);
circle(p2.x, p2.y, 5);
}
}
function mouseClicked() {
runGPU();
}
function getKeypointsFromTexture(camg) {
camg.loadPixels();
n_features = 0;
for(let i = 0; i < camg.pixels.length; i += 4) {
img[i / 4] = (camg.pixels[i] + camg.pixels[i+1] + camg.pixels[i+2]) / 3.0;
if(camg.pixels[i] == 255) {
keypoints[n_features] = i / 4;
n_features += 1;
if(n_features == MAX_FEATURES) break;
}
}
}
function windowResized(){
resizeCanvas(windowWidth, windowHeight);
}