* RANSAC Line Fitting

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

 * March 2024

 */

​

// the points we are fitting

let ps = [];

​

function setup() {

  createCanvas(640, 480);

  noStroke();

  background(255);

}

​

function mouseDragged() {

  ps.push(new p5.Vector(mouseX, mouseY, 1));

}

​

function compute_error(pts, l) {

  let tot_error = 0;

  for (let j = 0; j < pts.length; j++) {

    let error = l.dot(pts[j]);

    tot_error += sq(error);

  }

  return tot_error;

}

​

function shuffleArray(array) {

  for (let i = array.length - 1; i > 0; i--) {

    const j = Math.floor(Math.random() * (i + 1));

    [array[i], array[j]] = [array[j], array[i]];

  }

}

​

function estimate_ransac(pts, model_n = 5, t = 1, n_inners = 10, n_rounds = 150) {

  let min_error = Number.MAX_VALUE;

  let min_model = new p5.Vector(0, 0, 1);

  if (pts.length < model_n) return min_model;

  let inners = [];

  let min_inners = [];

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

    // pick a few random points for a candidate model

    shuffleArray(pts);

    let model_pts = pts.slice(0, model_n);

    let cand_model = estimate_least_squares(model_pts);

    let inners = [];

    for (let j = 0; j < pts.length; j++) {

      let error = sq(cand_model.dot(pts[j]));

      if (error < t) {

        inners.push(pts[j]);

      }

    }

    if (inners.length > n_inners) {

      cand_model = estimate_least_squares(inners);

      let tot_error = compute_error(inners, cand_model);

      if (tot_error < min_error) {

        min_error = tot_error;

        min_inners = inners;

        min_model = cand_model;

      }

    }

  }

  for (let i = 0; i < min_inners.length; i++) {

    ellipse(min_inners[i].x, min_inners[i].y, 12, 12);

  }

  return min_model;

}

​

function estimate_least_squares(p) {

  let x = new p5.Vector(0, 0, 0);

​

  if (p.length >= 2) {

    as = [];

    bs = [];

    // Creates the estimation matrix

    for (let i = 0; i < p.length; i++) {

      as.push([p[i].x, 1]);

      bs.push([p[i].y]);

    }

    let mls_a = new ML.Matrix(as);

    let mls_b = new ML.Matrix(bs);

    let mls_x = ML.MatrixLib.solve(mls_a, mls_b);

    let a = mls_x.data[0][0];

    let b = -1;

    let c = mls_x.data[1][0];

    x = new p5.Vector(a, b, c);

  }

  return x;

}

​

function draw() {

  background(255);

  noStroke();

  fill(0);

  for (let i = 0; i < ps.length; i++) {

    ellipse(ps[i].x, ps[i].y, 5, 5);

  }

​

  if (ps.length >= 2) {

    stroke(255, 0, 0);

    strokeWeight(2);

    let ln = estimate_least_squares(ps);

    plotLine(ln);

​

    textSize(24);

    noFill();

    text("least squares", 20, 20);

​

    stroke(0, 0, 255);

    strokeWeight(2);

    let rln = estimate_ransac(ps);

    plotLine(rln);

​

    text("ransac", 20, 40);

  }

}

​

function plotLine(v) {

  if (abs(v.x) < abs(v.y)) {

    // mostly horizontal

    line(0, v.z / -v.y, width, (v.x * width + v.z) / -v.y);

  } else {

    //mostly vertical

    line(v.z / -v.x, 0, (v.y * height + v.z) / -v.x, height);

  }

}

​