Ransac plane fitting

import random

import numpy as np
import functools

from scipy.optimize import leastsq



def homog_points(points):
  # Expand point vector x,y,z to x,y,z,1
  n = points.shape[0]
  return  np.hstack([points, np.ones( (n, 1) )])


def plane_distances(plane_eq, points):
  normal = plane_eq[0:3]

  distance = np.tensordot(normal, points, (0, 1)) + plane_eq[3]
  return distance / np.linalg.norm(normal)

def optimize_plane(plane_eq, points):
  def eval(params):
    return plane_distances(params, points)

  optimized, _ = leastsq(eval, plane_eq) 
  return optimized / np.linalg.norm(optimized[0:3])
  
def plane3(points):
  assert points.shape == (3, 3)

  vecA = points[1] - points[0]
  vecB = points[2] - points[0]

  # Now we compute the cross product of vecA and vecB to get vecC which is normal to the plane
  normal = np.cross(vecA, vecB)

  # The plane equation will be normal[0]*x + normal[1]*y + normal[0]*z + k = 0
  # We have to use a point to find k
  normal = normal / np.linalg.norm(normal)
  k = -np.sum(np.multiply(normal, points[1, :]))
  return np.array([*normal, k])  


def ransac_plane(points, thresh=0.05, min_points=100, max_iterations=1000):
    """
    Ransac plane fitting.

    :param points: 3D point cloud as a `np.array (N,3)`.
    :param thresh: Threshold distance from the plane which is considered inlier.
    :param maxIteration: Number of maximum iteration which RANSAC will loop over.
    :returns:
    - `self.equation`:  Parameters of the plane using Ax+By+Cy+D `np.array (1, 4)`
    - `self.inliers`: points from the dataset considered inliers

    ---
    """
    n_points = points.shape[0]
    for _ in range(max_iterations):

        # Samples 3 random points
        pt_samples = points[random.sample(range(0, n_points), 3)]

        # We have to find the plane equation described by those 3 points
        # We find first 2 vectors that are part of this plane
        plane_eq = plane3(pt_samples)
        distances = plane_distances(plane_eq, points)

        # Select indexes where distance is biggers than the threshold
        inliers = np.where(np.abs(distances) <= thresh)[0]
        if len(inliers) > min_points:
          return plane_eq, inliers
          
    return None


def random_plane_points(num_points = 1000, size=5.0, noise_level=0.1):
  plane_points = np.random.rand(3, 3)
  plane_eq = plane3(plane_points)

  vecA = plane_points[1] - plane_points[0]
  vecB = plane_points[2] - plane_points[0]

  vecA = vecA * size / np.linalg.norm(vecA)
  vecB = vecB * size / np.linalg.norm(vecB)

  x = np.random.randn(num_points, 2)

  points = x[:, 0:1] * vecA.reshape(-1, 3) + x[:, 1:2] * vecB.reshape(-1, 3) + plane_points[0]
  noise = np.random.randn(*points.shape) * noise_level

  return plane_eq, points + noise


if __name__ == '__main__':

  plane, points = random_plane_points(5000, size=5.0, noise_level=0.2)
  found = ransac_plane(points, thresh=0.4)

  if (found):
    found_plane, inliers = found
    opt_plane = optimize_plane(found_plane, points[inliers])

    print(plane)

    print(len(inliers))
    print(found_plane)
    print(opt_plane)

Fancier version which works as one function... I guess have a play and see which works best.

import random

import numpy as np
import functools

from scipy.optimize import least_squares



def homog_points(points):
  # Expand point vector x,y,z to x,y,z,1
  n = points.shape[0]
  return  np.hstack([points, np.ones( (n, 1) )])


def plane_distances(plane_eq, points):
  normal = plane_eq[0:3]

  distance = np.tensordot(normal, points, (0, 1)) + plane_eq[3]
  return distance / np.linalg.norm(normal)

def optimize_plane(plane_eq, points):
  def eval(params):
    return plane_distances(params, points)

  optimized, _ = least_squares(eval, plane_eq, x_scale='jac') 
  return optimized / np.linalg.norm(optimized[0:3])


def optimize_plane_soft(plane_eq, points, err_scale):
  def eval(params):
    return plane_distances(params, points)

  result = least_squares(eval, plane_eq, loss='soft_l1', f_scale=err_scale, x_scale='jac') 
  return result.x / np.linalg.norm(result.x[0:3])

  
def plane3(points):
  assert points.shape == (3, 3)

  vecA = points[1] - points[0]
  vecB = points[2] - points[0]

  # Now we compute the cross product of vecA and vecB to get vecC which is normal to the plane
  normal = np.cross(vecA, vecB)

  # The plane equation will be normal[0]*x + normal[1]*y + normal[0]*z + k = 0
  # We have to use a point to find k
  normal = normal / np.linalg.norm(normal)
  k = -np.sum(np.multiply(normal, points[1, :]))
  return np.array([*normal, k])  


def select_inliers(plane_eq, points, threshold):
  distances = plane_distances(plane_eq, points)
  return np.where(np.abs(distances) <= threshold)[0]


def ransac_plane(points, threshold=0.05, min_points=100, outlier_iters=1, max_iterations=100):
    """
    Ransac plane fitting.

    :param points: 3D point cloud as a `np.array (N,3)`.
    :param thresh: Threshold distance from the plane which is considered inlier.
    :param maxIteration: Number of maximum iteration which RANSAC will loop over.
    :returns:
    - `self.equation`:  Parameters of the plane using Ax+By+Cy+D `np.array (1, 4)`
    - `self.inliers`: points from the dataset considered inliers

    ---
    """
    n_points = points.shape[0]
    best_inliers = np.array([], np.int)
    best_plane = None

    for _ in range(max_iterations):

        # Samples 3 random points
        pt_samples = points[random.sample(range(0, n_points), 3)]

        # We have to find the plane equation described by those 3 points
        # We find first 2 vectors that are part of this plane
        plane_eq = plane3(pt_samples)

        inliers = select_inliers(plane_eq, points, threshold)
        # Select indexes where distance is biggers than the threshold
        for _ in range(outlier_iters):
          optimize_plane_soft(plane_eq, points[inliers], err_scale = threshold * 0.25)
          inliers = select_inliers(plane_eq, points, threshold)

        if len(inliers) > min_points:
          return plane_eq, inliers
        else:
          best_plane = plane_eq
          best_inliers = inliers
          
    return best_plane, best_inliers


def random_plane_points(num_points = 4000, size=5.0, noise_level=0.1):
  plane_points = np.random.rand(3, 3)
  plane_eq = plane3(plane_points)

  vecA = plane_points[1] - plane_points[0]
  vecB = plane_points[2] - plane_points[0]

  vecA = vecA * size / np.linalg.norm(vecA)
  vecB = vecB * size / np.linalg.norm(vecB)

  x = np.random.randn(num_points, 2)

  points = x[:, 0:1] * vecA.reshape(-1, 3) + x[:, 1:2] * vecB.reshape(-1, 3) + plane_points[0]
  noise = np.random.randn(*points.shape) * noise_level

  return plane_eq, points + noise


if __name__ == '__main__':
  np.set_printoptions(precision=2, suppress=True)


  for noise_level in np.linspace(0, 1.0, num=50):

    plane_truth, points = random_plane_points(5000, size=2.0, noise_level=0.2)
    plane_eq, inliers = ransac_plane(points, min_points=1000, threshold=0.1)

    print(plane_truth)
    print(plane_eq)

    print(f"noise={noise_level:.2f}, inliers: {len(inliers)} accuracy: { abs(np.dot(plane_truth[0:3], plane_eq[0:3] )):.2f}  ")