Created
May 8, 2025 10:20
-
-
Save zzhuolun/1552fc64c9be616abe05f962889a8e2c to your computer and use it in GitHub Desktop.
Ransac plane fitting with colors and normals
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import math | |
| import numpy as np | |
| import open3d as o3d | |
| from skimage.color import rgb2lab | |
| from dataclasses import dataclass | |
| @dataclass | |
| class O3DRANSACPlaneFitter: | |
| distance_threshold: float = 0.1 | |
| sample_size: int = 3 | |
| max_iterations: int = 1500 | |
| probability: float = 0.99999 | |
| def __post_init__(self): | |
| assert self.sample_size >= 3, "Number of inliers for plane fitting must be at least 3" | |
| assert 0 < self.probability <= 1, "Probability must be in the range (0, 1]" | |
| def fit_plane(self, points: np.ndarray) -> tuple[np.ndarray, np.ndarray, float, None]: | |
| """ | |
| Fit a plane to the points using Open3D's RANSAC (source code at https://github.com/isl-org/Open3D/blob/main/cpp/open3d/geometry/PointCloudSegmentation.cpp#L147) | |
| Returns: | |
| Tuple containing: | |
| - plane_coefficients: [a,b,c,d] where ax + by + cz + d = 0 | |
| - inlier_indices: Indices of inliers | |
| - fitness: Ratio of inliers | |
| - None: Not used | |
| """ | |
| pcd = o3d.geometry.PointCloud() | |
| pcd.points = o3d.utility.Vector3dVector(points) | |
| plane_model, inliers = pcd.segment_plane(distance_threshold=self.distance_threshold, ransac_n=self.sample_size, | |
| num_iterations=self.max_iterations, probability=self.probability) | |
| fitness = len(inliers) / len(points) | |
| return plane_model, inliers, fitness, None | |
| @dataclass | |
| class AugmentedPlaneFitter(O3DRANSACPlaneFitter): | |
| """Class for fitting planes using RANSAC with color and normal information""" | |
| distance_threshold: float = 0.1 | |
| sample_size: int = 3 | |
| max_iterations: int = 1500 | |
| probability: float = 0.99999 | |
| normal_threshold: float = np.pi / 6 # in radians | |
| lab_color_threshold: float = 50 # in LAB space | |
| def fit_plane(self, points: np.ndarray, normals: np.ndarray, | |
| colors: np.ndarray) -> tuple[np.ndarray, np.ndarray, float, float]: | |
| """ | |
| Fit a plane to the points using RANSAC with color and normal constraints | |
| Args: | |
| points: Nx3 array of 3D points | |
| normals: Nx3 array of point normals | |
| colors: Nx3 array of point colors in RGB [0,1] | |
| Returns: | |
| Tuple containing: | |
| - plane_coefficients: [a,b,c,d] where ax + by + cz + d = 0 | |
| - inlier_indices: Indices of inliers | |
| - fitness: Ratio of inliers | |
| - combined_metric: Combined metric of inliers | |
| """ | |
| num_points = len(points) | |
| if num_points < 3: | |
| raise ValueError("There must be at least 3 points.") | |
| colors_lab = rgb2lab(colors) | |
| best_fitness = 0 | |
| best_metric = float('inf') | |
| best_plane = None | |
| best_inliers = None | |
| break_iteration = self.max_iterations | |
| for iteration in range(self.max_iterations): | |
| if iteration > break_iteration: | |
| break | |
| # Randomly sample 3 points | |
| sample_indices = np.random.choice(num_points, self.sample_size, replace=False) | |
| sample_points = points[sample_indices] | |
| sample_colors = colors_lab[sample_indices] | |
| # Fit plane to these points | |
| if self.sample_size == 3: | |
| plane = self.compute_triangle_plane(sample_points[0], sample_points[1], sample_points[2]) | |
| else: | |
| plane = self.get_plane_from_points(sample_points) | |
| if np.allclose(plane, 0): | |
| continue | |
| # Find inliers and compute combined metric | |
| inliers, fitness, combined_metric = self._find_inliers_and_metric(points, normals, colors_lab, plane, | |
| sample_colors) | |
| if len(inliers) < 3: | |
| # print(f"At least 3 inliers required, got {len(inliers)}") | |
| continue | |
| # Update best model if better | |
| if fitness > best_fitness or (fitness == best_fitness and combined_metric < best_metric): | |
| best_fitness = fitness | |
| best_metric = combined_metric | |
| best_plane = plane | |
| best_inliers = inliers | |
| # Update break iteration based on probability | |
| if best_fitness < 1.0: | |
| break_iteration = min( | |
| int(math.log(1 - self.probability) / math.log(1 - math.pow(best_fitness, self.sample_size))), | |
| self.max_iterations) | |
| else: | |
| break_iteration = 0 | |
| if best_plane is None: | |
| raise ValueError("Failed to find a valid plane") | |
| # Final refinement using all inliers | |
| if best_inliers is not None and len(best_inliers) >= 3: | |
| if len(best_inliers) == 3: | |
| final_plane = self.compute_triangle_plane(points[best_inliers[0]], points[best_inliers[1]], | |
| points[best_inliers[2]]) | |
| else: | |
| final_plane = self.get_plane_from_points(points[best_inliers]) | |
| final_inliers, final_fitness, final_metric = self._find_inliers_and_metric( | |
| points, normals, colors_lab, final_plane, colors_lab[best_inliers]) | |
| return final_plane, final_inliers, final_fitness, final_metric | |
| else: | |
| return best_plane, best_inliers, best_fitness, best_metric | |
| def _find_inliers_and_metric(self, points: np.ndarray, normals: np.ndarray, colors: np.ndarray, plane: np.ndarray, | |
| sample_colors: np.ndarray) -> tuple[np.ndarray, float, float]: | |
| """ | |
| Find inliers based on distance, normal, and color (LAB space) constraints, | |
| and compute a combined metric for the inliers. | |
| Returns: | |
| - inlier indices | |
| - fitness (ratio of inliers) | |
| - combined metric (mean of normalized distances, angles, and color distances for inliers) | |
| """ | |
| plane_normal = plane[:3] | |
| plane_d = plane[3] | |
| plane_color = np.mean(sample_colors, axis=0) | |
| # Point-to-plane distances | |
| distances = np.abs(np.dot(points, plane_normal) + plane_d) | |
| # Angles between point normals and plane normal | |
| normal_angles = np.arccos(np.clip(np.dot(normals, plane_normal), -1.0, 1.0)) | |
| # Color distances in LAB space | |
| color_distances = np.linalg.norm(colors - plane_color, axis=1) | |
| # Masks for each constraint | |
| distance_mask = distances < self.distance_threshold | |
| normal_mask = normal_angles < self.normal_threshold | |
| color_mask = color_distances < self.lab_color_threshold | |
| inlier_mask = distance_mask & normal_mask & color_mask | |
| inlier_indices = np.where(inlier_mask)[0] | |
| # if len(inlier_indices) < 3: | |
| # print(distance_mask.sum(), normal_mask.sum(), color_mask.sum(), len(inlier_indices)) | |
| fitness = len(inlier_indices) / len(points) | |
| # --- Prefer horizontal planes --- | |
| # Compute angle between plane normal and [0,0,1] | |
| z_axis = np.array([0, 0, 1]) | |
| normal_alignment_angle = np.arccos(np.clip(np.dot(plane_normal, z_axis), -1.0, 1.0)) | |
| # Normalize by pi/2 (so 0 is best, pi/2 is worst for horizontal) | |
| normal_alignment_penalty = normal_alignment_angle / (np.pi / 2) | |
| if len(inlier_indices) == 0: | |
| combined_metric = float('inf') | |
| else: | |
| # Normalize each metric by its threshold, then average | |
| norm_dist = distances[inlier_indices] / self.distance_threshold | |
| norm_angle = normal_angles[inlier_indices] / self.normal_threshold | |
| norm_color = color_distances[inlier_indices] / self.lab_color_threshold | |
| combined_metric = np.mean(norm_dist + norm_angle + norm_color) | |
| # Add the normal alignment penalty (weighted, can tune if needed) | |
| combined_metric += normal_alignment_penalty * 0.5 | |
| return inlier_indices, fitness, combined_metric | |
| @staticmethod | |
| def compute_triangle_plane(p0: np.ndarray, p1: np.ndarray, p2: np.ndarray) -> np.ndarray: | |
| """ | |
| Fit a plane to exactly 3 points. Taken from open3d.TriangleMesh::ComputeTrianglePlane | |
| Returns [a, b, c, d] for the plane equation: ax + by + cz + d = 0 | |
| If the points are colinear, returns [0, 0, 0, 0]. | |
| """ | |
| e0 = p1 - p0 | |
| e1 = p2 - p0 | |
| normal = np.cross(e0, e1) | |
| norm = np.linalg.norm(normal) | |
| if norm == 0: | |
| return np.zeros(4) | |
| normal = normal / norm | |
| # Ensure normal points upwards | |
| if normal[2] < 0: | |
| normal = -normal | |
| d = -np.dot(normal, p0) | |
| return np.array([normal[0], normal[1], normal[2], d]) | |
| @staticmethod | |
| def get_plane_from_points(points: np.ndarray) -> np.ndarray: | |
| """Fit a plane to points using the improved method from ilikebigbits.com (2017) | |
| This method better handles noisy points by using weighted combinations of axis directions | |
| based on their determinants, rather than just picking the axis with the largest determinant. | |
| """ | |
| n = len(points) | |
| if n < 3: | |
| return np.zeros(4) | |
| # Calculate centroid | |
| centroid = np.mean(points, axis=0) | |
| # Calculate full 3x3 covariance matrix, excluding symmetries | |
| centered_points = points - centroid | |
| xx = np.sum(centered_points[:, 0] * centered_points[:, 0]) / n | |
| xy = np.sum(centered_points[:, 0] * centered_points[:, 1]) / n | |
| xz = np.sum(centered_points[:, 0] * centered_points[:, 2]) / n | |
| yy = np.sum(centered_points[:, 1] * centered_points[:, 1]) / n | |
| yz = np.sum(centered_points[:, 1] * centered_points[:, 2]) / n | |
| zz = np.sum(centered_points[:, 2] * centered_points[:, 2]) / n | |
| # Initialize weighted direction | |
| weighted_dir = np.zeros(3) | |
| # Calculate weighted direction for each axis | |
| # X-axis | |
| det_x = yy * zz - yz * yz | |
| axis_dir_x = np.array([det_x, xz * yz - xy * zz, xy * yz - xz * yy]) | |
| weight_x = det_x * det_x | |
| if np.dot(weighted_dir, axis_dir_x) < 0: | |
| weight_x = -weight_x | |
| weighted_dir += axis_dir_x * weight_x | |
| # Y-axis | |
| det_y = xx * zz - xz * xz | |
| axis_dir_y = np.array([xz * yz - xy * zz, det_y, xy * xz - yz * xx]) | |
| weight_y = det_y * det_y | |
| if np.dot(weighted_dir, axis_dir_y) < 0: | |
| weight_y = -weight_y | |
| weighted_dir += axis_dir_y * weight_y | |
| # Z-axis | |
| det_z = xx * yy - xy * xy | |
| axis_dir_z = np.array([xy * yz - xz * yy, xy * xz - yz * xx, det_z]) | |
| weight_z = det_z * det_z | |
| if np.dot(weighted_dir, axis_dir_z) < 0: | |
| weight_z = -weight_z | |
| weighted_dir += axis_dir_z * weight_z | |
| # Normalize the weighted direction to get the normal | |
| norm = np.linalg.norm(weighted_dir) | |
| if norm == 0: | |
| return np.zeros(4) | |
| normal = weighted_dir / norm | |
| # Ensure normal points upwards | |
| if normal[2] < 0: | |
| normal = -normal | |
| # Calculate d using the centroid | |
| d = -np.dot(normal, centroid) | |
| return np.append(normal, d) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment