class VSGazeboEnv(object):
def __init__(self):
self._robot = RobotControl(robotName=ROBOT_NAME, dof=DOF)
self._world = WorldControl()
self._pin_hole_left_cam = Camera(
rgbImageTopic=PIN_HOLE_LEFT_CAMERA_TOPIC)
self._pin_hole_right_cam = Camera(
rgbImageTopic=PIN_HOLE_RIGHT_CAMERA_TOPIC)
self._inspected_pin_leads_in_cameras = None
self._board_pose_mm_deg = None
self._robot_init_pose_mm_deg = None
self._left_pin_lead_pose_init_mm_deg = None
self._right_pin_lead_pose_init_mm_deg = None
# ----- get images ------
def getPinHoleLeftImage(self):
return self._pin_hole_left_cam.getRGBImage()
def getPinHoleRightImage(self):
return self._pin_hole_right_cam.getRGBImage()
# ----- robot manipulation -----
def getRobotPose(self):
return self._robot.getRobotPos()
def setRobotPose(self, pose_mm_deg, wait=True):
self._robot.setRobotPos(pose_mm_deg)
timeout = 0
while wait and self._robot.isMoving:
time.sleep(0.01)
timeout += 0.01
if timeout > ROBOT_TIMEOUT_SEC:
raise VSGazeboEnvError('Robot moving timeout')
def setRobotJointAngles(self, angles_deg, wait=True):
self._robot.setJointAngle(angles_deg)
timeout = 0
while wait and self._robot.isMoving:
time.sleep(0.01)
timeout += 0.01
if timeout > ROBOT_TIMEOUT_SEC:
raise VSGazeboEnvError('Robot moving timeout')
def resetRobotToHome(self):
self.setRobotJointAngles([0] * 6)
def goToLeadInspectionPose(self):
self.setRobotPose(LEAD_INSPECTION_POSE_MM_DEG)
def setRobotSpeed(self, speed_percentage):
self._robot.setSpeedRate(speed_percentage)
# ----- model manipulation -----
def _getModelPose_mm_deg(self, link_name):
return self._world.getLinkPos(link_name)
def _setModelPose(self, link_name, pose_mm_deg):
self._world.setLinkPos(link_name, pose_mm_deg)
# ----- board related actions -----
def getBoardPose_mm_deg(self):
return self._getModelPose_mm_deg(BOARD_LINK)
def setBoardPose(self, pose_mm_deg):
self._setModelPose(BOARD_LINK, pose_mm_deg)
def getHolePoses_mm_deg(self):
board_pose_mm_deg = self.getBoardPose_mm_deg()
t_board_to_ref = Pose2T(board_pose_mm_deg)
t_left_hole_to_board = Pose2T(LEFT_HOLE2BOARD_MM_DEG)
t_right_hole_to_board = Pose2T(RIGHT_HOLE2BOARD_MM_DEG)
t_left_hole_to_ref = t_board_to_ref.dot(t_left_hole_to_board)
t_right_hole_to_ref = t_board_to_ref.dot(t_right_hole_to_board)
left_hole_pose_mm_deg = T2Pose(t_left_hole_to_ref)
right_hole_pose_mm_deg = T2Pose(t_right_hole_to_ref)
return left_hole_pose_mm_deg.flatten(), right_hole_pose_mm_deg.flatten(
)
# ----- pin related operations -----
def getPinLink3PoseByCBase_mm_deg(self, pin_link3_to_c_base):
c_base_pose_mm_deg = self._getModelPose_mm_deg(C_BASE_LINK)
t_c_base_to_ref = Pose2T(c_base_pose_mm_deg)
t_pin_link3_to_c_base = Pose2T(pin_link3_to_c_base)
t_pin_link3_to_ref = t_c_base_to_ref.dot(t_pin_link3_to_c_base)
pin_link3_pose_mm_deg = T2Pose(t_pin_link3_to_ref)
return pin_link3_pose_mm_deg.flatten()
def getLeftPinLink3Pose_mm_deg(self, by_c_base=False):
return self.getPinLink3PoseByCBase_mm_deg(LEFT_PIN_LINK3_TO_CBASE_MM_DEG) if by_c_base else \
self._getModelPose_mm_deg(LEFT_PIN_LINK_3)
def getRightPinLink3Pose_mm_deg(self, by_c_base=False):
return self.getPinLink3PoseByCBase_mm_deg(RIGHT_PIN_LINK3_TO_CBASE_MM_DEG) if by_c_base else \
self._getModelPose_mm_deg(RIGHT_PIN_LINK_3)
def getPinLeadPoses_mm_deg(self, by_c_base=False):
left_pin_link3_pose_mm_deg = self.getLeftPinLink3Pose_mm_deg(by_c_base)
right_pin_link3_pose_mm_deg = self.getRightPinLink3Pose_mm_deg(
by_c_base)
t_left_pin_link3_to_ref = Pose2T(left_pin_link3_pose_mm_deg)
t_right_pin_link3_to_ref = Pose2T(right_pin_link3_pose_mm_deg)
t_lead_to_link3 = Pose2T(PIN_LEAD2LINK3_MM_DEG)
t_left_pin_lead_to_ref = t_left_pin_link3_to_ref.dot(t_lead_to_link3)
t_right_pin_lead_to_ref = t_right_pin_link3_to_ref.dot(t_lead_to_link3)
left_pin_lead_pose_mm_deg = T2Pose(t_left_pin_lead_to_ref)
right_pin_lead_pose_mm_deg = T2Pose(t_right_pin_lead_to_ref)
return left_pin_lead_pose_mm_deg.flatten(
), right_pin_lead_pose_mm_deg.flatten()
def getPinLeadsHolesInImage(self, camera_link):
f_width = IMAGE_WIDTH_PX / (np.tan(FOV / 2.0) * 2.0)
cameraMatrix = np.array(
[[f_width, 0, IMAGE_WIDTH_PX / 2.0],
[0, f_width, IMAGE_HEIGHT_PX / 2.0], [0, 0, 1]],
dtype=np.float32)
camera_pose_mm_deg = self._getModelPose_mm_deg(camera_link)
t_camera_to_ref = Pose2T(camera_pose_mm_deg)
t_ref_to_camera = np.linalg.inv(t_camera_to_ref)
# -----pin leads in image -----
left_pin_lead_pose_mm_deg, right_pin_lead_pose_mm_deg = self.getPinLeadPoses_mm_deg(
)
left_pin_in_camera_px = projectPtsToImg(left_pin_lead_pose_mm_deg[:3],
t_ref_to_camera, cameraMatrix,
[])
right_pin_in_camera_px = projectPtsToImg(
right_pin_lead_pose_mm_deg[:3], t_ref_to_camera, cameraMatrix, [])
# -----holes in image -----
left_hole_pose_mm_deg, right_hole_pose_mm_deg = self.getHolePoses_mm_deg(
)
left_hole_in_camera_px = projectPtsToImg(left_hole_pose_mm_deg[:3],
t_ref_to_camera, cameraMatrix,
[])
right_hole_in_camera_px = projectPtsToImg(right_hole_pose_mm_deg[:3],
t_ref_to_camera,
cameraMatrix, [])
return left_pin_in_camera_px.flatten(), right_pin_in_camera_px.flatten(), \
left_hole_in_camera_px.flatten(), right_hole_in_camera_px.flatten()
def getTcA2ref(self):
camera_pose_mm_deg = self._getModelPose_mm_deg(
PIN_HOLE_LEFT_CAMERA_LINK)
t_camera_to_ref = Pose2T(camera_pose_mm_deg)
return t_camera_to_ref
def getTcB2ref(self):
camera_pose_mm_deg = self._getModelPose_mm_deg(
PIN_HOLE_RIGHT_CAMERA_LINK)
t_camera_to_ref = Pose2T(camera_pose_mm_deg)
return t_camera_to_ref
def getPinLeadsHolesInLeftCamera(self):
return self.getPinLeadsHolesInImage(PIN_HOLE_LEFT_CAMERA_LINK)
def getPinLeadsHolesInRightCamera(self):
return self.getPinLeadsHolesInImage(PIN_HOLE_RIGHT_CAMERA_LINK)
def getPinLeadsHolesInCameras(self):
# ----- pin leads and holes in left camera -----
left_pin_in_left_camera_px, right_pin_in_left_camera_px, \
left_hole_in_left_camera_px, right_hole_in_left_camera_px = self.getPinLeadsHolesInLeftCamera()
# ----- pin leads and holes in right camera -----
left_pin_in_right_camera_px, right_pin_in_right_camera_px, \
left_hole_in_right_camera_px, right_hole_in_right_camera_px = self.getPinLeadsHolesInRightCamera()
return np.r_[left_pin_in_left_camera_px, right_pin_in_left_camera_px,
left_pin_in_right_camera_px, right_pin_in_right_camera_px], \
np.r_[left_hole_in_left_camera_px, right_hole_in_left_camera_px,
left_hole_in_right_camera_px, right_hole_in_right_camera_px]
# def setPinJointAngles(self, joint_angles_deg):
def getTc2r(self, camera_link):
camera_pose_mm_deg = self._getModelPose_mm_deg(camera_link)
t_camera_to_ref = Pose2T(camera_pose_mm_deg)
robot_base_pose = [0., 0., 600., 0., 0., 0.]
t_robot_base_to_ref = Pose2T(robot_base_pose)
t_ref_to_robot_base = np.linalg.inv(t_robot_base_to_ref)
t_camera_to_robot_base = np.dot(t_ref_to_robot_base, t_camera_to_ref)
return t_camera_to_robot_base
def getTcA2r(self):
return self.getTc2r(PIN_HOLE_LEFT_CAMERA_LINK)
def getTcB2r(self):
return self.getTc2r(PIN_HOLE_RIGHT_CAMERA_LINK)
def getTcA2cB(self):
TcA2r = self.getTcA2r()
TcB2r = self.getTcB2r()
Tr2cB = np.linalg.inv(TcB2r)
TcA2cB = np.dot(Tr2cB, TcA2r)
return TcA2cB
def getTref2r(self):
robot_base_pose = [0., 0., 600., 0., 0., 0.]
t_robot_base_to_ref = Pose2T(robot_base_pose)
t_ref_to_robot_base = np.linalg.inv(t_robot_base_to_ref)
return t_ref_to_robot_base
def setVirtualPointsInImg(self, delta_z, camera_link):
f_width = IMAGE_WIDTH_PX / (np.tan(FOV / 2.0) * 2.0)
cameraMatrix = np.array(
[[f_width, 0, IMAGE_WIDTH_PX / 2.0],
[0, f_width, IMAGE_HEIGHT_PX / 2.0], [0, 0, 1]],
dtype=np.float32)
camera_pose_mm_deg = self._getModelPose_mm_deg(camera_link)
t_camera_to_ref = Pose2T(camera_pose_mm_deg)
t_ref_to_camera = np.linalg.inv(t_camera_to_ref)
# -----pin leads in image -----
left_pin_lead_pose_mm_deg, right_pin_lead_pose_mm_deg = self.getPinLeadPoses_mm_deg(
)
left_pin_in_camera_px = projectPtsToImg(left_pin_lead_pose_mm_deg[:3],
t_ref_to_camera, cameraMatrix,
[])
right_pin_in_camera_px = projectPtsToImg(
right_pin_lead_pose_mm_deg[:3], t_ref_to_camera, cameraMatrix, [])
# -----holes in image -----
left_hole_pose_mm_deg, right_hole_pose_mm_deg = self.getHolePoses_mm_deg(
)
left_hole_pose_mm_deg[2] += delta_z
right_hole_pose_mm_deg[2] += delta_z
left_hole_in_camera_px = projectPtsToImg(left_hole_pose_mm_deg[:3],
t_ref_to_camera, cameraMatrix,
[])
right_hole_in_camera_px = projectPtsToImg(right_hole_pose_mm_deg[:3],
t_ref_to_camera,
cameraMatrix, [])
return left_pin_in_camera_px.flatten(), right_pin_in_camera_px.flatten(), \
left_hole_in_camera_px.flatten(), right_hole_in_camera_px.flatten()
def getVirtualPointsInImg(self, delta_z):
left_pin_in_left_camera, right_pin_in_left_camera, left_hole_in_left_camera, right_hole_in_left_camera = self.setVirtualPointsInImg(
delta_z, PIN_HOLE_LEFT_CAMERA_LINK)
left_pin_in_right_camera, right_pin_in_right_camera, left_hole_in_right_camera, right_hole_in_right_camera = self.setVirtualPointsInImg(
delta_z, PIN_HOLE_RIGHT_CAMERA_LINK)
virtualImgPtsA_hole = np.hstack(
(left_hole_in_left_camera.reshape(2, -1),
right_hole_in_left_camera.reshape(2, -1)))
virtualImgPtsB_hole = np.hstack(
(left_hole_in_right_camera.reshape(2, -1),
right_hole_in_right_camera.reshape(2, -1)))
return virtualImgPtsA_hole, virtualImgPtsB_hole
def getRealImgPts_hole(self):
_, _, left_hole_in_left_camera, right_hole_in_left_camera = self.getPinLeadsHolesInLeftCamera(
)
_, _, left_hole_in_right_camera, right_hole_in_right_camera = self.getPinLeadsHolesInRightCamera(
)
realImgPtsA_hole = np.hstack((left_hole_in_left_camera.reshape(2, -1),
right_hole_in_left_camera.reshape(2,
-1)))
realImgPtsB_hole = np.hstack(
(left_hole_in_right_camera.reshape(2, -1),
right_hole_in_right_camera.reshape(2, -1)))
return realImgPtsA_hole, realImgPtsB_hole
def getRealImgPts_pin(self):
left_pin_in_left_camera, right_pin_in_left_camera, _, _, = self.getPinLeadsHolesInLeftCamera(
)
left_pin_in_right_camera, right_pin_in_right_camera, _, _, = self.getPinLeadsHolesInRightCamera(
)
RealImgPtsA_pin = np.hstack((left_pin_in_left_camera.reshape(2, -1),
right_pin_in_left_camera.reshape(2, -1)))
RealImgPtsB_pin = np.hstack((left_pin_in_right_camera.reshape(2, -1),
right_pin_in_right_camera.reshape(2, -1)))
return RealImgPtsA_pin, RealImgPtsB_pin
def getRealPinPtsInRob(self):
left_pin_lead_pose_mm_deg, right_pin_lead_pose_mm_deg = self.getPinLeadPoses_mm_deg(
)
t_ref_to_robot_base = self.getTref2r()
RealPinPtsInRob_left = projectPts(
left_pin_lead_pose_mm_deg[:3].reshape(3, -1), t_ref_to_robot_base)
RealPinPtsInRob_right = projectPts(
right_pin_lead_pose_mm_deg[:3].reshape(3, -1), t_ref_to_robot_base)
RealPinPtsInRob = np.hstack((RealPinPtsInRob_left.reshape(3, -1),
RealPinPtsInRob_right.reshape(3, -1)))
return RealPinPtsInRob
# ----- operations for RL algorithms -----
def reset(self, board_pose_mm_deg, robot_pose_mm_deg):
self._board_pose_mm_deg = board_pose_mm_deg
self._robot_init_pose_mm_deg = robot_pose_mm_deg
self.setBoardPose(board_pose_mm_deg)
self.goToLeadInspectionPose()
self._inspected_pin_leads_in_cameras = self.getPinLeadsHolesInCameras()
self.setRobotPose(robot_pose_mm_deg)
self._left_pin_lead_pose_init_mm_deg, self._right_pin_lead_pose_init_mm_deg = \
self.getPinLeadPoses_mm_deg(GET_LINK3_POSE_BY_C_BASE)
return True
# return False if not self.straighten_pins() else True
def getInspectedPinLeadsInCameras(self):
return self._inspected_pin_leads_in_cameras
