MVA-2021 / npm3d / hw5_modelisation / code / ransac.py
ransac.py
Raw
#
#
#      0===========================================================0
#      |                      TP6 Modelisation                     |
#      0===========================================================0
#
#
#------------------------------------------------------------------------------------------
#
#      Plane detection with RANSAC
#

import numpy as np
from ply import write_ply, read_ply
import time
from sklearn.neighbors import KDTree


def compute_plane(points):
    
    point_plane = points[0]
    normal_plane = np.cross(points[1] - points[0], points[2] - points[0])
    normal_plane /= np.linalg.norm(normal_plane)
        
    return point_plane, normal_plane



def in_plane(points, pt_plane, normal_plane, threshold_in=0.1):
    
    dists = np.dot(points-pt_plane.T, normal_plane)
    dists = np.abs(dists)
    return dists <= threshold_in


def RANSAC(points, nb_draws=100, threshold_in=0.1):
    
    best_vote = 3
    best_pt_plane = np.zeros((3,1))
    best_normal_plane = np.zeros((3,1))
    
    for _ in range(nb_draws):
        samples = points[np.random.randint(points.shape[0], size=3), :]
        pt_plane, normal_plane = compute_plane(samples)
        vote = np.count_nonzero( in_plane(points, pt_plane, normal_plane, threshold_in) )
        if vote > best_vote:
            best_pt_plane = pt_plane
            best_normal_plane = normal_plane
            best_vote = vote
                
    return best_pt_plane, best_normal_plane, best_vote


def RANSAC_var(points, nb_draws=100, threshold_in=0.1, sample_radius=0.3):
    
    best_vote = 3
    best_pt_plane = np.zeros((3,1))
    best_normal_plane = np.zeros((3,1))

    tree = KDTree(points)
    
    for _ in range(nb_draws):
        samples = np.zeros((3,3))
        samples[0] = points[np.random.randint(points.shape[0]), :]
        neighbors = tree.query_radius(X=[samples[0]], r=sample_radius , return_distance=False)[0]
        samples[1:3,:] = points[np.random.choice(neighbors, size=2)]
        pt_plane, normal_plane = compute_plane(samples)
        vote = np.count_nonzero( in_plane(points, pt_plane, normal_plane, threshold_in) )
        if vote > best_vote:
            best_pt_plane = pt_plane
            best_normal_plane = normal_plane
            best_vote = vote
                
    return best_pt_plane, best_normal_plane, best_vote


def recursive_RANSAC(points, nb_draws=100, threshold_in=0.1, nb_planes=10, sample_radius=0.3):
    
    nb_points = len(points)
    plane_inds = np.arange(0,0)
    plane_labels = np.arange(0,0)
    remaining_inds = np.arange(0,nb_points)
	
    for label in range(nb_planes):
        pt_plane, normal_plane, vote = RANSAC_var(points[remaining_inds], nb_draws, 
                                        threshold_in= threshold_in, sample_radius=sample_radius)
        # pt_plane, normal_plane, vote = RANSAC(points[remaining_inds], nb_draws, threshold_in)
        in_plane_bools = in_plane(points[remaining_inds], pt_plane, normal_plane, threshold_in)
        in_plane_inds, remaining_inds = remaining_inds[in_plane_bools], remaining_inds[~ in_plane_bools]
        print(len(in_plane_inds), len(remaining_inds))
        plane_inds = np.append(plane_inds, in_plane_inds)
        plane_labels = np.append(plane_labels, [label for _ in range(in_plane_inds.shape[0])])
    
    return plane_inds, remaining_inds, plane_labels


if __name__ == '__main__':


    # Load point cloud
    # ****************
    #
    #   Load the file '../data/indoor_scan.ply'
    #   (See read_ply function)
    #

    # Path of the file
    file_path = '../data/indoor_scan.ply'
    # file_path = '../data/Lille_street_small.ply'

    # Load point cloud
    data = read_ply(file_path)

    # Concatenate data
    points = np.vstack((data['x'], data['y'], data['z'])).T
    # colors = np.vstack((data['red'], data['green'], data['blue'])).T
    # labels = data['label']
    nb_points = len(points)
    

    # # Computes the plane passing through 3 randomly chosen points
    # # ************************
    # #
    
    # print('\n--- 1) and 2) ---\n')
    
    # # Define parameter
    # threshold_in = 0.10

    # # Take randomly three points
    # pts = points[np.random.randint(0, nb_points, size=3)]
    
    # # Computes the plane passing through the 3 points
    # t0 = time.time()
    # pt_plane, normal_plane = compute_plane(pts)
    # t1 = time.time()
    # print('plane computation done in {:.3f} seconds'.format(t1 - t0))
    
    # # Find points in the plane and others
    # t0 = time.time()
    # points_in_plane = in_plane(points, pt_plane, normal_plane, threshold_in)
    # t1 = time.time()
    # print('plane extraction done in {:.3f} seconds'.format(t1 - t0))
    # plane_inds = points_in_plane.nonzero()[0]
    # remaining_inds = (1-points_in_plane).nonzero()[0]
    
    # # Save extracted plane and remaining points
    # write_ply('../plane.ply', [points[plane_inds], colors[plane_inds], labels[plane_inds]], ['x', 'y', 'z', 'red', 'green', 'blue', 'label'])
    # write_ply('../remaining_points_plane.ply', [points[remaining_inds], colors[remaining_inds], labels[remaining_inds]], ['x', 'y', 'z', 'red', 'green', 'blue', 'label'])
    

    # # Computes the best plane fitting the point cloud
    # # ***********************************
    # #
    # #
    
    # print('\n--- 3) ---\n')

    # # Define parameters of RANSAC
    # nb_draws = 100
    # threshold_in = 0.10

    # # Find best plane by RANSAC
    # t0 = time.time()
    # best_pt_plane, best_normal_plane, best_vote = RANSAC(points, nb_draws, threshold_in)
    # t1 = time.time()
    # print('RANSAC done in {:.3f} seconds'.format(t1 - t0))
    
    # # Find points in the plane and others
    # points_in_plane = in_plane(points, best_pt_plane, best_normal_plane, threshold_in)
    # plane_inds = points_in_plane.nonzero()[0]
    # remaining_inds = (1-points_in_plane).nonzero()[0]
    
    # # Save the best extracted plane and remaining points
    # write_ply('../best_plane.ply', [points[plane_inds], colors[plane_inds], labels[plane_inds]], ['x', 'y', 'z', 'red', 'green', 'blue', 'label'])
    # write_ply('../remaining_points_best_plane.ply', [points[remaining_inds], colors[remaining_inds], labels[remaining_inds]], ['x', 'y', 'z', 'red', 'green', 'blue', 'label'])
    

    # Find "all planes" in the cloud
    # ***********************************
    #
    #
    
    print('\n--- 4) ---\n')
    
    # Define parameters of recursive_RANSAC
    nb_draws = 100
    threshold_in = 0.02
    nb_planes = 10
    
    # Recursively find best plane by RANSAC
    t0 = time.time()
    plane_inds, remaining_inds, plane_labels = recursive_RANSAC(points, nb_draws, threshold_in, nb_planes)
    t1 = time.time()
    print('recursive RANSAC done in {:.3f} seconds'.format(t1 - t0))
                
    # Save the best planes and remaining points
    # write_ply('../best_planes.ply', [points[plane_inds], colors[plane_inds], labels[plane_inds], plane_labels.astype(np.int32)], ['x', 'y', 'z', 'red', 'green', 'blue', 'label', 'plane_label'])
    # write_ply('../remaining_points_best_planes.ply', [points[remaining_inds], colors[remaining_inds], labels[remaining_inds]], ['x', 'y', 'z', 'red', 'green', 'blue', 'label'])
    write_ply('../best_planes.ply', [points[plane_inds], plane_labels.astype(np.int32)], ['x', 'y', 'z', 'plane_label'])
    write_ply('../remaining_points_best_planes.ply', [points[remaining_inds]], ['x', 'y', 'z'])
    
    
    print('Done')