Segmentation mask of soft bodies

[pauljansonnie]

Hi,

I would like to get the segmentation mask of soft bodies in Pinch-v0. I am unable to find any example doing this and since soft bodies do not have any mesh-level or object/link id, it does not seem straight-forward. The fact that goals in Pinch-v0 are provided as point-clouds of the dough (or as images with segmented soft body) makes me believe that this should be feasible. Is it possible to get a segmentation mask of soft bodies? How? If not, how were the goals in Pinch-v0 actually recorded?

Thanks a lot

[xuanlinli17]

You can get the soft body segmentation through the actor id.

obs, _ = env.reset()
env.get_actors() # [Actor(name="ground", id="14"), Actor(name="", id="15")]
soft_body_seg = (obs['image']['hand_camera']['Segmentation'][..., 1] == 15)

[pauljansonnie]

This does not segment the soft-body but just removes the robot from the scene. I actually looped over all segmentation masks already and none of them corresponds to the soft-body.

For instance, based on your code snippet, let us apply the mask to the RGB image and save the masked image:

import matplotlib.pyplot as plt
obs, _ = env.reset()
env.get_actors() # [Actor(name="ground", id="14"), Actor(name="", id="15")]
camera_name = "hand_camera" # or "base_camera"
img_color = (obs['image'][camera_name]["Color"][:, :, :3] * 255).astype(np.uint8)
soft_body_seg = (obs['image'][camera_name]['Segmentation'][..., 1] == 15)
img_masked = img_color * soft_body_seg.reshape(*soft_body_seg.shape, 1)
plt.imsave("img_masked.png", img_masked)

We get the following images:

[xuanlinli17]

Oh my bad, the soft body segmentation isn't officially supported.

However I believe that @hzaskywalker has implemented one, probably he can provide you with a code snippet.