segmentation2

import pybullet as p
import pybullet_data
import open3d as o3d
import numpy as np

# Connect to PyBullet and set up the simulation environment
p.connect(p.GUI)
p.setAdditionalSearchPath(pybullet_data.getDataPath())

# Load a plane and two cubes onto it
plane = p.loadURDF("plane.urdf")
cube1 = p.loadURDF("cube.urdf", [0, 0, 0.5])
cube2 = p.loadURDF("cube.urdf", [1.5, 0, 0.5])
#cube3 = p.loadURDF("cube.urdf", [3, 0, 0.5])

# Step the simulation to update the positions of the objects
p.stepSimulation()

# Get the positions of the cubes
cube1_pos, _ = p.getBasePositionAndOrientation(cube1)
cube2_pos, _ = p.getBasePositionAndOrientation(cube2)
#cube3_pos, _ = p.getBasePositionAndOrientation(cube3)


# Create a point cloud of the plane and the two cubes
points = []
for x in np.linspace(-5, 5, 100):
    for y in np.linspace(-5, 5, 100):
        z = 0
        points.append([x, y, z])

for cube_pos in [cube1_pos, cube2_pos]:
    for x in np.linspace(-0.5, 0.5, 10):
        for y in np.linspace(-0.5, 0.5, 10):
            for z in np.linspace(-0.5, 0.5, 10):

                points.append([x + cube_pos[0], y + cube_pos[1], z + 1 + cube_pos[2]])

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

# Preprocess the data
pcd_filtered = pcd.voxel_down_sample(voxel_size=0.1)

# Apply segmentation algorithm
labels = pcd_filtered.cluster_dbscan(eps=0.5, min_points=10)


""""# 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))}")
"""
# Visualize the segmented point cloud
#colors = [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 0], [1, 0, 1], [0, 1, 1]]
#color_array = [colors[label % len(colors)] for label in 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_filtered.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")
    if label == -1:
        print("enter2")
        continue
    print("enter3")
    indices = np.where(labels == label)[0]
    print("indices")
    print(indices)
    points = pcd_filtered.select_by_index(indices)
    print('points')
    print(points)
    centroid = np.mean(points.points, axis=0)
    print(
        f"Object {label}: position={centroid}")
    #orientation={points.get_oriented_bounding_box().get_rotation_matrix()}"