def __init__(self, map_size, world_size_px, world_size_m):
super(MapAffine, self).__init__()
self.map_size = map_size
self.world_size_px = world_size_px
self.world_size_m = world_size_m
self.affine_2d = Affine2D()
self.prof = SimpleProfiler(torch_sync=PROFILE, print=PROFILE)
def __init__(self, map_size=MAP_SIZE, world_in_map_size=MAP_WORLD_SIZE):
super(MapAffine, self).__init__()
self.is_cuda = False
self.cuda_device = None
self.map_size = map_size
self.world_in_map_size = world_in_map_size
self.affine_2d = Affine2D()
self.prof = SimpleProfiler(torch_sync=PROFILE, print=PROFILE)
def __init__(self, source_map_size, dest_map_size, world_size_px,
world_size_m):
super(MapAffine, self).__init__()
self.source_map_size_px = source_map_size
self.dest_map_size_px = dest_map_size
self.world_in_map_size_px = world_size_px
self.world_size_m = world_size_m
self.affine_2d = Affine2D()
self.prof = SimpleProfiler(torch_sync=PROFILE, print=PROFILE)
pos = np.asarray(
[self.source_map_size_px / 2, self.source_map_size_px / 2])
rot = np.asarray([0])
self.canonical_pose_src = Pose(pos, rot)
pos = np.asarray(
[self.dest_map_size_px / 2, self.dest_map_size_px / 2])
rot = np.asarray([0])
self.canonical_pose_dst = Pose(pos, rot)
def train_top_down_pred():
P.initialize_experiment()
setup = P.get_current_parameters()["Setup"]
launch_ui()
env = PomdpInterface()
print("model_name:", setup["top_down_model"])
print("model_file:", setup["top_down_model_file"])
model, model_loaded = load_model(
model_name_override=setup["top_down_model"],
model_file_override=setup["top_down_model_file"])
exec_model, wrapper_model_loaded = load_model(
model_name_override=setup["wrapper_model"],
model_file_override=setup["wrapper_model_file"])
affine2d = Affine2D()
if model.is_cuda:
affine2d.cuda()
eval_envs = get_correct_eval_env_id_list()
print("eval_envs:", eval_envs)
train_instructions, dev_instructions, test_instructions, corpus = get_all_instructions(
max_size=setup["max_envs"])
all_instr = {
**train_instructions,
**dev_instructions,
**train_instructions
}
token2term, word2token = get_word_to_token_map(corpus)
dataset = model.get_dataset(envs=eval_envs,
dataset_name="supervised",
eval=True,
seg_level=False)
dataloader = DataLoader(dataset,
collate_fn=dataset.collate_fn,
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True)
for b, batch in list(enumerate(dataloader)):
print("batch:", batch)
images = batch["images"]
instructions = batch["instr"]
label_masks = batch["traj_labels"]
affines = batch["affines_g_to_s"]
env_ids = batch["env_id"]
set_idxs = batch["set_idx"]
seg_idxs = batch["seg_idx"]
env_id = env_ids[0][0]
set_idx = set_idxs[0][0]
print("env_id of this batch:", env_id)
env.set_environment(
env_id, instruction_set=all_instr[env_id][set_idx]["instructions"])
env.reset(0)
num_segments = len(instructions[0])
print("num_segments in this batch:", num_segments)
write_instruction("")
write_real_instruction("None")
instruction_str = read_instruction_file()
print("Initial instruction: ", instruction_str)
# TODO: Reset model state here if we keep any temporal memory etc
for s in range(num_segments):
start_state = env.reset(s)
keep_going = True
real_instruction = cuda_var(instructions[0][s], setup["cuda"], 0)
tmp = list(real_instruction.data.cpu()[0].numpy())
real_instruction_str = debug_untokenize_instruction(tmp)
write_real_instruction(real_instruction_str)
#write_instruction(real_instruction_str)
#instruction_str = real_instruction_str
image = cuda_var(images[0][s], setup["cuda"], 0)
label_mask = cuda_var(label_masks[0][s], setup["cuda"], 0)
affine_g_to_s = affines[0][s]
print("Your current environment:")
with open(
"/storage/dxsun/unreal_config_nl/configs/configs/random_config_"
+ str(env_id) + ".json") as fp:
config = json.load(fp)
print(config)
while keep_going:
write_real_instruction(real_instruction_str)
while True:
cv2.waitKey(200)
instruction = read_instruction_file()
if instruction == "CMD: Next":
print("Advancing")
keep_going = False
write_empty_instruction()
break
elif instruction == "CMD: Reset":
print("Resetting")
env.reset(s)
write_empty_instruction()
elif len(instruction.split(" ")) > 1:
instruction_str = instruction
print("Executing: ", instruction_str)
break
if not keep_going:
continue
#instruction_str = read_instruction_file()
# TODO: Load instruction from file
tok_instruction = tokenize_instruction(instruction_str,
word2token)
instruction_t = torch.LongTensor(tok_instruction).unsqueeze(0)
instruction_v = cuda_var(instruction_t, setup["cuda"], 0)
instruction_mask = torch.ones_like(instruction_v)
tmp = list(instruction_t[0].numpy())
instruction_dbg_str = debug_untokenize_instruction(
tmp, token2term)
# import matplotlib.pyplot as plt
#plt.plot(image.squeeze(0).permute(1,2,0).cpu().numpy())
#plt.show()
res = model(image, instruction_v, instruction_mask)
mask_pred = res[0]
shp = mask_pred.shape
mask_pred = F.softmax(mask_pred.view([2, -1]), 1).view(shp)
#mask_pred = softmax2d(mask_pred)
# TODO: Rotate the mask_pred to the global frame
affine_s_to_g = np.linalg.inv(affine_g_to_s)
S = 8.0
affine_scale_up = np.asarray([[S, 0, 0], [0, S, 0], [0, 0, 1]])
affine_scale_down = np.linalg.inv(affine_scale_up)
affine_pred_to_g = np.dot(
affine_scale_down, np.dot(affine_s_to_g, affine_scale_up))
#affine_pred_to_g_t = torch.from_numpy(affine_pred_to_g).float()
mask_pred_np = mask_pred.data.cpu().numpy()[0].transpose(
1, 2, 0)
mask_pred_g_np = apply_affine(mask_pred_np, affine_pred_to_g,
32, 32)
print("Sum of global mask: ", mask_pred_g_np.sum())
mask_pred_g = torch.from_numpy(
mask_pred_g_np.transpose(2, 0,
1)).float()[np.newaxis, :, :, :]
exec_model.set_ground_truth_visitation_d(mask_pred_g)
# Create a batch axis for pytorch
#mask_pred_g = affine2d(mask_pred, affine_pred_to_g_t[np.newaxis, :, :])
mask_pred_np[:, :, 0] -= mask_pred_np[:, :, 0].min()
mask_pred_np[:, :, 0] /= (mask_pred_np[:, :, 0].max() + 1e-9)
mask_pred_np[:, :, 0] *= 2.0
mask_pred_np[:, :, 1] -= mask_pred_np[:, :, 1].min()
mask_pred_np[:, :, 1] /= (mask_pred_np[:, :, 1].max() + 1e-9)
presenter = Presenter()
presenter.show_image(mask_pred_g_np,
"mask_pred_g",
torch=False,
waitkey=1,
scale=4)
#import matplotlib.pyplot as plt
#print("image.data shape:", image.data.cpu().numpy().shape)
#plt.imshow(image.data.squeeze().permute(1,2,0).cpu().numpy())
#plt.show()
# presenter.show_image(image.data, "mask_pred_g", torch=False, waitkey=1, scale=4)
#import pdb; pdb.set_trace()
pred_viz_np = presenter.overlaid_image(image.data,
mask_pred_np,
channel=0)
# TODO: Don't show labels
# TODO: OpenCV colours
#label_mask_np = p.data.cpu().numpy()[0].transpose(1,2,0)
labl_viz_np = presenter.overlaid_image(image.data,
label_mask.data,
channel=0)
viz_img_np = np.concatenate((pred_viz_np, labl_viz_np), axis=1)
viz_img_np = pred_viz_np
viz_img = presenter.overlay_text(viz_img_np,
instruction_dbg_str)
cv2.imshow("interactive viz", viz_img)
cv2.waitKey(100)
rollout_model(exec_model, env, env_ids[0][s], set_idxs[0][s],
seg_idxs[0][s], tok_instruction)
write_instruction("")
def affine_2d_test():
img = load_env_img(2, 128, 128)
img = standardize_image(img)
img = torch.from_numpy(img).float().unsqueeze(0)
px = 64
py = 64
theta = 0.5
c = math.cos(theta)
s = math.sin(theta)
t_p = torch.FloatTensor([[1, 0, px], [0, 1, py], [0, 0, 1]]).unsqueeze(0)
t_r = torch.FloatTensor([[c, -s, 0], [s, c, 0], [0, 0, 1]]).unsqueeze(0)
mat_np = np.dot(t_p.squeeze().numpy(), t_r.squeeze().numpy())
mat_np_t = torch.from_numpy(mat_np).unsqueeze(0)
# For some forsaken reason rightmultiplying seems to mean applying the transformation second
mat = torch.bmm(t_p, t_r)
#mat1 = t_p
#mat2 = t_r
affine_2d = Affine2D()
res1 = affine_2d(Variable(img), Variable(t_r))
res2 = affine_2d(res1, Variable(t_p))
res3 = affine_2d(img, Variable(mat))
res4 = affine_2d(img, Variable(mat_np_t))
res3_big = affine_2d(img, Variable(mat), out_size=[512, 512])
res3_small = affine_2d(img, Variable(mat), out_size=[128, 128])
Presenter().show_image(res1.data[0],
"res_1",
torch=True,
waitkey=False,
scale=4)
Presenter().show_image(res2.data[0],
"res_2",
torch=True,
waitkey=False,
scale=4)
Presenter().show_image(res3.data[0],
"res_3",
torch=True,
waitkey=False,
scale=4)
Presenter().show_image(res3_big.data[0],
"res3_big",
torch=True,
waitkey=False,
scale=4)
Presenter().show_image(res3_small.data[0],
"res3_small",
torch=True,
waitkey=False,
scale=4)
Presenter().show_image(res4.data[0],
"res_4",
torch=True,
waitkey=True,
scale=4)
print("res2 should be the same as res_3 and res_4")