diff --git a/CMakeLists.txt b/CMakeLists.txt
index 08d1a45..7c009bf 100644
--- a/CMakeLists.txt
+++ b/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
+)
diff --git a/bt_sim_node.py b/bt_sim_node.py
index 04c6991..6e88648 100755
--- a/bt_sim_node.py
+++ b/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)
diff --git a/msg/AABBox.msg b/msg/AABBox.msg
new file mode 100644
index 0000000..80b0d5b
--- /dev/null
+++ b/msg/AABBox.msg
@@ -0,0 +1,3 @@
+ geometry_msgs/Point min
+ geometry_msgs/Point max
+
\ No newline at end of file
diff --git a/package.xml b/package.xml
index 89dc1ea..6880fa4 100644
--- a/package.xml
+++ b/package.xml
@@ -6,7 +6,12 @@
Michel Breyer
TODO
+
catkin
+
+ message_generation
+
+ message_runtime
+
rospy
- kdl_parser_py
diff --git a/policies.py b/policies.py
index d1d350e..3044538 100644
--- a/policies.py
+++ b/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,
)
diff --git a/simulation.py b/simulation.py
index 6e3c831..676ff8d 100644
--- a/simulation.py
+++ b/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):
diff --git a/srv/Reset.srv b/srv/Reset.srv
new file mode 100644
index 0000000..b1b6924
--- /dev/null
+++ b/srv/Reset.srv
@@ -0,0 +1,2 @@
+---
+active_grasp/AABBox bbox
diff --git a/utils.py b/utils.py
index 7ef2723..319d13e 100644
--- a/utils.py
+++ b/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)