plane_ransac/classes.py at main · SlyPredator/plane_ransac · GitHub

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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import numpy as np
import open3d as o3d
from scipy.spatial import ConvexHull


"""
Summary:

This file goes about implementing plane fitting via two ways for a head-to-head comparison:
    - RANSAC via numpy and dual-SVD to try and optimize the detected inliers
    - ConvexHull to return plane area
    - Using Open3D to compare my solution vs optimized soln

exporter also pushes a photo of the detected plane in each frame, verify with that
"""

# Camera intrinsics weren't provided, so used D435i's intrinsics for 640x480

FX = 387.776184082031
FY = 387.776184082031
CX = 320.546813964844
CY = 246.21598815918
CAMERA_NORMAL = np.array([0, 0, 1])


class PlaneModel(object):
    def __init__(self, feat_dim=3):
        self._feat_dim = 3

    def fit_plane(self, points):
        points = self._check_data(points)
        X_mean = np.mean(points, axis=0)
        Xc = points - X_mean[np.newaxis, :]
        _, _, W = np.linalg.svd(Xc)
        plane_normal = W[-1, :]
        w_0 = np.dot(X_mean, -plane_normal)
        return np.concatenate(([w_0], plane_normal))

    def get_error(self, points, w):
        points = self._check_data(points)
        return np.abs(w[0] + points.dot(w[1:])) / np.linalg.norm(w[1:])

    def _check_data(self, points):
        if not self._feat_dim in points.shape:
            raise ValueError("Data shape is wrong, check again")
        return points if points.shape[1] == 3 else points.T


class RansacModel(object):
    """
    - Selects random points, fits an initial idea of the plane, calculates errors
    - This is done again if it passes a threshold as a polisher
    - Done iteratively over the max_iters defined
    """
    def fit(self, points, model, n_pts_fit_model, n_min_pts_inlier, max_iter, dist_thresh):
        if points.shape[1] != 3: points = points.T
        N = points.shape[0]
        if N < n_min_pts_inlier: return False, None, None

        best_w, best_res_num_inliers, best_res_inliers = None, -1, []
        for i in range(max_iter):
            rand_indices = np.random.permutation(N)
            maybe_idxs = rand_indices[:n_pts_fit_model + 1]
            maybe_data = points[maybe_idxs]
            maybe_w = model.fit_plane(maybe_data)
            all_error = model.get_error(points, maybe_w)

            inlier_mask = all_error < dist_thresh
            n_pts_inlier = np.count_nonzero(inlier_mask)

            if n_pts_inlier >= n_min_pts_inlier:
                also_idxs = np.arange(N)[inlier_mask]
                np.random.shuffle(also_idxs)
                also_idxs = also_idxs[:n_min_pts_inlier]
                also_w = model.fit_plane(points[also_idxs])

                if n_pts_inlier > best_res_num_inliers:
                    best_res_num_inliers = n_pts_inlier
                    best_w = also_w
                    best_res_inliers = np.arange(N)[model.get_error(points, also_w) < dist_thresh]

        return (True, best_w, best_res_inliers) if best_w is not None else (False, None, None)


"""
Helper functions
    - Calculation of area via ConvexHull (SVD + flattening 3D to 2D + find smallest ENCLOSED convex polygon)
    - Comparison function: using segment_plane from O3D (clip for [-1.0, 1.0] for arccos)
"""

def plane_area(pts):
    centr = np.mean(pts, axis=0)
    _, _, vh = np.linalg.svd(pts - centr, full_matrices=False)
    normal = vh[-1]
    ref = np.array([1, 0, 0]) if abs(normal[0]) < 0.9 else np.array([0, 1, 0])
    u = np.cross(normal, ref)
    u /= np.linalg.norm(u)
    v = np.cross(normal, u)
    proj = np.stack([pts @ u, pts @ v], axis=1)
    return ConvexHull(proj).area


def o3d_ransac(pts, dist_thresh=0.02):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(pts)
    try:
        plane_model, _ = pcd.segment_plane(distance_threshold=dist_thresh,
                                           ransac_n=3,
                                           num_iterations=500)
        n_o3d = np.array(plane_model[:3])
        n_o3d /= np.linalg.norm(n_o3d)
        angle = np.degrees(np.arccos(np.clip(np.abs(np.dot(n_o3d, CAMERA_NORMAL)), -1.0, 1.0)))
        return angle
    except:
        return 0.0