scene_type=scene_name,
start_scene_number=start_scene_number)
object_states = []
for _ in range(1):
env.reset(random_init=False)
# print(env.current_scene, env.scene_config['answer'], len(env.scene_config['goal']['action_list']))
scene_state = None
for i, x in enumerate(env.scene_config['goal']['action_list']):
# print(i)
if i == 0:
scene_state = SceneState(env.step_output)
else:
scene_state.update(env.step_output)
env.step(action=x[0])
object_states.append(scene_state)
env.controller.end_scene(None, None)
for i, scene_state in enumerate(object_states):
plt.figure()
for j, (id, obj_state) in enumerate(scene_state.object_state_dict.items()):
v_xs = [v[0] for v in obj_state.velocity_history]
v_ys = [v[1] for v in obj_state.velocity_history]
plt.scatter(v_xs,
v_ys,
label="Object {}".format(j),
alpha=0.3,
edgecolors='none')
h_t_1 = torch.zeros(size=(1, 1, HIDDEN_STATE_SIZE))
print(len(env_1.scene_config['goal']['action_list']))
if len(env_1.scene_config['goal']['action_list']) == 40:
for _ in range(20):
f_none = get_locomotion_feature(obj=None,
object_occluded=False,
object_in_scene=False)
f_none = torch.tensor(f_none)
f_none = f_none.unsqueeze(0).unsqueeze(0)
_, h_t_1 = net((f_none, h_t_1))
obj_in_scene = False
for i, action in enumerate(env_1.scene_config['goal']['action_list']):
env_1.step(action=action[0])
obj_in_view = None
if len(env_1.step_output.object_list) == 1:
f_1 = get_locomotion_feature(env_1.step_output.object_list[0],
object_occluded=False,
object_in_scene=True)
obj_in_view = True
obj_in_scene = True
elif len(env_1.step_output.object_list) == 0:
f_1 = get_locomotion_feature(
None,
object_occluded=True,
object_in_scene=(obj_in_scene > 0))
obj_in_view = False
scene_name = "github_scenes" + "/collect_object_shape_data"
for _, shape_type in enumerate(SHAPE_TYPES):
print("Scene: {}".format(shape_type))
os.makedirs(os.path.join("appearance", "object_mask_frame", shape_type), exist_ok=True)
env = McsEnv(task="intphys_scenes", scene_type=scene_name)
env.reset(random_init=False)
env.scene_config['objects'][0]['type'] = shape_type
env.step_output = env.controller.start_scene(env.scene_config)
object_frames = []
for scale in [0.2 + 0.1*i for i in range(10)]:
set_scale(env.scene_config, scale)
env.step_output = env.controller.start_scene(env.scene_config)
for i, action in enumerate(env.scene_config['goal']['action_list']):
env.step(action=action[0])
assert len(env.step_output.object_list) <= 1
if len(env.step_output.object_list) == 1:
obj_state = ObjectState(
env.step_output.object_list[0], env.step_output.depth_mask_list[-1], env.step_output.object_mask_list[-1]
)
x_s = obj_state.edge_pixels['x_min']
x_e = obj_state.edge_pixels['x_max']
y_s = obj_state.edge_pixels['y_min']
y_e = obj_state.edge_pixels['y_max']
new_size = (IMAGE_CROP_SIZE, IMAGE_CROP_SIZE)
obj_frame = np.array(env.step_output.object_mask_list[-1])
pixels_on_frame = get_object_match_pixels(env.step_output.object_list[0].color, obj_frame)
obj_frame[:,:] = [255, 255, 255]
for x,y in pixels_on_frame:
obj_frame[x,y,:] = [0, 0, 0]
}
return return_dict
scene_err = scene_steps = scene_path_err = scene_voe = 0
num_trials = 9
trial_info = None
for idx in range(num_trials):
env.reset(random_init=False)
path_taken = []
# _ret, cam_im = color_cap.read()
# _ret, mask_im = mask_cap.read()
first_action = env.scene_config['goal']['action_list'][0][0]
env.step(action=first_action)
cam_im = env.step_output.image_list[0]
cam_im = np.array(cam_im)[:, :, ::-1]
mask_im = env.step_output.object_mask_list[0]
mask_im = np.array(mask_im)[:, :, ::-1]
init_info_dict = {
"trial_num": idx,
"trial_err": 0,
"step_num": 0,
"arena": None,
"agent_ch_avgs": None,
"objs_ch_avgs": None,
"trial_1_objs_pos": None,
"path": None,
"wall_i_s": None