]> defiant.homedns.org Git - wt_open_manipulator.git/blob - scripts/grab_bottle.py
7744d7ad98d311e734016fdbf7393d1927e53035
[wt_open_manipulator.git] / scripts / grab_bottle.py
1 #!/usr/bin/env python
2 # -*- coding: iso-8859-15 -*-
3
4 """
5 roslaunch darknet_ros darknet_ros.launch image:=/camera/color/image_raw
6 """
7
8 import sys
9 import rospy
10 import tf
11 import tf2_ros
12 import tf2_geometry_msgs
13 import image_geometry
14 import moveit_commander
15 import numpy as np
16 from math import *
17 from time import sleep
18 from sensor_msgs.msg import Image, CameraInfo
19 from darknet_ros_msgs.msg import BoundingBoxes
20 from geometry_msgs.msg import PoseStamped
21 from cv_bridge import CvBridge
22
23 class OpenManipGrabBottle:
24         def __init__(self):
25                 rospy.init_node('open_manip_grab_bottle')
26                 moveit_commander.roscpp_initialize(sys.argv)
27                 self.sub_depthinfo = rospy.Subscriber("/camera/aligned_depth_to_color/camera_info", CameraInfo, self.depthinfo_received)
28                 self.tf_broadcaster = tf.broadcaster.TransformBroadcaster()
29                 self.depthimgmsg = None
30                 self.camera_model = None
31                 self.bridge = CvBridge()
32                 self.tfBuffer = tf2_ros.Buffer()
33                 self.listener = tf2_ros.TransformListener(self.tfBuffer)
34                 self.target_class = "bottle"
35
36                 self.gripper_pose_pub = rospy.Publisher('gripper_pose', PoseStamped)
37                 self.arm        = moveit_commander.MoveGroupCommander("arm")
38                 self.gripper    = moveit_commander.MoveGroupCommander("gripper")
39                 self.arm.set_goal_position_tolerance(0.01)
40                 self.arm.set_goal_orientation_tolerance(5*pi/180)
41                 self.home()
42                 sleep(10) # wait for image
43                 self.register_subscriber()
44                 rospy.spin()
45
46         def register_subscriber(self):
47                 self.sub_obj_detection = rospy.Subscriber("/darknet_ros/bounding_boxes", BoundingBoxes, self.detected_objects)
48                 self.sub_depthimage = rospy.Subscriber("/camera/aligned_depth_to_color/image_raw", Image, self.depthimage_received)
49
50         def home(self):
51                 self.gripper.set_named_target('close')
52                 self.gripper.go()
53                 #self.arm.set_named_target('home')
54                 self.arm.go([f * pi/180 for f in [90.0, -90, 77, 20]], wait=True)
55
56         def detected_objects(self, msg):
57                 if self.depthimgmsg is None or self.camera_model is None:
58                         return
59                 box_bottles = filter(lambda box: box.Class == self.target_class and box.probability > 0.5, msg.bounding_boxes)
60                 if len(box_bottles) > 0:
61                         box = box_bottles[0] # just use the first one found for now
62                         box_x = int(box.xmin + (box.xmax - box.xmin)/2.0)
63                         box_y = int(box.ymin + (box.ymax - box.ymin)/2.0)
64                         depthframe = self.bridge.imgmsg_to_cv2(self.depthimgmsg)
65                         dist = depthframe[box_y][box_x]/1000.0 # assume size rgb = size depth
66                         if dist == 0:
67                                 rospy.logerr("depth invalid")
68                                 return
69
70                         self.sub_obj_detection.unregister()
71                         self.sub_depthimage.unregister()
72
73                         coords3d = [f * dist for f in self.camera_model.projectPixelTo3dRay([box_x, box_y])]
74                         point = PoseStamped()
75                         point.header.frame_id = self.depthimgmsg.header.frame_id
76                         point.header.stamp = rospy.Time.now()
77                         point.pose.position.x = coords3d[0]
78                         point.pose.position.y = coords3d[1]
79                         point.pose.position.z = coords3d[2]
80                         point = self.tfBuffer.transform(point, "open_manipulator_base", rospy.Duration(1))
81                         point.pose.orientation.x, point.pose.orientation.y, point.pose.orientation.z, point.pose.orientation.w = [0, 0, 0, 1] # bottle stands up
82                         print point
83                         point.pose.position.z-=0.02
84
85                         self.tf_broadcaster.sendTransform([point.pose.position.x, point.pose.position.y, point.pose.position.z], [point.pose.orientation.x, point.pose.orientation.y, point.pose.orientation.z, point.pose.orientation.w], rospy.Time.now(), self.target_class, "open_manipulator_base")
86
87                         if self.pick(point):
88                                 rospy.signal_shutdown("done")
89                         else:
90                                 self.home()
91                                 self.register_subscriber()
92
93         def depthimage_received(self, msg):
94                 self.depthimgmsg = msg
95
96         def depthinfo_received(self, msg):
97                 self.camera_model = image_geometry.PinholeCameraModel()
98                 self.camera_model.fromCameraInfo(msg)
99                 self.sub_depthinfo.unregister()
100
101         def pick(self, pose):
102                 pos2d = [pose.pose.position.x, pose.pose.position.y]
103                 dist = np.linalg.norm(pos2d)
104                 pos2d_unit = pos2d / dist # unit vector
105
106                 # orientation towards base
107                 q = tf.transformations.quaternion_from_euler(0, 0, acos(np.dot(pos2d_unit, [1, 0]))) # use the dot product to get the angle
108                 pose.pose.orientation.x, pose.pose.orientation.y, pose.pose.orientation.z, pose.pose.orientation.w = [q[0], q[1], q[2], q[3]]
109
110                 # goal 5cm before
111                 pose.pose.position.x = pos2d_unit[0] * (dist-0.05)
112                 pose.pose.position.y = pos2d_unit[1] * (dist-0.05)
113                 self.arm.set_pose_target(pose)
114                 if not self.arm.go():
115                         print "before position failed"
116                         return False
117
118                 self.gripper.set_named_target('open')
119                 self.gripper.go()
120
121                 # goal on target +3cm
122                 pose.pose.position.x = pos2d_unit[0] * (dist+0.03)
123                 pose.pose.position.y = pos2d_unit[1] * (dist+0.03)
124                 self.arm.set_pose_target(pose)
125                 if not self.arm.go():
126                         print "position on target failed"
127                         return False
128
129                 self.gripper.set_named_target('close')
130                 self.gripper.go()
131
132                 # height
133                 angles = [0.0, 0.0, 0.0, 0.0]
134                 while angles == [0.0, 0.0, 0.0, 0.0]: # invalid
135                         angles = self.arm.get_current_joint_values()
136                 print [f * 180/pi for f in angles]
137                 angles[1]-=20*pi/180
138                 if not self.arm.go(angles, wait=True):
139                         print "higher position failed"
140                         return True
141
142                 if not self.arm.go([f * pi/180 for f in [-90, -60, 0, 0]], wait=True):
143                         print "final position failed"
144                         return True
145
146                 return True
147
148 if __name__ == "__main__":
149         OpenManipGrabBottle()