# Author: Jimmy Wu # Date: February 2023 import argparse import logging import math import time from collections import defaultdict from datetime import datetime from multiprocessing import Process from multiprocessing.connection import Client from pathlib import Path import cv2 as cv import numpy as np import utils from camera_server import CameraServer from constants import CAMERA_SERIALS, CAMERA_HEIGHT from constants import CONN_AUTHKEY from constants import FLOOR_LENGTH, FLOOR_WIDTH from constants import ROBOT_WIDTH, ROBOT_HEIGHT from controller_client import ControllerClient from controller_server import ControllerServer from image_client import ImageClient from marker_detector_server import MarkerDetectorServer from object_detector_server import ObjectDetectorServer, ObjectDetectorVisualizer from occupancy_map import OccupancyMap def get_grasp(mask): # Grasp center dist = cv.distanceTransform(mask, cv.DIST_L2, 5) # From https://stackoverflow.com/a/65409262 center = tuple(map(round, np.argwhere(dist == dist.max()).mean(axis=0))) center = (center[1], center[0]) # (i, j) to (x, y) # Grasp orientation contours, _ = cv.findContours(mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE) if len(contours) > 0: largest_contour = max(contours, key=cv.contourArea) moments = cv.moments(largest_contour) orientation = 0.5 * math.atan2(2 * moments['mu11'], moments['mu20'] - moments['mu02']) # From https://en.wikipedia.org/wiki/Image_moment orientation = -orientation # y goes up but j goes down else: orientation = None return center, orientation class Demo: def __init__(self, robot_idx, scenario_name, scale_factor=0.35, debug=False): self.robot_idx = robot_idx self.scale_factor = scale_factor self.debug = debug self.robot_state = 'idle' # States: idle, moving self.robot_pose = (0, 0, 0) self.robot_position = (0, 0) self.robot_gripper_state = None self.robot_arm_heading = None self.detected_object = None self.command = None self.last_command_time = 0 self.num_objs_per_recep = defaultdict(int) # Set up scenario scenario = utils.load_yaml(f'scenarios/{scenario_name}.yml') self.receptacles = scenario['receptacles'] self.occupancy_map = OccupancyMap(obstacles=self.receptacles.values(), debug=self.debug) self.preferences = utils.load_yaml(f'preferences/{scenario_name}.yml') # Image client for overhead images self.image_client = ImageClient(port1=6000, port2=6001, scale_factor=self.scale_factor) self.image_shape = self.image_client.image_shape # Graphical interface self.window_name = f'Demo (Robot {self.robot_idx + 1})' cv.namedWindow(self.window_name) # Logging log_dir = 'logs' Path(log_dir).mkdir(exist_ok=True) today = datetime.now().strftime('%Y-%m-%d') logging.basicConfig(filename=f'{log_dir}/{today}.log', format='%(asctime)s %(levelname)s %(filename)s: %(message)s', level=logging.INFO, encoding='utf-8') logging.info('Starting new session') def position_to_pixel_xy(self, position): return utils.position_to_pixel_xy(position, self.image_shape, self.scale_factor) def pixel_xy_to_position(self, pixel_xy): return utils.pixel_xy_to_position(pixel_xy, self.image_shape, self.scale_factor) def filter_detection_output(self, output): if len(output['boxes']) == 0: return output camera_center_offset_x = FLOOR_WIDTH * ((0.5486 + 0.5506) / 2 - 0.5) # 18 cm (from camera_params/*.json) ignore_boxes = [] # Base at floor height (assumes heading 0) min_xy = self.position_to_pixel_xy((self.robot_position[0] - ROBOT_WIDTH / 2, self.robot_position[1] + ROBOT_WIDTH / 2)) max_xy = self.position_to_pixel_xy((self.robot_position[0] + ROBOT_WIDTH / 2, self.robot_position[1] - ROBOT_WIDTH / 2)) ignore_boxes.append([min_xy[0], min_xy[1], max_xy[0], max_xy[1]]) # Top of base (assumes heading 0) height_ratio_base = CAMERA_HEIGHT / (CAMERA_HEIGHT - ROBOT_HEIGHT) if self.robot_position[1] < 0: robot_center_top = ( camera_center_offset_x + height_ratio_base * (self.robot_position[0] - camera_center_offset_x), (-FLOOR_LENGTH) / 4 + height_ratio_base * (self.robot_position[1] + FLOOR_LENGTH / 4)) else: robot_center_top = ( camera_center_offset_x + height_ratio_base * (self.robot_position[0] - camera_center_offset_x), FLOOR_LENGTH / 4 + height_ratio_base * (self.robot_position[1] - FLOOR_LENGTH / 4)) robot_width_top = height_ratio_base * ROBOT_WIDTH min_xy = self.position_to_pixel_xy((robot_center_top[0] - robot_width_top / 2, robot_center_top[1] + robot_width_top / 2)) max_xy = self.position_to_pixel_xy((robot_center_top[0] + robot_width_top / 2, robot_center_top[1] - robot_width_top / 2)) ignore_boxes.append([min_xy[0], min_xy[1], max_xy[0], max_xy[1]]) # Arm if self.robot_arm_heading is not None: arm_width = 0.15 # Joint 4 of 7 height_ratio_arm_joint_4 = CAMERA_HEIGHT / (CAMERA_HEIGHT - (ROBOT_HEIGHT + 0.69)) theta = self.robot_pose[2] - math.radians(self.robot_arm_heading) arm_position_joint_4 = (self.robot_pose[0] + 0.15 * math.cos(theta), self.robot_pose[1] + 0.15 * math.sin(theta)) if arm_position_joint_4[1] < 0: arm_center_joint_4 = ( camera_center_offset_x + height_ratio_arm_joint_4 * (arm_position_joint_4[0] - camera_center_offset_x), (-FLOOR_LENGTH) / 4 + height_ratio_arm_joint_4 * (arm_position_joint_4[1] + FLOOR_LENGTH / 4)) else: arm_center_joint_4 = ( camera_center_offset_x + height_ratio_arm_joint_4 * (arm_position_joint_4[0] - camera_center_offset_x), FLOOR_LENGTH / 4 + height_ratio_arm_joint_4 * (arm_position_joint_4[1] - FLOOR_LENGTH / 4)) arm_width_joint_4 = height_ratio_arm_joint_4 * arm_width min_xy = self.position_to_pixel_xy((arm_center_joint_4[0] - arm_width_joint_4 / 2, arm_center_joint_4[1] + arm_width_joint_4 / 2)) max_xy = self.position_to_pixel_xy((arm_center_joint_4[0] + arm_width_joint_4 / 2, arm_center_joint_4[1] - arm_width_joint_4 / 2)) ignore_boxes.append([min_xy[0], min_xy[1], max_xy[0], max_xy[1]]) # Joint 6 of 7 height_ratio_arm_joint_6 = CAMERA_HEIGHT / (CAMERA_HEIGHT - (ROBOT_HEIGHT + 0.51)) arm_position_joint_6 = (self.robot_pose[0] - 0.11 * math.cos(theta), self.robot_pose[1] - 0.11 * math.sin(theta)) if arm_position_joint_6[1] < 0: arm_center_joint_6 = ( camera_center_offset_x + height_ratio_arm_joint_6 * (arm_position_joint_6[0] - camera_center_offset_x), (-FLOOR_LENGTH) / 4 + height_ratio_arm_joint_6 * (arm_position_joint_6[1] + FLOOR_LENGTH / 4)) else: arm_center_joint_6 = ( camera_center_offset_x + height_ratio_arm_joint_6 * (arm_position_joint_6[0] - camera_center_offset_x), FLOOR_LENGTH / 4 + height_ratio_arm_joint_6 * (arm_position_joint_6[1] - FLOOR_LENGTH / 4)) arm_width_joint_6 = height_ratio_arm_joint_6 * arm_width min_xy = self.position_to_pixel_xy((arm_center_joint_6[0] - arm_width_joint_6 / 2, arm_center_joint_6[1] + arm_width_joint_6 / 2)) max_xy = self.position_to_pixel_xy((arm_center_joint_6[0] + arm_width_joint_6 / 2, arm_center_joint_6[1] - arm_width_joint_6 / 2)) ignore_boxes.append([min_xy[0], min_xy[1], max_xy[0], max_xy[1]]) # Receptacles # Note: Due to camera perspective distortion, receptacles should not be placed near edges of overhead camera image frame for receptacle in self.receptacles.values(): pos_x, pos_y = receptacle['position'] dim_x, dim_y = receptacle['dimensions'] min_xy = self.position_to_pixel_xy((pos_x - dim_x / 2, pos_y + dim_y / 2)) max_xy = self.position_to_pixel_xy((pos_x + dim_x / 2, pos_y - dim_y / 2)) ignore_boxes.append([min_xy[0], min_xy[1], max_xy[0], max_xy[1]]) # output['boxes'] = ignore_boxes # output['masks'] = len(ignore_boxes) * output['masks'][0] # output['categories'] = ['robot_floor', 'robot_top', 'arm_joint_4', 'arm_joint_6', *self.receptacles.keys()] # output['scores'] = len(ignore_boxes) * [1.0] # return output keep_indices = [] for i, box in enumerate(output['boxes']): skip = False # Skip detections that overlap with ignore boxes for ignore_box in ignore_boxes: w = max(0, min(box[2], ignore_box[2]) - max(box[0], ignore_box[0])) h = max(0, min(box[3], ignore_box[3]) - max(box[1], ignore_box[1])) intersection = w * h if intersection > 0: skip = True break if not skip: keep_indices.append(i) keep_indices = np.array(keep_indices, dtype=np.int64) # Explicit dtype needs to be specified for empty list output['boxes'] = output['boxes'][keep_indices] output['masks'] = output['masks'][keep_indices] output['categories'] = [output['categories'][i] for i in keep_indices] output['scores'] = output['scores'][keep_indices] return output def distance(self, position): path = self.occupancy_map.shortest_path(self.robot_position, position) return sum(utils.distance(path[i - 1], path[i]) for i in range(1, len(path))) def get_closest_grasp(self, boxes, masks): best_dist = math.inf center = None orientation = None for i, box in enumerate(boxes): xmin, ymin, xmax, ymax = box box_center = ((xmin + xmax) / 2, (ymin + ymax) / 2) dist = self.distance(self.pixel_xy_to_position(box_center)) # 30 ms if dist < best_dist: best_dist = dist center, orientation = get_grasp(masks[i]) # 4 ms return center, orientation def update_command(self, detection_output=None): if self.robot_state != 'idle': return if self.robot_gripper_state == 'open': # Pick up object if detection_output is None: # Wait for detection output return grasp_center, grasp_orientation = self.get_closest_grasp(detection_output['boxes'], detection_output['masks']) if grasp_center is not None: # Pick up closest object self.command = { 'primitive_name': 'pick', 'waypoints': self.occupancy_map.shortest_path(self.robot_position, self.pixel_xy_to_position(grasp_center)), 'grasp_orientation': grasp_orientation, 'categories': self.preferences['categories'], 'image_path': detection_output['image_path'], } else: # No objects found self.command = None elif self.robot_gripper_state == 'closed': if self.detected_object is not None: # Bring object to receptacle # Select receptacle receptacle_name = self.preferences['placements'][self.detected_object] receptacle = self.receptacles[receptacle_name] target_ee_pos = (receptacle['position'][0], receptacle['position'][1]) # Select primitive primitive_name = self.preferences['primitives'].get(self.detected_object, 'place') # place or toss if primitive_name not in receptacle['primitive_names']: primitive_name = 'place' if primitive_name == 'toss' and utils.distance(self.robot_position, target_ee_pos) < 0.9: # If close to receptacle, just place instead of toss primitive_name = 'place' if 'shelf' in receptacle_name: primitive_name = 'shelf' elif 'drawer' in receptacle_name: primitive_name = 'drawer' # Build command command = {'receptacle_name': receptacle_name, 'primitive_name': primitive_name} if 'center_offset' in receptacle: target_ee_pos = (receptacle['position'][0] + receptacle['center_offset'][0], receptacle['position'][1] + receptacle['center_offset'][1]) if receptacle_name in {'coffee table', 'sofa'}: # Avoid stacking objects on top of each other num_objs = self.num_objs_per_recep[receptacle_name] offset = ((num_objs + 1) // 2) * (1 if num_objs % 2 == 0 else -1) if receptacle['dimensions'][0] > receptacle['dimensions'][1]: # Longer in x-direction target_ee_pos = (target_ee_pos[0] + 0.1 * offset * receptacle['dimensions'][0], target_ee_pos[1]) else: # Longer in y-direction target_ee_pos = (target_ee_pos[0], target_ee_pos[1] + 0.1 * offset * receptacle['dimensions'][1]) if 'end_effector_offset' in receptacle: if receptacle['dimensions'][0] > receptacle['dimensions'][1]: # Longer in x-direction new_waypoint = (target_ee_pos[0], target_ee_pos[1] + math.copysign(1, receptacle['center_offset'][1]) * receptacle['end_effector_offset']) else: # Longer in y-direction new_waypoint = (target_ee_pos[0] + math.copysign(1, receptacle['center_offset'][0]) * receptacle['end_effector_offset'], target_ee_pos[1]) command['waypoints'] = self.occupancy_map.shortest_path(self.robot_position, new_waypoint) + [target_ee_pos] else: command['waypoints'] = self.occupancy_map.shortest_path(self.robot_position, target_ee_pos) if primitive_name == 'place': command['place_height'] = receptacle['place_height'] self.command = command else: # Object detection took too long, put object back down target_ee_pos = (self.robot_position[0] - 0.55 * math.cos(self.robot_pose[2]), self.robot_position[1] - 0.55 * math.sin(self.robot_pose[2])) self.command = {'primitive_name': 'place', 'waypoints': [self.robot_position, target_ee_pos]} def store_controller_data(self, controller_data): if 'base' in controller_data: self.robot_pose = controller_data['base']['pose'] self.robot_position = (self.robot_pose[0], self.robot_pose[1]) if 'arm' in controller_data: self.robot_gripper_state = controller_data['arm']['gripper_state'] self.robot_arm_heading = controller_data['arm']['arm_heading'] self.detected_object = controller_data.get('detected_object', None) def run(self): # Connect to controller controller = ControllerClient(self.robot_idx) # Set up object detection object_detector_conn = Client(('localhost', 6003), authkey=CONN_AUTHKEY) if self.debug: visualizer = ObjectDetectorVisualizer() # Set up the different operating modes default_idx = 0 operating_modes = ['supervised', 'autonomous'] operating_mode = operating_modes[default_idx] operating_mode_descs = [ 'Supervised mode', 'Autonomous mode', ] print('Operating mode:') for i, desc in enumerate(operating_mode_descs): if i == default_idx: print(f'{i}. {desc} (default)') else: print(f'{i}. {desc}') # Run GUI loop idle_steps = 0 esc_pressed_count = 0 autonomous_running = False detection_image_sent = False last_time = time.time() controller_start_time = None while True: step_time = time.time() - last_time if step_time > 0.5: # 2 Hz print(f'Warning: Step time {1000 * step_time:.1f} ms in {self.__class__.__name__}') last_time = time.time() # Update robot state controller_data = controller.get_controller_data() # 30 ms if controller_data['start_time'] != controller_start_time: # Controller was restarted controller_start_time = controller_data['start_time'] autonomous_running = False if controller_data['state'] == 'moving': self.robot_state = 'moving' elif controller_data['state'] == 'idle' and self.robot_state == 'moving': idle_steps += 1 if idle_steps > 3: idle_steps = 0 self.robot_state = 'idle' self.command = None # Store controller data self.store_controller_data(controller_data) # Check for keypress key = cv.waitKey(1) if key == ord('q') or cv.getWindowProperty(self.window_name, cv.WND_PROP_VISIBLE) < 0.5: break if key == 27: # controller.stop() esc_pressed_count += 1 if esc_pressed_count > 1: controller.reset_visualizer() autonomous_running = False elif key == 13: # esc_pressed_count = 0 autonomous_running = True elif ord(str(0)) <= key < ord(str(len(operating_modes))): idx = key - ord(str(0)) operating_mode = operating_modes[idx] print(operating_mode_descs[idx]) # Send command if self.robot_state == 'idle' and autonomous_running and self.command is not None and time.time() - self.last_command_time > 1.0: logging.info(f'Sending command: {self.command}') # pylint: disable=logging-fstring-interpolation controller.execute_command(self.command) self.last_command_time = time.time() # Wait at least 1 sec between commands if 'receptacle_name' in self.command: self.num_objs_per_recep[self.command['receptacle_name']] += 1 self.robot_state = 'moving' if operating_mode == 'supervised': # In supervised mode, every command is manually approved autonomous_running = False # Draw waypoints and grasp orientation image = self.image_client.get_image() image_copy = image.copy() if self.command is not None: waypoints = list(map(self.position_to_pixel_xy, self.command['waypoints'])) for i in range(1, len(waypoints)): cv.line(image_copy, waypoints[i - 1], waypoints[i], (255, 0, 0), 2) grasp_orientation = self.command.get('grasp_orientation', None) if grasp_orientation is not None: grasp_center = self.command['waypoints'][-1] robotiq_gripper_width = 0.085 grasp_line = [(grasp_center[0] - robotiq_gripper_width / 2 * math.sin(grasp_orientation), grasp_center[1] + robotiq_gripper_width / 2 * math.cos(grasp_orientation)), (grasp_center[0] + robotiq_gripper_width / 2 * math.sin(grasp_orientation), grasp_center[1] - robotiq_gripper_width / 2 * math.cos(grasp_orientation))] cv.line(image_copy, self.position_to_pixel_xy(grasp_line[0]), self.position_to_pixel_xy(grasp_line[1]), (255, 0, 0), 2) cv.imshow(self.window_name, image_copy) # Object detection detection_output = None if object_detector_conn.poll(): detection_output = object_detector_conn.recv() detection_image_sent = False detection_output = self.filter_detection_output(detection_output) if self.debug and self.robot_state == 'idle': visualizer.visualize(cv.cvtColor(image, cv.COLOR_BGR2RGB), detection_output) if self.robot_state == 'idle' and not detection_image_sent: _, encoded_image = cv.imencode('.jpg', image) object_detector_conn.send({'encoded_image': encoded_image, 'categories': ['object'], 'min_box_area': 36, 'max_box_area': 40000}) detection_image_sent = True # Update command self.update_command(detection_output=detection_output) cv.destroyAllWindows() def main(args): # Start camera servers def start_camera_server(serial, port): CameraServer(serial, port=port).run() for serial, port in [(CAMERA_SERIALS[0], 6000), (CAMERA_SERIALS[1], 6001)]: Process(target=start_camera_server, args=(serial, port), daemon=True).start() # Wait for camera servers to be ready time.sleep(1.5) # Start marker detector server def start_marker_detector_server(): MarkerDetectorServer(hostname='0.0.0.0').run() Process(target=start_marker_detector_server, daemon=True).start() # Start controller server robot_idx = args.robot_num - 1 def start_controller_server(): ControllerServer(robot_idx).run() Process(target=start_controller_server, daemon=True).start() # Start object detector server # def start_object_detector_server(): # ObjectDetectorServer(hostname='0.0.0.0').run() # Process(target=start_object_detector_server, daemon=True).start() # Run demo Demo(robot_idx, args.scenario_name, debug=args.debug).run() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--robot-num', type=int, default=1) parser.add_argument('--scenario-name', default='test') parser.add_argument('--debug', action='store_true') main(parser.parse_args())