-
Notifications
You must be signed in to change notification settings - Fork 0
/
pick_and_place_utils.py
149 lines (115 loc) · 4.76 KB
/
pick_and_place_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
import rospy
import numpy as np
from sensor_msgs import msg
import cv2
import cv_bridge
from copy import deepcopy
import dvrk
import PyKDL
import tf
import image_geometry
from feature_processor import feature_processor, FeatureType
from tf_conversions import posemath
from math import pi
from collections import OrderedDict
'''
Some useful methods and constants for picking up a ball with dVRK and CoppeliaSim
'''
PSM_J1_TO_BASE_LINK_ROT = PyKDL.Rotation.RPY(pi / 2, - pi, 0)
PSM_J1_TO_BASE_LINK_TF = PyKDL.Frame(PSM_J1_TO_BASE_LINK_ROT, PyKDL.Vector())
# TODO: make this less hardcoded
RED_BALL_FEAT_PATH = './red_ball.csv'
BLUE_BALL_FEAT_PATH = './blue_ball.csv'
GREEN_BALL_FEAT_PATH = './green_ball.csv'
FEAT_PATHS = [RED_BALL_FEAT_PATH, BLUE_BALL_FEAT_PATH, GREEN_BALL_FEAT_PATH]
CV_TO_CAM_FRAME_ROT = np.asarray([
[-1, 0, 0],
[0, -1, 0],
[0, 0, 1]
])
class Object3d:
def __init__(self, pos, type, color):
self.pos = pos
self.type = type
self.color = color
def __str__(self):
return "<Object3d pos: {} type: {} color: {}>".format(self.pos, self.type, self.color)
def __repr__(self):
return self.__str__()
def clamp_image_coords(pt, im_shape):
return tuple(np.clip(pt, (0, 0), np.array(im_shape)[:2] - np.array([1, 1])))
def get_objects_and_img(left_image_msg, right_image_msg, stereo_cam_model, cam_to_world_tf):
# this gets the position of the red ball thing in the camera frame
# and the image with X's on the desired features
fp = feature_processor(FEAT_PATHS)
left_feats, left_frame = fp.FindImageFeatures(left_image_msg)
right_feats, right_frame = fp.FindImageFeatures(right_image_msg)
# discard features with image y > bowl y
left_bowl = None
for left_feat in left_feats:
if left_feat.type == FeatureType.BOWL:
left_bowl = left_feat
left_feats = filter(lambda feat : feat.pos[1] >= left_bowl.pos[1], left_feats)
left_feats.sort(key=lambda feat: feat.pos[1], reverse=False)
right_bowl = None
for right_feat in right_feats:
if right_feat.type == FeatureType.BOWL:
right_bowl = right_feat
right_feats = filter(lambda feat : feat.pos[1] >= right_bowl.pos[1], right_feats)
right_feats.sort(key=lambda feat: feat.pos[1], reverse=False)
matched_feats = []
for left_feat in left_feats:
same_color_feats = filter(lambda right_feat: right_feat.color == left_feat.color,
right_feats)
if not same_color_feats:
rospy.logwarn(
"Failed to match left detection {} to a right detection!".format(left_feat))
continue
matched_feats.append((left_feat,
min(same_color_feats,
key=lambda right_feat: (right_feat.pos[0] - left_feat.pos[0]) ** 2 \
+ (right_feat.pos[1] - left_feat.pos[1]) ** 2)))
objects = []
for left_feat, right_feat in matched_feats:
disparity = abs(left_feat.pos[0] - right_feat.pos[0])
pos_cv = stereo_cam_model.projectPixelTo3d(left_feat.pos, float(disparity))
# there's a fixed rotation to convert this to the camera coordinate frame
pos_cam = np.matmul(CV_TO_CAM_FRAME_ROT, pos_cv)
pos = None
if cam_to_world_tf is not None:
pos = cam_to_world_tf * PyKDL.Vector(*pos_cam)
else:
pos = PyKDL.Vector(*pos_cam)
if left_feat.color != right_feat.color:
rospy.loginfo("Color mismatch between left and right detection")
objects.append(Object3d(pos, left_feat.type, left_feat.color))
return objects, np.hstack((left_frame, right_frame))
def tf_to_pykdl_frame(tfl_frame):
pos, rot_quat = tfl_frame
pos2 = PyKDL.Vector(*pos)
rot = PyKDL.Rotation.Quaternion(*rot_quat)
return PyKDL.Frame(rot, pos2)
class World:
def __init__(self, all_detections):
self.objects = []
for object in all_detections:
if object.type == FeatureType.BOWL:
self.bowl = object
else:
self.objects.append(object)
def __str__(self):
return "<World objects: {}\nbowl: {}>".format(self.objects, self.bowl)
def __repr__(self):
return self.__str__()
def get_objects_for_psms(objects, world_to_psm_tfs):
'''
Returns a dict of PSM index -> list of objects that are closest to that PSM
'''
result = dict()
for object in objects:
closest_psm_idx = world_to_psm_tfs.index(
min(world_to_psm_tfs, key=lambda tf : (tf * object.pos).Norm()))
if closest_psm_idx not in result:
result[closest_psm_idx] = list()
result[closest_psm_idx].append(object)
return result