Select target

CMakeLists.txt

@@ -1,4 +1,22 @@
 cmake_minimum_required(VERSION 3.1)
 project(active_grasp)
 
-find_package(catkin REQUIRED)
+find_package(catkin REQUIRED COMPONENTS
+  geometry_msgs
+  message_generation
+)
+
+add_message_files(
+  FILES
+  AABBox.msg
+)
+
+add_service_files(
+  FILES
+  Reset.srv
+)
+
+generate_messages(
+  DEPENDENCIES
+  geometry_msgs
+)

bt_sim_node.py

@@ -11,6 +11,7 @@ from sensor_msgs.msg import JointState, Image, CameraInfo
 import std_srvs.srv as std_srvs
 import tf2_ros
 
+import active_grasp.srv
 from robot_utils.ros.conversions import *
 from simulation import Simulation
 
@@ -29,7 +30,7 @@ class BtSimNode:
         self.broadcast_transforms()
 
     def advertise_services(self):
-        rospy.Service("reset", std_srvs.Trigger, self.reset)
+        rospy.Service("reset", active_grasp.srv.Reset, self.reset)
 
     def broadcast_transforms(self):
         self.static_broadcaster = tf2_ros.StaticTransformBroadcaster()
@@ -44,10 +45,11 @@ class BtSimNode:
     def reset(self, req):
         self.reset_requested = True
         rospy.sleep(1.0)  # wait for the latest sim step to finish
-        self.sim.reset()
+        bbox = self.sim.reset()
+        res = active_grasp.srv.ResetResponse(bbox.to_msg())
         self.step_cnt = 0
         self.reset_requested = False
-        return std_srvs.TriggerResponse(success=True)
+        return res
 
     def run(self):
         rate = rospy.Rate(self.sim.rate)
@@ -158,8 +160,8 @@ class CameraInterface:
         stamp = rospy.Time.now()
         self.cam_info_msg.header.stamp = stamp
         self.cam_info_pub.publish(self.cam_info_msg)
-        _, depth = self.camera.get_image()
-        depth_msg = self.cv_bridge.cv2_to_imgmsg(depth)
+        img = self.camera.get_image()
+        depth_msg = self.cv_bridge.cv2_to_imgmsg(img.depth)
         depth_msg.header.stamp = stamp
         self.depth_pub.publish(depth_msg)
 

msg/AABBox.msg

@@ -0,0 +1,3 @@
+  geometry_msgs/Point min
+  geometry_msgs/Point max
+  

package.xml

@@ -6,7 +6,12 @@
     <maintainer email="mbreyer@ethz.ch">Michel Breyer</maintainer>
     <license>TODO</license>
     <url></url>
+
     <buildtool_depend>catkin</buildtool_depend>
+
+    <build_depend>message_generation</build_depend>
+
+    <exec_depend>message_runtime</exec_depend>
+
     <depend>rospy</depend>
-    <depend>kdl_parser_py</depend>
 </package>

policies.py

@@ -2,9 +2,10 @@ import cv_bridge
 import numpy as np
 from pathlib import Path
 import rospy
-from sensor_msgs.msg import Image, CameraInfo
+import sensor_msgs.msg
 from visualization_msgs.msg import Marker, MarkerArray
 
+from robot_utils.perception import Image
 from robot_utils.ros import tf
 from robot_utils.ros.conversions import *
 from robot_utils.ros.rviz import *
@@ -45,12 +46,16 @@ class BasePolicy:
         self.T_B_task = tf.lookup(self.base_frame, self.task_frame)
 
     def connect_to_camera(self):
-        msg = rospy.wait_for_message(self.info_topic, CameraInfo, rospy.Duration(2.0))
+        msg = rospy.wait_for_message(
+            self.info_topic, sensor_msgs.msg.CameraInfo, rospy.Duration(2.0)
+        )
         self.intrinsic = from_camera_info_msg(msg)
-        rospy.Subscriber(self.depth_topic, Image, self.sensor_cb, queue_size=1)
+        rospy.Subscriber(
+            self.depth_topic, sensor_msgs.msg.Image, self.sensor_cb, queue_size=1
+        )
 
     def sensor_cb(self, msg):
-        self.depth_img = self.cv_bridge.imgmsg_to_cv2(msg).astype(np.float32)
+        self.img = Image(depth=self.cv_bridge.imgmsg_to_cv2(msg).astype(np.float32))
         self.extrinsic = tf.lookup(
             self.cam_frame,
             self.task_frame,
@@ -72,7 +77,7 @@ class BasePolicy:
     def integrate_latest_image(self):
         self.viewpoints.append(self.extrinsic.inv())
         self.tsdf.integrate(
-            self.depth_img,
+            self.img,
             self.intrinsic,
             self.extrinsic,
         )

simulation.py

@@ -5,6 +5,7 @@ import rospkg
 from robot_utils.bullet import *
 from robot_utils.controllers import CartesianPoseController
 from robot_utils.spatial import Rotation, Transform
+from utils import AABBox
 
 
 class Simulation(BtSim):
@@ -47,9 +48,9 @@ class Simulation(BtSim):
     def reset(self):
         self.remove_all_objects()
         self.set_initial_arm_configuration()
-        urdfs = np.random.choice(self.urdfs, 4)
-        origin = np.r_[self.origin[:2], 0.625]
-        self.random_object_arrangement(urdfs, origin, self.length, 0.8)
+        self.load_random_object_arrangement()
+        uid = self.select_target()
+        return self.get_target_bbox(uid)
 
     def set_initial_arm_configuration(self):
         q = self.arm.configurations["ready"]
@@ -75,7 +76,10 @@ class Simulation(BtSim):
         for uid in list(self.object_uids):
             self.remove_object(uid)
 
-    def random_object_arrangement(self, urdfs, origin, length, scale=1.0, attempts=10):
+    def load_random_object_arrangement(self, attempts=10):
+        origin = np.r_[self.origin[:2], 0.625]
+        scale = 0.8
+        urdfs = np.random.choice(self.urdfs, 4)
         for urdf in urdfs:
             # Load the object
             uid = self.load_object(urdf, Rotation.identity(), [0.0, 0.0, 0.0], scale)
@@ -85,7 +89,7 @@ class Simulation(BtSim):
             for _ in range(attempts):
                 # Try to place the object without collision
                 ori = Rotation.from_rotvec([0.0, 0.0, np.random.uniform(0, 2 * np.pi)])
-                offset = np.r_[np.random.uniform(0.2, 0.8, 2) * length, z_offset]
+                offset = np.r_[np.random.uniform(0.2, 0.8, 2) * self.length, z_offset]
                 p.resetBasePositionAndOrientation(uid, origin + offset, ori.as_quat())
                 self.step()
                 if not p.getContactPoints(uid):
@@ -96,6 +100,17 @@ class Simulation(BtSim):
                 # No placement found, remove the object
                 self.remove_object(uid)
 
+    def select_target(self):
+        img = self.camera.get_image()
+        uids, counts = np.unique(img.mask, return_counts=True)
+        target_uid = uids[np.argmin(counts)]
+        p.changeVisualShape(target_uid, -1, rgbaColor=[1, 0, 0, 1])
+        return target_uid
+
+    def get_target_bbox(self, uid):
+        aabb_min, aabb_max = p.getAABB(uid)
+        return AABBox(aabb_min, aabb_max)
+
 
 class CartesianPoseController:
     def __init__(self, model, frame, x0):

srv/Reset.srv

@@ -0,0 +1,2 @@
+---
+active_grasp/AABBox bbox

utils.py

@@ -1,6 +1,7 @@
 from geometry_msgs.msg import PoseStamped
 import rospy
 
+import active_grasp.msg
 from robot_utils.ros.conversions import *
 
 
@@ -11,3 +12,24 @@ class CartesianPoseControllerClient:
     def send_target(self, pose):
         msg = to_pose_stamped_msg(pose, "panda_link0")
         self.target_pub.publish(msg)
+
+
+class AABBox:
+    def __init__(self, bbox_min, bbox_max):
+        self.min = bbox_min
+        self.max = bbox_max
+
+    @classmethod
+    def from_msg(cls, msg):
+        aabb_min = from_point_msg(msg.min)
+        aabb_max = from_point_msg(msg.max)
+        return cls(aabb_min, aabb_max)
+
+    def to_msg(self):
+        msg = active_grasp.msg.AABBox()
+        msg.min = to_point_msg(self.min)
+        msg.max = to_point_msg(self.max)
+        return msg
+
+    def is_inside(self, p):
+        return np.all(p > self.min) and np.all(p < self.max)