-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathhandle_pointcloud.py
106 lines (97 loc) · 4.44 KB
/
handle_pointcloud.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
from env.move_box_env import MoveBoxEnv
from env.robot_env_cube_move import RobotCubeMove
import robosuite as suite
import numpy as np
import open3d as o3d
import matplotlib.pyplot as plt
import robosuite.utils.camera_utils as camera_utils
import sys
np.set_printoptions(threshold=sys.maxsize)
env = suite.make(
"RobotCubeMove",
robots="Panda",
has_renderer=True,
has_offscreen_renderer=True,
use_camera_obs=True,
camera_segmentations = "element",
camera_depths = True,
# reward_shaping=True,
control_freq=20,
horizon=100,
camera_names = ["birdview", "agentview", "frontview", "sideview"],
camera_heights = 256,
camera_widths = 256
)
def filp_image(image):
return np.flip(image, 0)
if __name__ == "__main__":
obs = env.reset()
print(obs.keys())
# seg_image = obs["agentview_segmentation_element"]
# depth_image = obs['agentview_depth']
# rgb_image = obs['agentview_image']
seg_image = obs["agentview_segmentation_element"]
depth_image = obs['agentview_depth']
depth_image = filp_image(depth_image)
depth_image = camera_utils.get_real_depth_map(env.sim, depth_image)
rgb_image = obs['agentview_image']
# element_id = env.sim.model.geom_name2id('table_visual')
# element_id = env.sim.model.geom_name2id('Door_handle_visual')
print(env.sim.model.geom_names)
# robot_geom_id = [env.sim.model.geom_name2id(robot_geo) for robot_geo in env.sim.model.geom_names if 'robot' in robot_geo or 'gripper' in robot_geo or 'mount' in robot_geo]
pc_geom_id = [env.sim.model.geom_name2id(pc_geo) for pc_geo in env.sim.model.geom_names if 'cube' in pc_geo or 'table' in pc_geo]
print(pc_geom_id)
element_id = env.sim.model.geom_name2id('cube_g0_vis')
masked_segmentation = np.isin(seg_image, pc_geom_id)
masked_segmentation = filp_image(masked_segmentation)
print(masked_segmentation.shape)
# masked_segmentation = np.where(seg_image == int(element_id), 1.0, 0.0)
# masked_segmentation = np.multiply(seg_image, robot_mask)
# print("max depth",np.max(depth_image))
# print(masked_segmentation.shape )
masked_rgb = np.multiply(rgb_image, masked_segmentation)
# to int image
masked_rgb = masked_rgb.astype(np.uint8)
# print(masked_rgb)
# show the masked rgb image
masked_depth = np.multiply(depth_image, masked_segmentation).astype(np.float32)
# masked_depth = np.zeros_like(depth_image,dtype=np.float32)
# masked_depth = masked_depth + 0.5
# for i in range(len(depth_image)):
# for j in range(len(depth_image[0])):
# if masked_segmentation[i][j] == 0:
# depth_image[i][j][0] = 0
# else:
# masked_depth[i][j][0] = depth_image[i][j][0]
# masked_depth[masked_depth > 0.9] = 0.0
print(np.max(masked_depth))
print(np.min(masked_depth))
# normalize to 0-1
# masked_depth = masked_depth - 0.97
# masked_depth = masked_depth / np.max(masked_depth)
# masked_depth = np.clip(masked_depth, 0, 1)
masked_depth_image = o3d.geometry.Image(masked_depth)
intrinisc_cam_parameters_numpy = camera_utils.get_camera_intrinsic_matrix(env.sim, "agentview", 256, 256)
extrinsic_cam_parameters= camera_utils.get_camera_extrinsic_matrix(env.sim, "agentview")
cx = intrinisc_cam_parameters_numpy[0][2]
cy = intrinisc_cam_parameters_numpy[1][2]
fx = intrinisc_cam_parameters_numpy[0][0]
fy = intrinisc_cam_parameters_numpy[1][1]
intrinisc_cam_parameters = o3d.camera.PinholeCameraIntrinsic(256, #width
256, #height
fx,
fy,
cx,
cy)
masked_pcd = o3d.geometry.PointCloud.create_from_depth_image(masked_depth_image,
intrinisc_cam_parameters
)
masked_pcd.transform(extrinsic_cam_parameters)
#estimate normals
masked_pcd.estimate_normals()
#orientation normals to camera
masked_pcd.orient_normals_towards_camera_location(extrinsic_cam_parameters[:3,3])
# o3d.visualization.draw_geometries([masked_pcd])
o3d.io.write_point_cloud('point_clouds/test.pcd', masked_pcd)
plt.imshow(masked_depth)
plt.show()