// 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);

}