Plane fitting with RANSAC

ransac-2.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math


def RANSAC(x, y, z, num_iter, threshold):
    best_inliers = []
    num_best_inliers = 0

    num_points = len(x)
    for i in range(0, num_iter):
        rand_idxs = np.random.randint(0, num_points, 3)
        x1 = x[rand_idxs]
        y1 = y[rand_idxs]
        z1 = z[rand_idxs]

        A = np.stack([x1, y1, z1], axis=1)
        b = np.ones((3, 1))

        A_inv = np.linalg.pinv(A)
        res = np.matmul(A_inv, b)

        dists = abs(res[0]*x + res[1]*y + res[2]*z - 1)/math.sqrt(res[0]*res[0] + res[1]*res[1] + res[2]*res[2])
        
        inliers = dists < threshold #Threshols
        num_inliers = sum(inliers)

        if num_inliers > num_best_inliers:
            num_best_inliers = num_inliers
            best_inliers = inliers
            print(i, num_inliers)

    return num_best_inliers, best_inliers



f = open('cluttered_hallway.txt', "r")
x = []
y = []
z = []
for line in f:
    tokens = line.split(" ")
    tokens[-1] = tokens[-1][:-1]
    x.append(float(tokens[0]))
    y.append(float(tokens[1]))
    z.append(float(tokens[2]))

f.close()


x = np.array(x)
y = np.array(y)
z = np.array(z)


num_iter = 1000
threshold = 0.008

num_best_inliers, best_inliers = RANSAC(x, y, z, num_iter, threshold)

x_plane1 = x[best_inliers]
y_plane1 = y[best_inliers]
z_plane1 = z[best_inliers]

outlier_x = x[np.logical_not(best_inliers)]
outlier_y = y[np.logical_not(best_inliers)]
outlier_z = z[np.logical_not(best_inliers)]

A = np.stack([x_plane1, y_plane1, z_plane1], axis=1)
b = np.ones((num_best_inliers, 1))

A_inv = np.linalg.pinv(A)
res_1 = np.matmul(A_inv, b)

print(res_1)

x_pts_1 = np.linspace(np.min(x_plane1), np.max(x_plane1), 100)
y_pts_1 = np.linspace(np.min(y_plane1), np.max(y_plane1), 100)

X_1, Y_1 = np.meshgrid(x_pts_1, y_pts_1)
P_1 = (1 - res_1[0]*X_1 - res_1[1]*Y_1)/res_1[2]


# 2nd plane
num_best_inliers, best_inliers = RANSAC(outlier_x, outlier_y, outlier_z, num_iter, threshold)

x_plane2 = outlier_x[best_inliers]
y_plane2 = outlier_y[best_inliers]
z_plane2 = outlier_z[best_inliers]

outlier_x = outlier_x[np.logical_not(best_inliers)]
outlier_y = outlier_y[np.logical_not(best_inliers)]
outlier_z = outlier_z[np.logical_not(best_inliers)]

A = np.stack([x_plane2, y_plane2, z_plane2], axis=1)
b = np.ones((num_best_inliers, 1))

A_inv = np.linalg.pinv(A)
res_2 = np.matmul(A_inv, b)

print(res_2)

x_pts_2 = np.linspace(np.min(x_plane2), np.max(x_plane2), 100)
y_pts_2 = np.linspace(np.min(y_plane2), np.max(y_plane2), 100)

X_2, Y_2 = np.meshgrid(x_pts_2, y_pts_2)
P_2 = (1 - res_2[0]*X_2 - res_2[1]*Y_2)/res_2[2]


# 3rd plane
num_best_inliers, best_inliers = RANSAC(outlier_x, outlier_y, outlier_z, num_iter, threshold)

x_plane3 = outlier_x[best_inliers]
y_plane3 = outlier_y[best_inliers]
z_plane3 = outlier_z[best_inliers]

outlier_x = outlier_x[np.logical_not(best_inliers)]
outlier_y = outlier_y[np.logical_not(best_inliers)]
outlier_z = outlier_z[np.logical_not(best_inliers)]

A = np.stack([x_plane3, y_plane3, z_plane3], axis=1)
b = np.ones((num_best_inliers, 1))

A_inv = np.linalg.pinv(A)
res_3 = np.matmul(A_inv, b)

print(res_3)

x_pts_3 = np.linspace(np.min(x_plane3), np.max(x_plane3), 100)
y_pts_3 = np.linspace(np.min(y_plane3), np.max(y_plane3), 100)

X_3, Y_3 = np.meshgrid(x_pts_3, y_pts_3)
P_3 = (1 - res_3[0]*X_3 - res_3[1]*Y_3)/res_3[2]


#4th plane
num_best_inliers, best_inliers = RANSAC(outlier_x, outlier_y, outlier_z, num_iter, threshold)

x_plane4 = outlier_x[best_inliers]
y_plane4 = outlier_y[best_inliers]
z_plane4 = outlier_z[best_inliers]

A = np.stack([x_plane4, y_plane4, z_plane4], axis=1)
b = np.ones((num_best_inliers, 1))

A_inv = np.linalg.pinv(A)
res_4 = np.matmul(A_inv, b)

print(res_4)

x_pts_4 = np.linspace(np.min(x_plane4), np.max(x_plane4), 100)
y_pts_4 = np.linspace(np.min(y_plane4), np.max(y_plane4), 100)

X_4, Y_4 = np.meshgrid(x_pts_4, y_pts_4)
P_4 = (1 - res_4[0]*X_4 - res_4[1]*Y_4)/res_4[2]


fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(x, y, z, color='black')



ax.plot_surface(X_1, Y_1, P_1, color='red')
ax.plot_surface(X_2, Y_2, P_2, color='green')
ax.plot_surface(X_3, Y_3, P_3, color='yellow')
ax.plot_surface(X_4, Y_4, P_4, color='blue')

plt.show()

ransac.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math


def RANSAC(x, y, z, num_iter, threshold):
    best_inliers = []
    num_best_inliers = 0

    num_points = len(x)
    for i in range(0, num_iter):
        rand_idxs = np.random.randint(0, num_points, 3)
        x1 = x[rand_idxs]
        y1 = y[rand_idxs]
        z1 = z[rand_idxs]

        A = np.stack([x1, y1, z1], axis=1)
        b = np.ones((3, 1))

        A_inv = np.linalg.pinv(A)
        res = np.matmul(A_inv, b)

        dists = abs(res[0]*x + res[1]*y + res[2]*z - 1)/math.sqrt(res[0]*res[0] + res[1]*res[1] + res[2]*res[2])
        
        inliers = dists < threshold #Threshols
        num_inliers = sum(inliers)

        if num_inliers > num_best_inliers:
            num_best_inliers = num_inliers
            best_inliers = inliers
            print(i, num_inliers)

    return num_best_inliers, best_inliers



f = open('clean_hallway.txt', "r")
x = []
y = []
z = []
for line in f:
    tokens = line.split(" ")
    tokens[-1] = tokens[-1][:-1]
    x.append(float(tokens[0]))
    y.append(float(tokens[1]))
    z.append(float(tokens[2]))

f.close()


x = np.array(x)
y = np.array(y)
z = np.array(z)


num_iter = 500
threshold = 0.008

num_best_inliers, best_inliers = RANSAC(x, y, z, num_iter, threshold)

x_plane1 = x[best_inliers]
y_plane1 = y[best_inliers]
z_plane1 = z[best_inliers]

outlier_x = x[np.logical_not(best_inliers)]
outlier_y = y[np.logical_not(best_inliers)]
outlier_z = z[np.logical_not(best_inliers)]

A = np.stack([x_plane1, y_plane1, z_plane1], axis=1)
b = np.ones((num_best_inliers, 1))

A_inv = np.linalg.pinv(A)
res_1 = np.matmul(A_inv, b)

print(res_1)

x_pts_1 = np.linspace(np.min(x_plane1), np.max(x_plane1), 100)
y_pts_1 = np.linspace(np.min(y_plane1), np.max(y_plane1), 100)

X_1, Y_1 = np.meshgrid(x_pts_1, y_pts_1)
P_1 = (1 - res_1[0]*X_1 - res_1[1]*Y_1)/res_1[2]


# 2nd plane
num_best_inliers, best_inliers = RANSAC(outlier_x, outlier_y, outlier_z, num_iter, threshold)

x_plane2 = outlier_x[best_inliers]
y_plane2 = outlier_y[best_inliers]
z_plane2 = outlier_z[best_inliers]

outlier_x = outlier_x[np.logical_not(best_inliers)]
outlier_y = outlier_y[np.logical_not(best_inliers)]
outlier_z = outlier_z[np.logical_not(best_inliers)]

A = np.stack([x_plane2, y_plane2, z_plane2], axis=1)
b = np.ones((num_best_inliers, 1))

A_inv = np.linalg.pinv(A)
res_2 = np.matmul(A_inv, b)

print(res_2)

x_pts_2 = np.linspace(np.min(x_plane2), np.max(x_plane2), 100)
y_pts_2 = np.linspace(np.min(y_plane2), np.max(y_plane2), 100)

X_2, Y_2 = np.meshgrid(x_pts_2, y_pts_2)
P_2 = (1 - res_2[0]*X_2 - res_2[1]*Y_2)/res_2[2]


# 3rd plane
num_best_inliers, best_inliers = RANSAC(outlier_x, outlier_y, outlier_z, num_iter, threshold)

x_plane3 = outlier_x[best_inliers]
y_plane3 = outlier_y[best_inliers]
z_plane3 = outlier_z[best_inliers]

outlier_x = outlier_x[np.logical_not(best_inliers)]
outlier_y = outlier_y[np.logical_not(best_inliers)]
outlier_z = outlier_z[np.logical_not(best_inliers)]

A = np.stack([x_plane3, y_plane3, z_plane3], axis=1)
b = np.ones((num_best_inliers, 1))

A_inv = np.linalg.pinv(A)
res_3 = np.matmul(A_inv, b)

print(res_3)

x_pts_3 = np.linspace(np.min(x_plane3), np.max(x_plane3), 100)
y_pts_3 = np.linspace(np.min(y_plane3), np.max(y_plane3), 100)

X_3, Y_3 = np.meshgrid(x_pts_3, y_pts_3)
P_3 = (1 - res_3[0]*X_3 - res_3[1]*Y_3)/res_3[2]


#4th plane
num_best_inliers, best_inliers = RANSAC(outlier_x, outlier_y, outlier_z, num_iter, threshold)

x_plane4 = outlier_x[best_inliers]
y_plane4 = outlier_y[best_inliers]
z_plane4 = outlier_z[best_inliers]

A = np.stack([x_plane4, y_plane4, z_plane4], axis=1)
b = np.ones((num_best_inliers, 1))

A_inv = np.linalg.pinv(A)
res_4 = np.matmul(A_inv, b)

print(res_4)

x_pts_4 = np.linspace(np.min(x_plane4), np.max(x_plane4), 100)
y_pts_4 = np.linspace(np.min(y_plane4), np.max(y_plane4), 100)

X_4, Y_4 = np.meshgrid(x_pts_4, y_pts_4)
P_4 = (1 - res_4[0]*X_4 - res_4[1]*Y_4)/res_4[2]


fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(x, y, z, color='black')



ax.plot_surface(X_1, Y_1, P_1, color='red')
ax.plot_surface(X_2, Y_2, P_2, color='green')
ax.plot_surface(X_3, Y_3, P_3, color='yellow')
ax.plot_surface(X_4, Y_4, P_4, color='blue')

plt.show()