utils.py

##https://medium.com/@ajithraj_gangadharan/3d-ransac-algorithm-for-lidar-pcd-segmentation-315d2a51351
import random
import numpy as np
import math

# test data in polar coordinates
rho_test = np.array([[10, 11, 11.7, 13, 14, 15, 16, 17, 17, 17, 16.5,
17, 17, 16, 14.5, 14, 13]]).T
n = rho_test.shape[0]
theta_test = (math.pi/180)*np.linspace(0, 85, n).reshape(-1,1)

def polar_to_cartesian(points):
    #From the lidar data, the range is the distance (r) and the index represents the angle of the

    #PDF says -45 to 180 vut I think that is wrong
    start_angle = -135
    end_angle = 135

    angle_step = (end_angle - start_angle)/len(points) #Degree of step over for each lidar measurement

    rho = points
    thetas = np.deg2rad(np.linspace(start_angle, end_angle, len(points)))

    x = rho*np.cos(thetas)
    y = rho*np.sin(thetas)

    return x, y 

#points is point cloud in cartesian coordinates
def ransac(x, y, max_iterations, inlier_threshold):
    """
    RANSAC algorithm to fit a line to 2D points

    Args:
        x: x-coordinates of points
        y: y-coordinates of points
        max_iterations: maximum number of RANSAC iterations
        inlier_threshold: distance threshold to consider a point as an inlier

    Returns:
        p0: (x, y) start point of the line, or None if no line found
        p1: (x, y) end point of the line, or None if no line found
    """
    n_points = len(x)
    best_inlier_count = 0
    best_line = None
    best_inliers = None

    for iteration_num in range(max_iterations):
        # Randomly select two different points
        indices = random.sample(range(n_points), 2)
        num1, num2 = indices[0], indices[1]

        p1 = np.array([x[num1], y[num1]])
        p2 = np.array([x[num2], y[num2]])

        # Fit a line between the two points
        # General form of line in 2d - Ax + By + C = 0
        # Using two-point form: (y2-y1)x - (x2-x1)y + (x2-x1)y1 - (y2-y1)x1 = 0
        A = p2[1] - p1[1]  # y2 - y1
        B = p1[0] - p2[0]  # x1 - x2
        C = p2[0]*p1[1] - p1[0]*p2[1]  # x2*y1 - x1*y2

        # Normalize the line equation for consistent distance calculation
        norm = np.sqrt(A**2 + B**2)
        if norm < 1e-10:  # Skip if points are too close
            continue

        A, B, C = A/norm, B/norm, C/norm

        # Calculate distance from all points to the line
        # Distance = |Ax + By + C| / sqrt(A^2 + B^2), but we already normalized
        distances = np.abs(A*x + B*y + C)

        # Count inliers
        inliers = distances < inlier_threshold
        inlier_count = np.sum(inliers)

        # Update best model if this one has more inliers
        if inlier_count > best_inlier_count:
            best_inlier_count = inlier_count
            best_line = (A, B, C)
            best_inliers = inliers

    # Convert line coefficients to start/end points using actual inlier points
    if best_line is None or best_inliers is None:
        return None, None

    A, B, C = best_line
    x_inliers = x[best_inliers]
    y_inliers = y[best_inliers]

    if len(x_inliers) < 2:
        return None, None

    # Find the two inlier points that are farthest apart
    # This gives us the actual extent of the fitted line based on real data
    max_dist = 0
    p0_idx = 0
    p1_idx = 0

    for i in range(len(x_inliers)):
        for j in range(i + 1, len(x_inliers)):
            dist = np.sqrt((x_inliers[i] - x_inliers[j])**2 + (y_inliers[i] - y_inliers[j])**2)
            if dist > max_dist:
                max_dist = dist
                p0_idx = i
                p1_idx = j

    # Return the two farthest inlier points as endpoints
    x0 = float(x_inliers[p0_idx])
    y0 = float(y_inliers[p0_idx])
    x1 = float(x_inliers[p1_idx])
    y1 = float(y_inliers[p1_idx])

    return (x0, y0), (x1, y1)


def segment_lidar(ranges, jump_threshold=1.0, min_range=0.1, max_range=50.0, max_segment_size=50):
    """
    Segment laser scan based on Euclidean distance AND maximum segment size.

    This prevents huge segments from spanning multiple walls.

    Returns: list of index arrays for each segment
    """
    segments = []
    current_segment = []

    # Lidar parameters (adjust based on your sensor)
    start_angle = -135  # degrees
    end_angle = 135     # degrees
    n_points = len(ranges)

    # Calculate angle for each measurement
    angles = np.deg2rad(np.linspace(start_angle, end_angle, n_points))

    for i in range(len(ranges)):
        r = ranges[i]

        # Skip invalid ranges
        if r < min_range or r > max_range or np.isnan(r) or np.isinf(r):
            if len(current_segment) >= 2:
                segments.append(current_segment)
            current_segment = []
            continue

        # Force segment split if it gets too large
        if len(current_segment) >= max_segment_size:
            if len(current_segment) >= 2:
                segments.append(current_segment)
            current_segment = [i]
            continue

        # Check Euclidean distance to previous point
        if len(current_segment) > 0:
            last_idx = current_segment[-1]

            # Convert both points to Cartesian
            x1 = ranges[last_idx] * np.cos(angles[last_idx])
            y1 = ranges[last_idx] * np.sin(angles[last_idx])
            x2 = r * np.cos(angles[i])
            y2 = r * np.sin(angles[i])

            # Euclidean distance between consecutive points
            euclidean_dist = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)

            if euclidean_dist > jump_threshold:
                # Large jump - start new segment
                if len(current_segment) >= 2:
                    segments.append(current_segment)
                current_segment = [i]
            else:
                current_segment.append(i)
        else:
            current_segment.append(i)

    # Save final segment
    if len(current_segment) >= 2:
        segments.append(current_segment)

    return segments


# Test with the provided data
if __name__ == "__main__":
    import matplotlib.pyplot as plt

    # Convert test data from polar to cartesian
    x_test = rho_test * np.cos(theta_test)
    y_test = rho_test * np.sin(theta_test)

    # Flatten the arrays for RANSAC
    x_test = x_test.flatten()
    y_test = y_test.flatten()

    # Run RANSAC
    max_iterations = 10000
    inlier_threshold = 5

    best_line, best_inliers = ransac(x_test, y_test, max_iterations, inlier_threshold)

    print(f"Best line coefficients (Ax + By + C = 0):")
    print(f"A = {best_line[0]:.4f}")
    print(f"B = {best_line[1]:.4f}")
    print(f"C = {best_line[2]:.4f}")
    print(f"\nNumber of inliers: {np.sum(best_inliers)}/{len(x_test)}")
    print(f"Inlier indices: {np.where(best_inliers)[0]}")

    # Visualization
    plt.figure(figsize=(10, 8))

    # Plot inliers and outliers
    plt.scatter(x_test[best_inliers], y_test[best_inliers],
                c='green', label='Inliers', s=100, alpha=0.7)
    plt.scatter(x_test[~best_inliers], y_test[~best_inliers],
                c='red', label='Outliers', s=100, alpha=0.7)

    # Plot the fitted line
    # Line equation: Ax + By + C = 0 => y = -(A*x + C)/B
    A, B, C = best_line
    x_range = np.linspace(x_test.min() - 2, x_test.max() + 2, 100)

    if abs(B) > 1e-10:  # Not a vertical line
        y_line = -(A * x_range + C) / B
        plt.plot(x_range, y_line, 'b-', linewidth=2, label='RANSAC Fitted Line')
    else:  # Vertical line
        x_line = -C / A
        plt.axvline(x=x_line, color='b', linewidth=2, label='RANSAC Fitted Line')

    plt.xlabel('X', fontsize=12)
    plt.ylabel('Y', fontsize=12)
    plt.title('RANSAC Line Fitting Results', fontsize=14)
    plt.legend(fontsize=10)
    plt.grid(True, alpha=0.3)
    plt.axis('equal')
    plt.tight_layout()
    plt.show()