Merge pull request #105 from Skylark0924/jjl-dev

Support open-vocabulary part-level segmentation by using VLPart and SAM

examples/simulator/example_object_show.py

@@ -9,17 +9,20 @@
 import rofunc as rf
 
 args = rf.config.get_sim_config("Objects")
-asset_files = ["urdf/ycb/002_master_chef_can/002_master_chef_can.urdf",
-               "urdf/ycb/003_cracker_box/003_cracker_box.urdf",
-               "urdf/ycb/004_sugar_box/004_sugar_box.urdf",
-               "urdf/ycb/005_tomato_soup_can/005_tomato_soup_can.urdf",
-               "urdf/ycb/006_mustard_bottle/006_mustard_bottle.urdf",
-               "urdf/ycb/007_tuna_fish_can/007_tuna_fish_can.urdf",
-               "urdf/ycb/008_pudding_box/008_pudding_box.urdf",
-               "urdf/ycb/009_gelatin_box/009_gelatin_box.urdf",
-               "urdf/ycb/010_potted_meat_can/010_potted_meat_can.urdf",
-               "urdf/ycb/011_banana/011_banana.urdf",
-               ]
+asset_files = [
+    # "urdf/ycb/002_master_chef_can/002_master_chef_can.urdf",
+    # "urdf/ycb/003_cracker_box/003_cracker_box.urdf",
+    # "urdf/ycb/004_sugar_box/004_sugar_box.urdf",
+    # "urdf/ycb/005_tomato_soup_can/005_tomato_soup_can.urdf",
+    # "urdf/ycb/006_mustard_bottle/006_mustard_bottle.urdf",
+    # "urdf/ycb/007_tuna_fish_can/007_tuna_fish_can.urdf",
+    # "urdf/ycb/008_pudding_box/008_pudding_box.urdf",
+    # "urdf/ycb/009_gelatin_box/009_gelatin_box.urdf",
+    # "urdf/ycb/010_potted_meat_can/010_potted_meat_can.urdf",
+    # "urdf/ycb/011_banana/011_banana.urdf",
+    # "urdf/ycb/048_hammer/048_hammer.urdf",
+    "urdf/ycb/019_pitcher_base/019_pitcher_base.urdf",
+]
 object_sim = rf.sim.ObjectSim(args, asset_file=asset_files)
 object_sim.create_track_cameras()
-object_sim.show(mode="seg")
+object_sim.show(mode="rgb")

examples/visualab/example_vlpart_sam_seg_w_prompt.py

@@ -0,0 +1,20 @@
+"""
+Part-level segmentation using SAM and VLPart with prompt
+============================================================
+
+This example allows user to provide a text prompt and generate a mask for the corresponding object part.
+"""
+
+import os
+
+import cv2
+
+import rofunc as rf
+
+# obtain rgb and depth image of the grasping scene
+image_path = os.path.join(rf.oslab.get_rofunc_path(), "../examples/data/visualab/p2.png")
+image = cv2.imread(image_path)
+image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+text_prompt = "bottle "
+affordance_masks = rf.visualab.vlpart_sam_predict(image, text_prompt)

requirements.txt

@@ -1,3 +1,5 @@
 https://github.com/Skylark0924/Rofunc/releases/download/v0.0.2.3/pbdlib-0.1-py3-none-any.whl
 https://github.com/Skylark0924/Rofunc/releases/download/v0.0.0.9/isaacgym-1.0rc4-py3-none-any.whl
-git+https://github.com/facebookresearch/segment-anything.git
+git+https://github.com/facebookresearch/segment-anything.git
+git+https://github.com/facebookresearch/detectron2.git
+git+https://github.com/openai/CLIP.git

rofunc/config/simulator/Objects.yaml

@@ -11,7 +11,7 @@ env:
   asset:
     assetFile:
     assetRoot:
-    init_pose: [0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 1.0]
+    init_pose: [0.0, 0.0, 1.1, 0.0, 0.0, 0.0, 1.0]
     armature: 0.01   # The value added to the diagonal elements of inertia tensors for all of the asset’s rigid bodies/links. Could improve simulation stability
     collapse_fixed_joints: # Merge links that are connected by fixed joints.
     convex_decomposition_from_submeshes: False  # Whether to treat submeshes in the mesh as the convex decomposition of the mesh. Default False.

rofunc/simulator/object_sim.py

@@ -203,12 +203,12 @@ def create_track_cameras(self):
         for env_idx in range(self.num_envs):
             env_ptr = self.envs[env_idx]
             camera_handle0 = self.gym.create_camera_sensor(env_ptr, camera_props)
-            self.gym.set_camera_location(camera_handle0, env_ptr, gymapi.Vec3(0.5, -0.8, 1.5), gymapi.Vec3(0, 0, 0))
+            self.gym.set_camera_location(camera_handle0, env_ptr, gymapi.Vec3(0.5, -0.5, 1.3), gymapi.Vec3(0, 0, 0))
 
         for env_idx in range(self.num_envs):
             env_ptr = self.envs[env_idx]
             camera_handle1 = self.gym.create_camera_sensor(env_ptr, camera_props)
-            self.gym.set_camera_location(camera_handle1, env_ptr, gymapi.Vec3(0.5, 0.8, 1.5), gymapi.Vec3(0, 0, 0))
+            self.gym.set_camera_location(camera_handle1, env_ptr, gymapi.Vec3(0.5, 0.5, 1.3), gymapi.Vec3(0, 0, 0))
 
         self.gym.render_all_camera_sensors(self.sim)
         self.visual_obs_flag = True

rofunc/utils/visualab/__init__.py

@@ -1,9 +1,21 @@
 from __future__ import absolute_import
 
+# Visualize package
 from .trajectory import *
 from .ellipsoid import *
 from .distribution import *
 from .utils import *
-from .sam_seg import *
 from .image import *
 from .interact import *
+
+# Segment package
+from .segment import *
+from .segment.sam_seg import *
+from .segment.vlpart_sam_seg import *
+
+# Detect package
+
+# Point cloud package
+
+# SLAM package
+

rofunc/utils/visualab/sam_seg.py → rofunc/utils/visualab/segment/sam_seg.py

@@ -120,26 +120,3 @@ def sam_predict(image,
             show_star_points_w_labels(input_point, input_label, plt.gca())
         plt.axis('off')
         plt.show()
-
-
-if __name__ == '__main__':
-    # image_path = "/home/ubuntu/Documents/tmp/images/dog.jpg"
-    # image = cv2.imread(image_path)
-    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    # plt.figure(figsize=(10, 10))
-    # plt.imshow(image)
-    # plt.axis('off')
-    # plt.show()
-    # sam_generate(image)
-
-    image_path = "/home/ubuntu/Documents/tmp/images/truck.jpg"
-    image = cv2.imread(image_path)
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    plt.figure(figsize=(10, 10))
-    plt.imshow(image)
-    plt.axis('off')
-    plt.show()
-    sam_predict(image,
-                use_point=True,
-                use_box=False,
-                choose_best_mask=True)