def generate_data(arg_step=10):
# argparser
parser = argparse.ArgumentParser()
parser.add_argument('--step', type=int, default=arg_step)
parser.add_argument('--nogui', action='store_true')
parser.add_argument('--path', type=str, default='./log')
parser.add_argument('--file_name', type=str, default=None)
args = parser.parse_args()
# logger
if args.file_name is not None:
mkdir(args.path)
logger = DataLogger()
# set robot & camera
cfg = get_sim_config()
view_matrix = tuple(invRt(cfg.Camera_pose).T.reshape(16))
imsize_w, imsize_h = 640, 360
f = 462
fov_w = math.atan(imsize_w / 2 / f) * 2 / math.pi * 180
fov_h = math.atan(imsize_h / 2 / f) * 2 / math.pi * 180
proj_matrix = p.computeProjectionMatrixFOV(fov=fov_h,
aspect=imsize_w / imsize_h,
nearVal=0.01,
farVal=5.0)
cam = simCam(imsize=[640, 360],
view_matrix=view_matrix,
proj_matrix=proj_matrix,
z_near=0.01,
z_far=5.0)
robot = SimUR5(center=cfg.Table_center,
prepare_height=cfg.Gripper_prepare_height,
real_time_sim=False,
GUI=not args.nogui)
obj_num = cfg.Object_num
# set object
scene = Sim_Scene(view_matrix=view_matrix,
fov_w=fov_w,
imsize=[imsize_w, imsize_h],
cube_num=cfg.Object_num)
image_RGB_ini, image_depth_ini, _ = cam.get_data()
# set property
friction_discrete, mass_discrete = set_friction_and_mass(
obj_num, scene.object_ids, cfg)
set_position_and_color(obj_num, scene.object_ids)
# always go home first
robot.open_gripper(blocking=True)
robot.close_gripper(blocking=True)
logger.save_attribute('friction', friction_discrete)
logger.save_attribute('mass', mass_discrete)
cnt_dict = {'collide': 0, 'push': 0}
all_proposal = {'collide': [], 'push': []}
# start action
for step in range(args.step):
robot.close_gripper(blocking=True)
image_RGB, image_depth, _ = cam.get_data()
mask = get_mask_real(image_RGB=image_RGB,
image_RGB_ini=image_RGB_ini,
image_depth=image_depth,
image_depth_ini=image_depth_ini,
Object_num=cfg.Object_num)
# depth normalization
tmp_depth = image_depth.copy()
tmp_depth -= np.min(tmp_depth)
tmp_depth /= np.max(tmp_depth)
obj_pos, points_list = [], []
for index in range(cfg.Object_num):
cur_pos, points = get_pos((mask == (index + 1)).astype(np.int),
image_RGB, image_depth, cfg.Camera_pose,
cfg.Camera_intrinsic)
obj_pos.append(cur_pos)
points_list.append(points)
scene.reset_coord_prev()
# generate proposal
# collide
for i in range(cfg.Object_num):
for dir in range(2):
if can_collide(obj_pos, i, cfg.Collide_region, dir):
all_proposal['collide'].append({
'action_type': 'collide',
'obj_id': i,
'dir': dir
})
# push
for i in range(cfg.Object_num):
for dir in range(cfg.direction_num):
if can_push(obj_pos, points_list, i,
dir / cfg.direction_num * 2 * np.pi,
cfg.Table_center):
all_proposal['push'].append({
'action_type': 'push',
'obj_id': i,
'dir': dir,
})
np.random.shuffle(all_proposal['push'])
np.random.shuffle(all_proposal['collide'])
proposal_list = all_proposal['push'] + all_proposal['collide']
if cnt_dict['push'] > cnt_dict['collide']:
proposal_list = all_proposal['collide'] + all_proposal['push']
if get_random(0, 1) < 0.2:
np.random.shuffle(proposal_list)
# choose action randomly
if len(proposal_list) == 0:
print('No Proposal')
exit()
action = proposal_list[0]
cnt_dict[action['action_type']] += 1
# prepare
cur_pos = obj_pos[action['obj_id']]
points = points_list[action['obj_id']]
logger.save_data(step, 'depth', image_depth)
logger.save_data(step, 'mask', mask)
logger.save_data(step, 'action_type', action['action_type'])
logger.save_data(step, 'obj_id', action['obj_id'])
if action['action_type'] == 'collide':
dir = action['dir']
aux_init = cfg.Auxiliary_initialization[dir]
# set the auxiliary cylinder
p.resetBasePositionAndOrientation(
scene.object_ids[cfg.Object_num],
posObj=aux_init + [cfg.Auxiliary_initialization_height],
ornObj=p.getQuaternionFromEuler([0, 0, np.pi / 2]))
# start action
robot.primitive_topdown_grasp(
aux_init + [cfg.Gripper_grasp_height], np.pi / 2)
robot.primitive_place([
cur_pos[0], cfg.Gripper_place_y[dir], cfg.Gripper_place_height
], np.pi / 2, cfg.Slow_speed)
robot.close_gripper(True)
# finish action
robot.primitive_gohome(speed=cfg.Fast_speed)
p.resetBasePositionAndOrientation(
scene.object_ids[cfg.Object_num],
posObj=cfg.Auxiliary_initialization[0] +
[cfg.Auxiliary_initialization_height],
ornObj=cfg.Orientation_initialization)
# save info
logger.save_data(step, 'action', [1] + [0, 0, 0] + [0] * 30 + [0] +
[cur_pos[0], action['dir']])
logger.save_data(step, 'direction', action['dir'])
elif action['action_type'] == 'push':
angle = action['dir'] / cfg.direction_num * 2 * np.pi
delta = -1 * np.asarray([
np.cos(angle), np.sin(angle)
]) * (get_dist(points, cur_pos,
[np.cos(angle), np.sin(angle)]) + 0.075)
speed_discrete = np.random.choice(len(cfg.Push_speed))
push_speed = cfg.Push_speed[speed_discrete]
start = [
cur_pos[0] + delta[0], cur_pos[1] + delta[1], cfg.Push_height
]
robot.primitive_push_tilt(start, angle, push_speed)
robot.primitive_gohome(speed=cfg.Fast_speed)
# save info
logger.save_data(step, 'action',
[0] + start + get_onehot(action['dir'], 30) +
[push_speed] + [0, 0])
logger.save_data(step, 'direction', action['dir'])
logger.save_data(step, 'speed', speed_discrete)
# calc the optical flow
flowim, depth_im_prev, depth_im_curr = scene.comput_2dflow(
get_mask=False)
_, flow = depth_smooth(depth_im_prev, flowim)
flow = flow * (mask == (action['obj_id'] + 1)).astype(np.float32)
logger.save_data(step, 'flow', flow)
if args.file_name is not None:
joblib.dump((logger.attribute_dict, logger.data_dict),
os.path.join(args.path, args.file_name),
compress=True)
