* Keith O'Hara <kohara2@swarthmore.edu>

 *

 * K-Means Image Segmentation

 *

 * When clicked, the mouse's y-coordinate decides how many colors (clusters) to use

 * mouse's x-coordinate decides how much to weight the spatial features

 *

 * 's': save

 * 'c': show centers

 */

​

let L = 2;              // distance norm

let space = 0.5;        // spatial weighting factor

let showCenters = false;

​

let numClusters = 5;

let classified = []; // which cluster is each pixel in (index)

let clusters = []; // what does each cluster look like ([x, y, r, g, b])

​

let video;

​

function assignClosest() {

  video.loadPixels();

  for (let x = 0; x < video.width; x++) {

    for (let y = 0; y < video.height; y++) {

      let i = (y * video.width + x) * 4;

      let r = video.pixels[i];

      let g = video.pixels[i + 1];

      let b = video.pixels[i + 2];

      let closest_k = 0;

      let closest_dist = Number.MAX_VALUE;

      for (let j = 0; j < numClusters; j++) {

        let p = clusters[j];

        let cx = p[0];

        let cy = p[1];

        let cr = p[2];

        let cg = p[3];

        let cb = p[4];

​

        let d = pow(

          0 +

            pow(abs(cr - r) / 255.0, L) +

            pow(abs(cb - b) / 255.0, L) +

            pow(abs(cg - g) / 255.0, L) +

            space * pow(abs(x - cx) / video.width, L) +

            space * pow(abs(y - cy) / video.height, L),

          1.0 / L

        );

​

        if (d < closest_dist) {

          closest_dist = d;

          closest_k = j;

        }

      }

      classified[y * width + x] = closest_k;

    }

  }

}

​

function computeCentroids() {

  let count = new Array(numClusters).fill(0);

​

  // clear cluster means

  for (let i = 0; i < numClusters; i++) {

    for (let j = 0; j < clusters[i].length; j++) {

      clusters[i][j] = 0;

    }

  }

​

  for (let x = 0; x < video.width; x++) {

    for (let y = 0; y < video.height; y++) {

      let i = (y * video.width + x) * 4;

      let r = video.pixels[i];

      let g = video.pixels[i + 1];

      let b = video.pixels[i + 2];

​

      let v = [x, y, r, g, b];

​

      let clusterCenter = classified[y*width+x];

      count[clusterCenter]++;

      for (let z = 0; z < clusters[clusterCenter].length; z++) {

        clusters[clusterCenter][z] += v[z];

      }

    }

  }

​

  for (let j = 0; j < numClusters; j++) {

    for (let z = 0; z < clusters[j].length; z++) {

      if (count[j] > 0) clusters[j][z] = floor(clusters[j][z] / count[j]);

    }

  }

}

​

function setup() {

  createCanvas(640, 480);

  video = createCapture(VIDEO);

  video.size(width, height);

  video.hide();

  reset();

}

​

function reset() {

  classified = new Array(width * height).fill(0);

  clusters = new Array(numClusters).fill().map(() => new Array(5).fill(0));

  genCenters();

}

​

function keyPressed() {

  if (key == "s") {

    print("saving");

    saveCanvas("frame.jpg");

  }

  else if (key == "c"){

    showCenters = !showCenters;

  }

}

​

function genCenters() {

  video.loadPixels();

​

  for (let i = 0; i < numClusters; i++) {

    let x = int(random(width));

    let y = int(random(height));

    let z = (y * width + x)*4;

    let r = video.pixels[z];

    let g = video.pixels[z + 1];

    let b = video.pixels[z + 2];

    clusters[i] = [x, y, r, g, b];

  }

}

​

function mousePressed() {

  numClusters = int(map(mouseY, 0, height, 0, 120));

  print("number of colors: " + numClusters + " " + space);

  reset();

}

​

function draw() {

  background(0);

  space = map(mouseX, 0, width, 0, 2);

  video.loadPixels();

​

  assignClosest();

  computeCentroids();

​

  // Create a new image to store the resulting clustered image

  let cImg = createImage(video.width, video.height);

  cImg.loadPixels();

  for (let x = 0; x < video.width; x++) {

    for (let y = 0; y < video.height; y++) {

      let i = (y * video.width  + x) * 4;

      let v = clusters[classified[y * width + x]];

​

      cImg.pixels[i] = v[2];

      cImg.pixels[i + 1] = v[3];

      cImg.pixels[i + 2] = v[4];

      cImg.pixels[i + 3] = 255;

    }

  }

​

  cImg.updatePixels();

  image(cImg, 0, 0);

​

  if (showCenters){

    // draw cluster centers

    stroke(255, 255, 0, 128);

    strokeWeight(2);

    for (let i = 0; i < numClusters; i++) {

      let v = clusters[i];

      fill(v[2], v[3], v[4]);

      ellipse(clusters[i][0] , clusters[i][1] , 10, 10);

    }

  }

}

​