xxxxxxxxxx
369
// TODO:
// Make two images with matches between both
// How to make this better?
// 1. Look at distribution of hamming distances
let cam;
let camg1;
let blurredCam;
let MIN_BITS_C = 20;
const MAX_FEATURES = 5000;
const MAX_IMAGE_SIZE = 1000000;
const N_PAIRS = 512;
const S = 64; // BRIEF diameter
let n_features;
let n_matchFeatures;
let keypoints = new Int32Array(MAX_FEATURES);
let matchKeypoints = new Int32Array(MAX_FEATURES);
let closestFeatures = new Uint32Array(MAX_FEATURES);
let img = new Float32Array(MAX_IMAGE_SIZE);
let matchImg = new Float32Array(MAX_IMAGE_SIZE);
let descriptorResults;
let matchDescriptorResults;
let descriptor = new Uint32Array(MAX_FEATURES*12);
let pairs = new Int32Array(N_PAIRS*4);
let resolution = new Int32Array(2);
let device, commandBuffer, keypointBuffer, imageBuffer, resultBuffer, descriptorBuffer, descriptorStagingBuffer, pairsBuffer, resolutionBuffer, pipeline, bindGroup;
function preload(){
camShader1 = loadShader('effect.vert', 'fast.frag');
}
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<u32>;
@group(0) @binding(4) var<storage, read_write> resolution: array<i32>;
@compute @workgroup_size(1) fn brief(
@builtin(global_invocation_id) id: vec3<u32>
) {
let feature_id = id.x; // keypoint feature id
let i = i32(feature_id) % resolution[0]; // column id
let j = i32(feature_id) / resolution[0]; // row id
// If I don't do this, the program crashes. Why?
keypoints[feature_id] = keypoints[feature_id];
img[feature_id] = img[feature_id];
pairs[feature_id] = pairs[feature_id];
resolution[feature_id] = resolution[feature_id];
var descriptor: u32 = 0;
var n: u32 = 0;
for(var pair_id: u32 = 0; pair_id < 512; pair_id += 1) {
let pxx = pairs[4*pair_id];
let pxy = pairs[4*pair_id+1];
let pyx = pairs[4*pair_id+2];
let pyy = pairs[4*pair_id+3];
if(pair_id % 32 == 0){
descriptor = 0;
n = 0;
}
let intensity_x = img[(i + pxx) + (j + pxy)*resolution[0]];
let intensity_y = img[(i + pyx) + (j + pyy)*resolution[0]];
if(intensity_x > intensity_y) {
descriptor += (u32(1) << n);
}
n += u32(1);
if(pair_id % 32 == 31){
descriptors[12*feature_id + (pair_id / 32)] = descriptor;
}
}
}
`,
});
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: pairs.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
});
resolutionBuffer = device.createBuffer({
label: 'resolution buffer',
size: resolution.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,
});
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 } },
{ binding: 4, resource: { buffer: resolutionBuffer } },
],
});
}
function mouseClicked() {
getKeypointsFromTexture(blurredCam, camg1, true);
runGPU(true);
}
async function runGPU(freeze) {
submitCommand(freeze);
}
async function submitCommand(freeze) {
// 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
l_n_features = freeze ? n_matchFeatures : n_features;
pass.dispatchWorkgroups(l_n_features);
pass.end();
// Encode a command to copy the results to a mappable buffer.
encoder.copyBufferToBuffer(descriptorBuffer, 0, descriptorStagingBuffer, 0, descriptorStagingBuffer.size);
// Finish encoding and submit the commands
commandBuffer = encoder.finish();
l_keypoints = freeze ? matchKeypoints : keypoints;
l_img = freeze ? matchImg : img;
device.queue.writeBuffer(keypointBuffer, 0, l_keypoints);
device.queue.writeBuffer(imageBuffer, 0, l_img);
device.queue.writeBuffer(pairsBuffer, 0, pairs);
device.queue.writeBuffer(resolutionBuffer, 0, resolution);
device.queue.submit([commandBuffer]);
await descriptorStagingBuffer.mapAsync(GPUMapMode.READ);
if(freeze) {
matchDescriptorResults = new Uint32Array(descriptorStagingBuffer.getMappedRange().slice());
} else {
descriptorResults = new Uint32Array(descriptorStagingBuffer.getMappedRange().slice());
}
await descriptorStagingBuffer.unmap()
if(!freeze) {
// console.log(matchDescriptorResults);
// console.log(descriptorResults);
for(let i = 0; i < n_features; ++i) {
let closest = findClosestFeature(i);
console.log(closest[1])
if(closest[1] < MIN_BITS_C) {
closestFeatures[i] = closest[0]
} else {
closestFeatures[i] = -1;
}
}
}
}
function features_to_line(i_match, i) {
let px = featureId_to_coordinate(i_match, matchKeypoints);
let py = featureId_to_coordinate(i, keypoints);
// console.log(i_match, i)
// console.log(matchKeypoints, keypoints)
// console.log(px, py)
stroke(0, 255, 0)
noFill();
line(px.x, px.y, py.x, py.y);
circle(px.x, px.y, 5);
circle(py.x, py.y, 5);
}
function featureId_to_coordinate(keypoint_id, l_keypoints) {
let pixel_id = l_keypoints[keypoint_id]
return {x: pixel_id % width, y: pixel_id / width};
}
function setup() {
initWebGPU();
// Generate random pairs
// TODO: Check if you have twice the same
// pairs[i] = px.x, pairs[i+1] = px.y, pairs[i+2] = py.x, pairs[i+3] = py.y
for(let i = 0; i < 4*N_PAIRS; i += 1) {
pairs[i] = floor(random() * S) - (S / 2);
}
console.log(pairs)
createCanvas(windowWidth, windowHeight);
camg1 = createGraphics(windowWidth, windowHeight, WEBGL);
camg1.pixelDensity(1);
camg1.noStroke();
blurredCam = createGraphics(windowWidth, windowHeight);
blurredCam.pixelDensity(1);
blurredCam.noStroke();
resolution[0] = width; //TODO: make this a 2vec ?
resolution[1] = height;
cam = createCapture(VIDEO);
cam.size(windowWidth, windowHeight);
cam.hide();
}
function draw() {
blurredCam.image(cam, 0, 0);
// console.log(cam)
blurredCam.filter(BLUR, 1.0);
camg1.shader(camShader1);
camShader1.setUniform('tex0', cam);
camShader1.setUniform('iResolutionX', width);
camShader1.setUniform('iResolutionY', height);
camg1.rect(0, 0, width, height);
image(camg1, 0, 0);
if(frameCount < 10) {
getKeypointsFromTexture(blurredCam, camg1, true);
} else {
getKeypointsFromTexture(blurredCam, camg1, false);
}
if(frameCount == 10) runGPU(true);
if(frameCount % 30 == 0 && frameCount != 0) runGPU(false);
for(let i = 0; i < n_features; ++i) {
if(closestFeatures[i] != -1)
features_to_line(i, closestFeatures[i]);
}
// 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 hammingDistance(n1, n2) {
let x = n1 ^ n2;
let setBits = 0;
while (x > 0) {
setBits += x & 1;
x >>= 1;
}
return setBits;
}
function findClosestFeature(feature_id) {
// console.log(b0, b1, b2, b3);
let min_diff = Infinity;
let min_diff_id = -1;
for(let i = 0; i < n_features; ++i) {
let diff = 0.0;
for(let j = 0; j < 12; ++j) {
diff += hammingDistance(matchDescriptorResults[12*feature_id + j], descriptorResults[12*i + j]);
}
if(diff < min_diff) {
min_diff = diff;
min_diff_id = i;
}
}
return [min_diff_id, min_diff];
}
function getKeypointsFromTexture(tex, texg, freeze) {
tex.loadPixels();
texg.loadPixels();
// console.log(tex.pixels.length, width*height*4)
l_n_features = 0;
l_img = freeze ? matchImg: img;
l_keypoints = freeze ? matchKeypoints: keypoints;
for(let i = 0; i < tex.pixels.length; i += 4) {
l_img[i / 4] = (tex.pixels[i] + tex.pixels[i+1] + tex.pixels[i+2]) / 3.0;
if(texg.pixels[i] == 255) {
l_keypoints[l_n_features] = i / 4;
l_n_features += 1;
if(l_n_features == MAX_FEATURES) break;
}
}
if(freeze){
n_matchFeatures = l_n_features;
} else {
n_features = l_n_features;
}
// console.log(img)
}
function windowResized(){
resizeCanvas(windowWidth, windowHeight);
}