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

}