def render(self, mode='rgb_array'): if mode == 'rgb_array': img, depth = pyflex.render() width, height = self.camera_params['default_camera']['width'], self.camera_params['default_camera']['height'] img = img.reshape(height, width, 4)[::-1, :, :3] # Need to reverse the height dimension return img elif mode == 'human': raise NotImplementedError
def set_scene(self, config, states=None): ''' Construct the passing water scence. ''' # create fluid super().set_scene(config) # do not sample fluid parameters, as it's very likely to generate very strange fluid # compute box params if states is None: self.set_box_params(config) else: box_params = states['box_params'] self.height = box_params['height'] self.box_dis_x = box_params['box_dis_x'] self.box_dis_z = box_params['box_dis_z'] self.box_params = box_params self.x_center = 0 # create box self.create_box(self.box_dis_x, self.box_dis_z, self.height) # move box to be at ground or on the table self.box_states = self.init_box_state(self.x_center, 0, self.box_dis_x, self.box_dis_z, self.height) pyflex.set_shape_states(self.box_states) # record box floor center x self.box_x = self.x_center # no cached init states passed in if states is None: for _ in range(50): pyflex.step() pyflex.render() return True else: # set to passed-in cached init states self.set_state(states)
def show_depth(): # render rgb and depth img, depth = pyflex.render() img = img.reshape((720, 720, 4))[::-1, :, :3] depth = depth.reshape((720, 720))[::-1] # get foreground mask rgb, depth = pyflex.render_cloth() depth = depth.reshape(720, 720)[::-1] # mask = mask[:, :, 3] # depth[mask == 0] = 0 # show rgb and depth(masked) depth[depth > 5] = 0 fig, axes = plt.subplots(1, 2, figsize=(12, 5)) axes[0].imshow(img) axes[1].imshow(depth) plt.show()
def _get_center_point(self, pos): pos = np.reshape(pos, [-1, 4]) min_x = np.min(pos[:, 0]) min_y = np.min(pos[:, 2]) max_x = np.max(pos[:, 0]) max_y = np.max(pos[:, 2]) return 0.5 * (min_x + max_x), 0.5 * (min_y + max_y) if __name__ == '__main__': env = RopeNewEnv(observation_mode='key_point', action_mode='picker', num_picker=2, render=True, headless=False, horizon=75, action_repeat=8, num_variations=10, use_cached_states=False, save_cached_states=False, deterministic=False) env.reset(config=env.get_default_config()) for i in range(1000): print(i) print("right before pyflex step") pyflex.step() print("right after pyflex step") print("right before pyflex render") pyflex.render() print("right after pyflex render")
v_box += rand_float(-0.15, 0.15) - x_box * 0.1 shape_states_ = calc_shape_states(x_box, x_box_last, scene_params[-2:]) pyflex.set_shape_states(shape_states_) positions[i] = pyflex.get_positions().reshape(-1, dim_position) velocities[i] = pyflex.get_velocities().reshape(-1, dim_velocity) shape_states[i] = pyflex.get_shape_states().reshape(-1, dim_shape_state) if i == 0: print(np.min(positions[i], 0), np.max(positions[i], 0)) print(x_box, box_dis_x, box_dis_z) pyflex.step() # playback pyflex.set_scene(6, scene_params, 0) for i in range(len(boxes) - 1): halfEdge = boxes[i][0] center = boxes[i][1] quat = boxes[i][2] pyflex.add_box(halfEdge, center, quat) for i in range(time_step): pyflex.set_positions(positions[i]) pyflex.set_shape_states(shape_states[i, :-1]) pyflex.render(capture=1, path=os.path.join(des_dir, 'render_%d.tga' % i)) pyflex.clean()
def gen_Cloth(info): env, env_idx = info['env'], info['env_idx'] thread_idx, data_dir, data_names = info['thread_idx'], info['data_dir'], info['data_names'] n_rollout, time_step = info['n_rollout'], info['time_step'] dt, args, phase = info['dt'], info['args'], info['phase'] vis_width, vis_height = info['vis_width'], info['vis_height'] state_dim = args.state_dim action_dim = args.action_dim dt = 1. / 60. np.random.seed(round(time.time() * 1000 + thread_idx) % 2 ** 32) stats = [init_stat(state_dim), init_stat(action_dim)] engine = ClothEngine(dt, state_dim, action_dim) import pyflex pyflex.init() # bar = ProgressBar() for i in range(n_rollout): rollout_idx = thread_idx * n_rollout + i rollout_dir = os.path.join(data_dir, str(rollout_idx)) os.system('mkdir -p ' + rollout_dir) engine.init(pyflex) scene_params = engine.scene_params action = np.zeros(4) states_all = np.zeros((time_step, engine.n_particles, state_dim)) actions_all = np.zeros((time_step, 1, action_dim)) # drop the cloth down engine.set_action(action) engine.step() for j in range(time_step): positions = pyflex.get_positions().reshape(-1, 4)[:, :3] # sample the action if j % 5 == 0: ctrl_pts = rand_int(0, 8) act_lim = 0.05 dx = rand_float(-act_lim, act_lim) dz = rand_float(-act_lim, act_lim) dy = 0.05 action = np.array([ctrl_pts, dx, dy, dz]) else: action[2] = 0. # store the rollout information state = engine.get_state() states_all[j] = state tga_path = os.path.join(rollout_dir, '%d.tga' % j) pyflex.render(capture=True, path=tga_path) tga = Image.open(tga_path) tga = np.array(tga)[:, 60:780, :3][:, :, ::-1] tga = cv2.resize(tga, (vis_width, vis_height), interpolation=cv2.INTER_AREA) os.system('rm ' + tga_path) jpg_path = os.path.join(rollout_dir, 'fig_%d.jpg' % j) cv2.imwrite(jpg_path, tga) actions_all[j, 0] = action.copy() engine.set_action(action) engine.step() datas = [states_all, actions_all, scene_params] store_data(data_names, datas, rollout_dir + '.h5') datas = [datas[j].astype(np.float64) for j in range(len(datas))] for j in range(len(stats)): stat = init_stat(stats[j].shape[0]) stat[:, 0] = np.mean(datas[j], axis=(0, 1))[:] stat[:, 1] = np.std(datas[j], axis=(0, 1))[:] stat[:, 2] = datas[j].shape[0] stats[j] = combine_stat(stats[j], stat) pyflex.clean() return stats
print(pyflex.get_sceneParams()) pyflex.step() for i in range(time_step - 1): idx = 10 cnt = 0 for j in range(n_rigids): st, ed = rigid_offsets[r, i, j, 0], rigid_offsets[r, i, j + 1, 0] for k in range(st, ed): if rigid_indices[r, i, cnt, 0] == idx: print(i, j, positions[r, i, rigid_indices[r, i, cnt, 0], :3], rigid_globalPositions[r, i, rigid_indices[r, i, cnt, 0]], rotate(rigid_localPositions[r, i, cnt], rigid_rotations[r, i, j]) + \ rigid_translations[r, i, j]) cnt += 1 pyflex.set_scene(5, scene_params, 0) pyflex.add_box(halfEdge, center, quat) pyflex.add_box(halfEdge, center, quat) for r in range(grip_time): for i in range(time_step): pyflex.set_positions(positions[r, i]) pyflex.set_shape_states(shape_states[r, i]) pyflex.render(capture=1, path=os.path.join(des_dir, 'render_%d.tga' % (r * time_step + i))) pyflex.clean()
# Set all the particles to unseen locations def clear_pos(): pos[:, 0:3] = np.ones((num_particles, 3)) * 20 # Function for putting particle positions into the array def load_pos(particles): print(particles.shape[0]) for i in range(particles.shape[0]): if i < num_particles: pos[i, 0:3] = particles[i, 0:3] # # Load in the particles # path = particle_path_prefix+"frame_{}.npy".format(599) # particles = np.load(path) # load_pos(particles) # pyflex.set_positions(pos) for i in range(time_step): # Load in the particles path = particle_path_prefix + "frame_{}.npy".format(i) particles = np.load(path) / scale clear_pos() load_pos(particles) pyflex.set_positions(pos) # pyflex.step(capture=1, path=os.path.join(des_dir, "render_%04d.tga" % i)) pyflex.render(capture=1, path=os.path.join(des_dir, "render_%04d.tga" % i)) pyflex.clean()
for step in range(args.time_step - 1): if args.env == 'RiceGrip': pyflex.set_shape_states(s_gt[step]) elif args.env == 'FluidShake': pyflex.set_shape_states(s_gt[step, :-1]) mass = np.zeros((n_particles, 1)) if args.env == 'RiceGrip': p = np.concatenate([p_gt[step, :n_particles, -3:], mass], 1) else: p = np.concatenate([p_gt[step, :n_particles], mass], 1) pyflex.set_positions(p) pyflex.render(capture=1, path=os.path.join(des_dir, 'gt_%d.tga' % step)) ##### render for the predictions pyflex.set_scene(env_idx, scene_params, 0) if args.env == 'RiceGrip': pyflex.add_box(halfEdge, center, quat) pyflex.add_box(halfEdge, center, quat) elif args.env == 'FluidShake': for box_idx in range(len(boxes) - 1): halfEdge = boxes[box_idx][0] center = boxes[box_idx][1] quat = boxes[box_idx][2] pyflex.add_box(halfEdge, center, quat) for step in range(args.time_step - 1):