segmentation

import open3d as o3d
import open3d.core as o3c
import numpy as np
import matplotlib.pyplot as plt
import pybullet as p
import pybullet_data
import time

from sklearn.neighbors import NearestNeighbors
import plotly.express as px

# Initialize the physics simulation
physicsClient = p.connect(p.GUI)
p.setAdditionalSearchPath(pybullet_data.getDataPath())
p.setGravity(0, 0, -9.81)

# Load the objects into the simulation
planeId = p.loadURDF("plane.urdf")
sphereId = p.loadURDF("sphere_small.urdf", [0, 1, 0.01])
sphereId2 = p.loadURDF("sphere_small.urdf", [0.5, 1, 0.01])
#robotId = p.loadURDF("r2d2.urdf", [0.5,0.5,0.5],useFixedBase=True)
cubeId = p.loadURDF("cube_small.urdf", [0, 0.5, 0.01])
cubeId2 = p.loadURDF("cube_small.urdf", [0, 0.9, 0.01])

while True:
    # Get the image from the camera
    camera_pos = [0, 1, 2]
    viewMatrix = p.computeViewMatrix(
        cameraEyePosition=camera_pos,
        cameraTargetPosition=[0, 1, 0],
        cameraUpVector=[0, 1, 0])

    # Set the intrinsic parameters of the camera
    fov = 60
    aspect = 320.0 / 240.0
    near = 0.01
    far = 100

    # Get the image from the camera
    width, height, rgbImg, depthImg, segImg = p.getCameraImage(width=320, height=240, viewMatrix=viewMatrix,
                                                               projectionMatrix=p.computeProjectionMatrixFOV(fov,
                                                                                                             aspect,
                                                                                                             near, far))
    # Convert the depth image to a point cloud
    points = np.zeros((height * width, 3))
    for y in range(height):
        for x in range(width):
            depth = depthImg[y, x]
            if depth != 0:
                points[y * width + x, 0] = (x - 160) * depth / 160.0
                points[y * width + x, 1] = (y - 120) * depth / 120.0
                points[y * width + x, 2] = depth

    # Create Open3D point cloud object
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)


    """# Compute the distance graph
    neigh = NearestNeighbors(n_neighbors=5)
    nbrs = neigh.fit(pcd.points)
    distances, indices = nbrs.kneighbors(pcd.points)
    distances = np.sort(distances, axis=0)
    distances = distances[:, 1]

    # Plot the distance graph
    plt.plot(distances)
    plt.xlabel('Data point index')
    plt.ylabel('Distance to the nearest neighbor')
    plt.show()"""
    # Preprocess the data
    #pcd_filtered = pcd.voxel_down_sample(voxel_size=0.021)

    # Apply segmentation algorithm
    labels = pcd.cluster_dbscan(eps=0.0103655, min_points=9)

    # Assign random colors to each segment
    num_segments = len(set(labels))
    print("Number of segments")
    print(num_segments)
    colors = [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 0], [1, 0, 1], [0, 1, 1]]
    color_array = [colors[label % num_segments] for label in np.asarray(labels).flatten()]
    pcd.colors = o3d.utility.Vector3dVector(color_array)
    o3d.visualization.draw_geometries([pcd])


    """"# Segment the point cloud into different regions using ECE algorithm
    ece = pcd.cluster_dbscan(eps=1, min_points=5)
    max_label = np.max(ece)
    print('max label')
    print(max_label)
    for i in range(max_label + 1):
        indices = np.where(ece == i)[0]
        segment = pcd.select_by_index(indices)
        labels = np.array(segment.cluster_dbscan(eps=1, min_points=5))
        print("ok")
        print(f"Number of clusters in segment {i}: {len(set(labels))}")
    """
    """"# Assign random colors to each segment
    num_segments = len(set(labels))
    print("Number of segments")
    print(num_segments)
    colors = [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 0], [1, 0, 1], [0, 1, 1]]
    color_array = [colors[label % num_segments] for label in np.asarray(labels).flatten()]
    pcd_filtered.colors = o3d.utility.Vector3dVector(color_array)
    o3d.visualization.draw_geometries([pcd_filtered])"""
    # Assign random colors to each segment
    """num_segments = len(set(labels))
    print("num seg")
    print(num_segments)
    colors = [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 0], [1, 0, 1], [0, 1, 1]]
    color_array = [colors[label % num_segments] for label in np.asarray(labels).flatten()]

    # Assign a random color to each point
    pcd_filtered.colors = o3d.utility.Vector3dVector(np.random.rand(len(points), 3))
    o3d.visualization.draw_geometries([pcd_filtered])"""

    # Extract object positions and orientations
    """for label in set(labels):
        if label == -1:
            continue
        points = pcd.select_down_sample(np.where(labels == label)[0])
        centroid = np.mean(points.points, axis=0)
        print(f"Object {label}: position={centroid}, orientation={points.get_oriented_bounding_box().get_rotation_matrix()}")"""

    # Extract object positions and orientations
    for label in set(labels):
        print("enter1")
        print(set(labels))
        if label == -1:
            print("enter2")
            continue
        print("enter3")
        #for i in range(len(points)):
            #print(points[i])
        points = pcd.points[labels == label]
        print('points')
        print(points)
        cluster_center = np.mean(points, axis=0)
        # Print centroid coordinates
        print(
            f"Object {label}: position={cluster_center}")