-
Notifications
You must be signed in to change notification settings - Fork 622
/
ransac.py
94 lines (81 loc) · 2.86 KB
/
ransac.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
from __future__ import division
import random
import numpy as np
def fit_plane(xyz,z_pos=None):
"""
if z_pos is not None, the sign
of the normal is flipped to make
the dot product with z_pos (+).
"""
mean = np.mean(xyz,axis=0)
xyz_c = xyz - mean[None,:]
l,v = np.linalg.eig(xyz_c.T.dot(xyz_c))
abc = v[:,np.argmin(l)]
d = -np.sum(abc*mean)
# unit-norm the plane-normal:
abcd = np.r_[abc,d]/np.linalg.norm(abc)
# flip the normal direction:
if z_pos is not None:
if np.sum(abcd[:3]*z_pos) < 0.0:
abcd *= -1
return abcd
def fit_plane_ransac(pts, neighbors=None,z_pos=None, dist_inlier=0.05,
min_inlier_frac=0.60, nsample=3, max_iter=100):
"""
Fits a 3D plane model using RANSAC.
pts : (nx3 array) of point coordinates
"""
n,_ = pts.shape
ninlier,models = [],[]
for i in xrange(max_iter):
if neighbors is None:
p = pts[np.random.choice(pts.shape[0],nsample,replace=False),:]
else:
p = pts[neighbors[:,i],:]
m = fit_plane(p,z_pos)
ds = np.abs(pts.dot(m[:3])+m[3])
nin = np.sum(ds < dist_inlier)
if nin/pts.shape[0] >= min_inlier_frac:
ninlier.append(nin)
models.append(m)
if models == []:
print "RANSAC plane fitting failed!"
return #None
else: #refit the model to inliers:
ninlier = np.array(ninlier)
best_model_idx = np.argsort(-ninlier)
n_refit, m_refit, inliers = [],[],[]
for idx in best_model_idx[:min(10,len(best_model_idx))]:
# re-estimate the model based on inliers:
dists = np.abs(pts.dot(models[idx][:3])+models[idx][3])
inlier = dists < dist_inlier
m = fit_plane(pts[inlier,:],z_pos)
# compute new inliers:
d = np.abs(pts.dot(m[:3])+m[3])
inlier = d < dist_inlier/2 # heuristic
n_refit.append(np.sum(inlier))
m_refit.append(m)
inliers.append(inlier)
best_plane = np.argmax(n_refit)
return m_refit[best_plane],inliers[best_plane]
if __name__ == '__main__':
from matplotlib import pylab
from mpl_toolkits import mplot3d
fig = pylab.figure()
ax = mplot3d.Axes3D(fig)
def plot_plane(a, b, c, d):
xx, yy = np.mgrid[10:20, 10:20]
return xx, yy, (-d - a * xx - b * yy) / c
n = 100
max_iterations = 100
goal_inliers = n * 0.3
# test data
xyzs = np.random.random((n, 3)) * 10 + 10
xyzs[:90, 2:] = xyzs[:90, :1]
ax.scatter3D(xyzs.T[0], xyzs.T[1], xyzs.T[2])
# RANSAC
m, b = run_ransac(xyzs, estimate, lambda x, y: is_inlier(x, y, 0.01), 3, goal_inliers, max_iterations)
a, b, c, d = m
xx, yy, zz = plot_plane(a, b, c, d)
ax.plot_surface(xx, yy, zz, color=(0, 1, 0, 0.5))
plt.show()