update

src/active_grasp/active_perception/configs/local_inference_config.yaml

@@ -12,7 +12,7 @@ settings:
   experiment:
     name: test_inference
     root_dir: "experiments"
-    model_path: "/media/hofee/data/weights/nbv_grasping/grasping_full_centralized_ood_111.pth"
+    model_path: "/media/hofee/data/weights/nbv_grasping/241204_centralized_ood.pth"
     use_cache: True
     small_batch_overfit: False
     

src/active_grasp/active_perception_policy.py

@@ -187,16 +187,33 @@ class ActivePerceptionSingleViewPolicy(SingleViewPolicy):
                         gsnet_input_points = PtsUtil.random_downsample_point_cloud(gsnet_input_points, required_num_points)
            
             self.publish_pointcloud(gsnet_input_points)
-    
+            
             gsnet_grasping_poses = np.asarray(self.request_grasping_pose(gsnet_input_points.tolist()))
             print(gsnet_grasping_poses[0].keys())
-            #import ipdb; ipdb.set_trace()
 
-            # DEBUG: publish grasps
+            # ----------------------------------------------------------------
+            # FIXME: 
             # self.publish_grasps(gsnet_grasping_poses)
 
             # Convert all grasping poses' reference frame to arm frame
             current_cam_pose = torch.from_numpy(x.as_matrix()).float().to("cuda:0")
+            
+            # 将点云转换为齐次坐标
+            points_homogeneous = np.ones((gsnet_input_points.shape[0], 4))
+            points_homogeneous[:, 0:3] = gsnet_input_points
+            
+            # 使用相机位姿矩阵转换到世界坐标系
+            world_points = (current_cam_pose.cpu().numpy() @ points_homogeneous.T).T
+            
+            # 只保留xyz坐标
+            world_points = world_points[:, 0:3]
+            
+            # 保存为txt文件
+            np.savetxt('gsnet_input_world_points.txt', world_points)
+            
+            import ipdb; ipdb.set_trace()
+            # ----------------------------------------------------------------
+
             for gg in gsnet_grasping_poses:
                 gg['T'] = current_cam_pose.cpu().numpy().dot(np.asarray(gg['T']))
                 R = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
@@ -211,10 +228,17 @@ class ActivePerceptionSingleViewPolicy(SingleViewPolicy):
                 grasp = ParallelJawGrasp(T, width)
                 grasps.append(grasp)
                 qualities.append(gg['score'])
-            
+            # ----------------------------------------------------------------
             # Visualize grasps
-            # self.vis.grasps(self.base_frame, grasps, qualities)
-            
+            import ipdb; ipdb.set_trace()
+            grasp_points = []
+            for grasp, quality in zip(grasps, qualities):
+                pose = grasp.pose
+                # tip = pose.rotation.apply([0, 0, 0.05]) + pose.translation
+                tip = pose.translation
+                grasp_points.append(tip)
+            np.savetxt('grasp_points.txt', grasp_points)
+            # ----------------------------------------------------------------
             # Filter grasps
             filtered_grasps = []
             filtered_qualities = []

src/active_grasp/controller.py

@@ -17,6 +17,9 @@ from robot_helpers.ros.moveit import MoveItClient, create_collision_object_from_
 from robot_helpers.spatial import Rotation, Transform
 from vgn.utils import look_at, cartesian_to_spherical, spherical_to_cartesian
 
+RED = '\033[91m'
+RESET = '\033[0m'
+
 
 class GraspController:
     def __init__(self, policy):
@@ -87,7 +90,9 @@ class GraspController:
         self.latest_seg_msg = msg
 
     def run(self):
+        print("before reset")
         bbox = self.reset()
+        print("after reset")
         self.switch_to_cartesian_velocity_control()
         with Timer("search_time"):
             grasp = self.search_grasp(bbox)
@@ -97,6 +102,7 @@ class GraspController:
                 res = self.execute_grasp(grasp)
         else:
             res = "aborted"
+            print(f"{RED}[{res}]No grasp found{RESET}")
         return self.collect_info(res)
 
     def reset(self):
@@ -125,7 +131,7 @@ class GraspController:
                     current_p = pose.as_matrix()[:3,3]
                     target_p = self.policy.x_d.as_matrix()[:3,3]
                     linear_d = np.linalg.norm(current_p - target_p)
-                    if(linear_d < self.move_to_target_threshold):
+                    if(linear_d < self.move_to_target_threshold*2):
                         # Arrived
                         moving_to_The_target = False
                 # sleep 3s to wait for the arrival of the robot
@@ -135,15 +141,19 @@ class GraspController:
         elif(self.policy.policy_type=="multi_view"):
             while not self.policy.done:
                 depth_img, seg_image, pose, q = self.get_state()
+                #print("before get target id in search_grasp()")
+                #import ipdb; ipdb.set_trace()
                 target_seg_id = self.get_target_id(TargetIDRequest()).id
                 support_seg_id = self.get_support_id(TargetIDRequest()).id
+                #print("after get target id in search_grasp()")
+
                 # print(target_seg_id, support_seg_id)
                 # import ipdb; ipdb.set_trace()
                 self.policy.update(depth_img, seg_image, target_seg_id, pose, q)
                 r.sleep()
         else:
             print("Unsupported policy type: "+str(self.policy.policy_type))
-
+        import ipdb; ipdb.set_trace()
         # Wait for a zero command to be sent to the robot.
         rospy.sleep(0.2)
         self.policy.deactivate()

src/active_grasp/nbv.py

@@ -85,6 +85,7 @@ class NextBestView(MultiViewPolicy):
         super().activate(bbox, view_sphere)
 
     def update(self, img, seg, target_id, x, q):
+        #import ipdb; ipdb.set_trace()
         if len(self.views) > self.max_views or self.best_grasp_prediction_is_stable():
             self.done = True
         else:
@@ -100,7 +101,7 @@ class NextBestView(MultiViewPolicy):
             self.vis.ig_views(self.base_frame, self.intrinsic, views, utilities)
             i = np.argmax(utilities)
             nbv, gain = views[i], gains[i]
-
+            import ipdb; ipdb.set_trace()
             if gain < self.min_gain and len(self.views) > self.T:
                 self.done = True
 
@@ -174,6 +175,7 @@ class NextBestView(MultiViewPolicy):
         bbox_min = self.T_task_base.apply(self.bbox.min) / voxel_size
         bbox_max = self.T_task_base.apply(self.bbox.max) / voxel_size
         mask = np.array([((i > bbox_min) & (i < bbox_max)).all() for i in indices])
+        #import ipdb; ipdb.set_trace()
         i, j, k = indices[mask].T
         tsdfs = tsdf_grid[i, j, k]
         ig = np.logical_and(tsdfs > -1.0, tsdfs < 0.0).sum()

src/active_grasp/simulation.py

@@ -2,6 +2,8 @@ from pathlib import Path
 import pybullet as p
 import pybullet_data
 import rospkg
+import pickle
+import os
 
 from active_grasp.bbox import AABBox
 from robot_helpers.bullet import *
@@ -29,6 +31,7 @@ class Simulation:
         self.load_robot()
         self.load_vgn(Path(vgn_path))
         self.scene = get_scene(scene_id)
+        self.target_uid = None
 
     def configure_physics_engine(self, gui, rate, sub_step_count):
         self.rate = rate
@@ -48,6 +51,16 @@ class Simulation:
     def load_robot(self):
         panda_urdf_path = urdfs_dir / "franka/panda_arm_hand.urdf"
         self.arm = BtPandaArm(panda_urdf_path)
+        print(f"Robot base uid: {self.arm.uid}")
+        
+        # 打印机械臂的所有部件信息
+        num_joints = p.getNumJoints(self.arm.uid)
+        print(f"Number of joints: {num_joints}")
+        for i in range(num_joints):
+            joint_info = p.getJointInfo(self.arm.uid, i)
+            link_name = joint_info[12].decode('utf-8')
+            print(f"Joint {i}: {link_name}")
+        
         self.gripper = BtPandaGripper(self.arm)
         self.model = KDLModel.from_urdf_file(
             panda_urdf_path, self.arm.base_frame, self.arm.ee_frame
@@ -63,7 +76,7 @@ class Simulation:
         sleep_ticks = self.rate * secs
         for i in range(int(sleep_ticks)):
             p.stepSimulation()
-        
+        self.target_uid = None
         # Reset the scene
         valid = False
         while not valid:
@@ -72,13 +85,17 @@ class Simulation:
             q = self.scene.generate(self.rng)
             self.set_arm_configuration(q)
             uid = self.select_target()
-            support_id = self.select_support()
+            if uid == -1:
+                continue
+            support_id = self.select_support()  
             bbox = self.get_target_bbox(uid)
+            #import ipdb; ipdb.set_trace()
             valid = True
             # valid = self.check_for_grasps(bbox)
         return bbox
 
     def set_arm_configuration(self, q):
+        #import ipdb; ipdb.set_trace()
         for i, q_i in enumerate(q):
             p.resetJointState(self.arm.uid, i, q_i, 0)
         p.resetJointState(self.arm.uid, 9, 0.04, 0)
@@ -86,15 +103,25 @@ class Simulation:
         self.gripper.set_desired_width(0.4)
 
     def select_target(self):
+        if self.target_uid is not None:
+            p.changeVisualShape(self.target_uid, -1, rgbaColor=[1, 0, 0, 1])
+            return self.target_uid
         _, _, mask = self.camera.get_image()
         uids, counts = np.unique(mask, return_counts=True)
+        
         mask = np.isin(uids, self.scene.object_uids)  # remove ids of the floor, etc
         uids, counts = uids[mask], counts[mask]
         # If no objects are detected, try again
-        if(len(uids) == 0 and len(counts) == 0):
-            return self.select_target()
-        target_uid = uids[np.argmin(counts)]
+        container_idx = np.where(np.isin(uids, self.scene.container_uid))[0]
+        object_uids = np.delete(uids, container_idx)
+        object_counts = np.delete(counts, container_idx)
+        target_uid = object_uids[np.argmin(object_counts)]
+        if(len(object_uids) == 0 and len(object_counts) == 0):
+            return -1
+        
+        print("select target uid: "+str(target_uid))
         p.changeVisualShape(target_uid, -1, rgbaColor=[1, 0, 0, 1])
+        self.target_uid = target_uid
         return target_uid
     
     def select_support(self):
@@ -104,6 +131,7 @@ class Simulation:
     def get_target_bbox(self, uid):
         aabb_min, aabb_max = p.getAABB(uid)
         # enlarge the bounding box
+        #import ipdb; ipdb.set_trace()
         aabb_min = np.asarray(aabb_min)
         aabb_max = np.asarray(aabb_max)
         aabb_min -= 0.0
@@ -150,6 +178,16 @@ class Scene:
         self.support_urdf = urdfs_dir / "plane/model.urdf"
         self.support_uid = -1
         self.object_uids = []
+        self.container_uid = -1
+        self.container_list = [
+                                "chair",
+                                "table",
+                                "sofa",
+                                "picnic_basket",
+                                "cabinet",
+                                "flower_pot",
+                                "blender_cabinet"
+                            ]
 
     def clear(self):
         self.remove_support()
@@ -167,6 +205,9 @@ class Scene:
 
     def add_object(self, urdf, ori, pos, scale=1.0):
         uid = p.loadURDF(str(urdf), pos, ori.as_quat(), globalScaling=scale)
+        print('add object uid: '+str(uid))
+        import ipdb; ipdb.set_trace()
+        
         self.object_uids.append(uid)
         return uid
 
@@ -175,13 +216,18 @@ class Scene:
         self.object_uids.remove(uid)
 
     def remove_all_objects(self):
+        print(f"Objects to remove: {self.object_uids}")
         for uid in list(self.object_uids):
+            # 获取物体信息
+            body_info = p.getBodyInfo(uid)
+            print(f'Removing object uid: {uid}, Name: {body_info[1].decode("utf-8")}')
             self.remove_object(uid)
 
 
 class YamlScene(Scene):
     def __init__(self, config_name):
         super().__init__()
+        #import ipdb; ipdb.set_trace()
         self.config_path = pkg_root / "cfg/sim" / config_name
 
     def load_config(self):
@@ -208,93 +254,111 @@ class YamlScene(Scene):
             p.stepSimulation()
         return self.scene["q"]
 
-
-class RandomScene(Scene):
-    def __init__(self):
+class PickleScene(Scene):
+    def __init__(self, config_name):
         super().__init__()
-        self.center = np.r_[0.5, 0.0, 0.2]
-        self.length = 0.3
-        self.origin = self.center - np.r_[0.5 * self.length, 0.5 * self.length, 0.0]
-        self.object_urdfs = find_urdfs(urdfs_dir / "test")
-
-    def generate(self, rng, object_count=4, attempts=10):
-        self.add_support(self.center)
-        urdfs = rng.choice(self.object_urdfs, object_count)
-        for urdf in urdfs:
-            scale = rng.uniform(0.8, 1.0)
-            uid = self.add_object(urdf, Rotation.identity(), np.zeros(3), scale)
-            lower, upper = p.getAABB(uid)
-            z_offset = 0.5 * (upper[2] - lower[2]) + 0.002
-            state_id = p.saveState()
-            for _ in range(attempts):
-                # Try to place and check for collisions
-                ori = Rotation.from_euler("z", rng.uniform(0, 2 * np.pi))
-                pos = np.r_[rng.uniform(0.2, 0.8, 2) * self.length, z_offset]
-                p.resetBasePositionAndOrientation(uid, self.origin + pos, ori.as_quat())
-                p.stepSimulation()
-                if not p.getContactPoints(uid):
-                    break
-                else:
-                    p.restoreState(stateId=state_id)
-            else:
-                # No placement found, remove the object
-                self.remove_object(uid)
-        q = [0.0, -1.39, 0.0, -2.36, 0.0, 1.57, 0.79]
-        q += rng.uniform(-0.08, 0.08, 7)
-        return q
-
-
-class ManualScene(Scene):
-    def __init__(self):
-        super().__init__()
-        self.config_path = pkg_root / "cfg/sim"
-        self.scene_index = 0
-        # Visit the directory and read all the yaml files
-        self.scenes = []
-        for file in self.config_path.iterdir():
-            if file.suffix == ".yaml":
-                self.scenes.append(file)
-        self.num_scenes = len(self.scenes)
-    
-    def load_config(self):
-        self.scene = load_yaml(self.scenes[self.scene_index])
-        self.center = np.asarray(self.scene["center"])
+        self.config_path = pkg_root / "cfg/sim" / config_name
+        self.center = np.r_[0.5, 0.0, 0.0]
         self.length = 0.3
         self.origin = self.center - np.r_[0.5 * self.length, 0.5 * self.length, 0.0]
         
 
+    def load_scene(self, scene_dir, model_dir, rot_z=0):
+        scene_path = os.path.join(scene_dir, "scene.pickle")
+        self.scene = pickle.load(open(scene_path, "rb"))
+        self.fall_objects_dict = self.get_fall_objects_dict(scene_dir)
+        self._initialize_pybullet_scene(model_dir, rot_z)
+
+    def _initialize_pybullet_scene(self, model_dir, rot_z=0):
+        p.setAdditionalSearchPath(pybullet_data.getDataPath())
+        # 修改平面加载,添加碰撞属性
+        planeId = p.loadURDF("plane100.urdf", useMaximalCoordinates=True, flags=p.URDF_USE_IMPLICIT_CYLINDER)
+        p.changeDynamics(planeId, -1, 
+            lateralFriction=1.0,
+            restitution=0.5,
+            contactStiffness=1000000,
+            contactDamping=100
+        )
+        
+        # 创建旋转矩阵
+        rot_matrix = Rotation.from_euler("z", rot_z, degrees=True)
+        
+        for obj_name in self.scene.keys():
+            if obj_name in self.fall_objects_dict:
+                continue
+            # 处理旋转
+            origin_orientation = self.scene[obj_name]["rotation"] # quaternion
+            origin_orientation = Rotation.from_quat(origin_orientation)
+            orientation = rot_matrix * origin_orientation
+            orientation = orientation.as_quat()
+            
+            # 处理位置
+            original_position = self.scene[obj_name]["position"]
+            # 将位置向量绕z轴旋转相同的角度
+            rotated_position = rot_matrix.apply(original_position)
+            position = rotated_position + self.center
+            
+            class_name = obj_name[:-4]
+            #obj_path = os.path.join(model_dir, class_name, obj_name, "model.obj")
+            
+            urdf_path = os.path.join(model_dir, class_name, obj_name, "model.urdf")
+            #import ipdb; ipdb.set_trace()
+            uid = p.loadURDF(urdf_path, position, orientation, globalScaling=1.0)
+            # uid = self._load_obj_to_pybullet(
+            #     obj_file_path=obj_path,
+            #     position=position,
+            #     orientation=orientation,
+            #     scale=1.0,
+            # )
+            
+            self.object_uids.append(uid)
+            for container in self.container_list:
+                if container in obj_name:
+                    self.container_uid = uid
+                    print('container id: '+str(self.container_uid))
+
+    def get_fall_objects_dict(self, scene_path):
+        fall_objects_path = os.path.join(scene_path, "fall_objects.pickle")
+        fall_objects_dict = pickle.load(open(fall_objects_path, "rb"))
+        return fall_objects_dict
+    
+    def _load_obj_to_pybullet(self, obj_file_path, position, orientation, scale):
+        visual_ind = p.createVisualShape(
+            shapeType=p.GEOM_MESH,
+            fileName=obj_file_path,
+            rgbaColor=[1, 1, 1, 1],
+            specularColor=[0.4, 0.4, 0],
+            visualFramePosition=[0, 0, 0],
+            meshScale=scale,
+        )
+        collision_ind = p.createCollisionShape(
+            shapeType=p.GEOM_MESH,
+            fileName=obj_file_path,
+            meshScale=scale,
+        )
+        uid = p.createMultiBody(
+            baseMass=0,
+            baseVisualShapeIndex=visual_ind,
+            baseCollisionShapeIndex=collision_ind,
+            basePosition=position,
+            baseOrientation=orientation,
+            useMaximalCoordinates=True,
+        )
+        return uid
+
     def generate(self, rng):
-        self.load_config()
-        self.add_support(self.center)
-
-        for object in self.scene["objects"]:
-            urdf = urdfs_dir / object["object_id"] / "model.urdf"
-            ori = Rotation.from_euler("xyz", object["rpy"], degrees=True)
-            pos = self.center + np.asarray(object["xyz"])
-            scale = object.get("scale", 1)
-            if randomize := object.get("randomize", False):
-                angle = rng.uniform(-randomize["rot"], randomize["rot"])
-                ori = Rotation.from_euler("z", angle, degrees=True) * ori
-                b = np.asarray(randomize["pos"])
-                pos += rng.uniform(-b, b)
-            self.add_object(urdf, ori, pos, scale)
-        for _ in range(60):
-            p.stepSimulation()
-        
-        self.scene_index += 1
-        if(self.scene_index >= self.num_scenes):
-            self.scene_index = 0
-
-        return self.scene["q"]
-
+        model_dir = "/media/hofee/data/data/nbv/nbv_objects_urdf"
+        self.load_scene(self.config_path, model_dir, rot_z=90)
+        self.add_support(self.center - np.array([0.0, 0.0, 0.05]))
+        # for _ in range(60):
+        #     p.stepSimulation()
+            
+        q = np.asarray([0.0, -1.39, 0.0, -2.36, 0.0, 1.57, 0.79])
+        q += np.asarray(rng.uniform(-0.08, 0.08, 7))
+        return q.tolist()
 
 def get_scene(scene_id):
     if scene_id.endswith(".yaml"):
         return YamlScene(scene_id)
-    elif scene_id == "random":
-        return RandomScene()
-    elif scene_id == "manual":
-        return ManualScene()
-
     else:
-        raise ValueError("Unknown scene {}.".format(scene_id))
+        return PickleScene(scene_id)