def forward(self, image_g, pose): self.set_map(image_g, None) image_r, _ = self.get_map(pose) presenter = Presenter() presenter.show_image(image_g[0].data, "img_g", torch=True, waitkey=False, scale=2) presenter.show_image(image_r[0].data, "img_r", torch=True, waitkey=100, scale=2) features_r = self.feature_net(image_r) coverage = torch.ones_like(features_r) return features_r, coverage, image_r
def browse_pvn_dataset(): P.initialize_experiment() setup = P.get_current_parameters()["Setup"] model_sim, _ = load_model(setup["model"], setup["sim_model_file"], domain="sim") data_params = P.get_current_parameters()["Training"] print("Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() #dom="real" dom = "sim" dataset = model_sim.get_dataset( data=None, envs=train_envs, domain=dom, dataset_names=data_params[f"{dom}_dataset_names"], dataset_prefix="supervised", eval=False, halfway_only=False) p = Presenter() for example in dataset: if example is None: continue md = example["md"][0] print( f"Showing example: {md['env_id']}:{md['set_idx']}:{md['seg_idx']}") print(f" instruction: {md['instruction']}") exec_len = len(example["images"]) for i in range(exec_len): print(f" timestep: {i}") img_i = example["images"][i] lm_fpv_i = example["lm_pos_fpv"][i] if lm_fpv_i is not None: img_i = p.plot_pts_on_torch_image(img_i, lm_fpv_i.long()) p.show_image(img_i, "fpv_img_i", scale=4, waitkey=True)
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("")
env = PomdpInterface() count = 0 presenter = Presenter() def save_landmark_img(state, landmark_name, i, eval): data_dir = get_landmark_images_dir(landmark_name, eval) os.makedirs(data_dir, exist_ok=True) full_path = os.path.join(data_dir, landmark_name + "_" + str(i) + ".jpg") scipy.misc.imsave(full_path, state.image) for landmark_name, landmark_radius in LANDMARK_RADII.items(): for i in range(IMAGES_PER_LANDMARK_TRAIN): print("Generating train image " + str(i) + " for landmark '" + landmark_name + "'") state = env.reset_to_random_cv_env(landmark_name) presenter.show_image(state.image) save_landmark_img(state, landmark_name, i, False) for i in range(IMAGES_PER_LANDMARK_TEST): print("Generating test image " + str(i) + " for landmark '" + landmark_name + "'") state = env.reset_to_random_cv_env(landmark_name) presenter.show_image(state.image) save_landmark_img(state, landmark_name, i, True)
def show_landmark_locations(self, loop=True, states=None): # Show landmark locations in first-person images img_all = self.tensor_store.get("images") img_w_all = self.tensor_store.get("images_w") import rollout.run_metadata as md if md.IS_ROLLOUT: # TODO: Discard this and move this to PomdpInterface or something # (it's got nothing to do with the model) # load landmark positions from configs from data_io.env import load_env_config from learning.datasets.aux_data_providers import get_landmark_locations_airsim from learning.models.semantic_map.pinhole_camera_inv import PinholeCameraProjection projector = PinholeCameraProjection( map_size_px=self.params["global_map_size"], world_size_px=self.params["world_size_px"], world_size_m=self.params["world_size_m"], img_x=self.params["img_w"], img_y=self.params["img_h"], cam_fov=self.params["cam_h_fov"], #TODO: Handle correctly domain="sim", use_depth=False) conf_json = load_env_config(md.ENV_ID) landmark_names, landmark_indices, landmark_pos = get_landmark_locations_airsim( conf_json) cam_poses = self.cam_poses_from_states(states) cam_pos = cam_poses.position[0] cam_rot = cam_poses.orientation[0] lm_pos_map_all = [] lm_pos_img_all = [] for i, landmark_in_world in enumerate(landmark_pos): lm_pos_img, landmark_in_cam, status = projector.world_point_to_image( cam_pos, cam_rot, landmark_in_world) lm_pos_map = torch.from_numpy( transformations.pos_m_to_px( landmark_in_world[np.newaxis, :], self.params["global_map_size"], self.params["world_size_m"], self.params["world_size_px"])) lm_pos_map_all += [lm_pos_map[0]] if lm_pos_img is not None: lm_pos_img_all += [lm_pos_img] lm_pos_img_all = [lm_pos_img_all] lm_pos_map_all = [lm_pos_map_all] else: lm_pos_img_all = self.tensor_store.get("lm_pos_fpv_img") lm_pos_map_all = self.tensor_store.get("lm_pos_map") print("Plotting landmark points") for i in range(len(img_all)): p = Presenter() overlay_fpv = p.overlay_pts_on_image(img_all[i][0], lm_pos_img_all[i]) overlay_map = p.overlay_pts_on_image(img_w_all[i][0], lm_pos_map_all[i]) p.show_image(overlay_fpv, "landmarks_on_fpv_img", scale=8) p.show_image(overlay_map, "landmarks_on_map", scale=20) if not loop: break
def sup_loss_on_batch(self, batch, eval=False, viz=False): if eval: self.eval() else: self.train() images = cuda_var(batch["images"], self.is_cuda, self.cuda_device) instructions = cuda_var(batch["instr"], self.is_cuda, self.cuda_device) instruction_masks = cuda_var(batch["instr_mask"], self.is_cuda, self.cuda_device) label_masks = cuda_var(batch["traj_labels"], self.is_cuda, self.cuda_device) # Each of the above is a list of lists of tensors, where the outer list is over the batch and the inner list # is over the segments. Loop through and accumulate loss for each batch sequentially, and for each segment. # Reset model state (embedding etc) between batches, but not between segments. # We don't process each batch in batch-mode, because it's complicated, with the varying number of segments and all. batch_size = len(images) total_class_loss = Variable(empty_float_tensor([1], self.is_cuda, self.cuda_device), requires_grad=True) total_ground_loss = Variable(empty_float_tensor([1], self.is_cuda, self.cuda_device), requires_grad=True) count = 0 label_masks = self.label_pool(label_masks) mask_pred, features, emb_loss = self(images, instructions, instruction_masks) if BCE: mask_pred_flat = mask_pred.view(-1, 1) label_masks_flat = label_masks - torch.min(label_masks) label_masks_flat = label_masks_flat / ( torch.max(label_masks_flat) + 1e-9) label_masks_flat = label_masks_flat.view(-1, 1).clamp(0, 1) main_loss = self.mask_loss(mask_pred_flat, label_masks_flat) elif NLL: mask_pred_1 = F.softmax(mask_pred, 1, _stacklevel=5) mask_pred_2 = 1 - mask_pred_1 mask_pred_1 = mask_pred_1.unsqueeze(1) mask_pred_2 = mask_pred_2.unsqueeze(1) mask_pred = torch.cat((mask_pred_1, mask_pred_2), dim=1) label_masks = label_masks.clamp(0, 1) if self.is_cuda: label_masks = label_masks.type(torch.cuda.LongTensor) else: label_masks = label_masks.type(torch.LongTensor) main_loss = self.mask_loss(mask_pred, label_masks) elif CE: # Crossentropy2D internally applies logsoftmax to mask_pred, # but labels are already assumed to be a valid probability distribution, so no softmax is applied main_loss = self.mask_loss(mask_pred, label_masks) # So for nice plotting, we must manually do it mask_pred = self.spatialsoftmax(mask_pred) else: main_loss = self.mask_loss(mask_pred, label_masks) # sum emb loss if batch size > 1 if type(emb_loss) == tuple: emb_loss = sum(emb_loss) # Extract the feature vectors corresponding to every landmark's location in the map # Apply a linear layer to classify which of the 64 landmarks it is # The landmark positions have to be divided by the same factor as the ResNet scaling factor lcount = 0 for i in range(batch_size): if self.class_loss and len(batch["lm_pos"][i]) > 0: lcount += 1 landmark_pos = cuda_var(batch["lm_pos"][i], self.is_cuda, self.cuda_device) landmark_indices = cuda_var(batch["lm_indices"][i], self.is_cuda, self.cuda_device) landmark_coords = (landmark_pos / 8).long() lm_features = self.gather2d(features[i:i + 1, 0:32], landmark_coords) lm_pred = self.aux_class_linear(lm_features) class_loss = self.aux_loss(lm_pred, landmark_indices) total_class_loss = total_class_loss + class_loss if self.ground_loss and len(batch["lm_pos"][i]) > 0: landmark_pos = cuda_var(batch["lm_pos"][i], self.is_cuda, self.cuda_device) landmark_mentioned = cuda_var(batch["lm_mentioned"][i], self.is_cuda, self.cuda_device) landmark_coords = (landmark_pos / 8).long() g_features = self.gather2d(features[i:i + 1, 32:35], landmark_coords) lm_pred = self.aux_ground_linear(g_features) ground_loss = self.aux_loss(lm_pred, landmark_mentioned) total_ground_loss = total_ground_loss + ground_loss total_class_loss = total_class_loss / (lcount + 1e-9) total_ground_loss = total_ground_loss / (lcount + 1e-9) count += 1 # Just visualization and debugging code if self.get_iter() % 50 == 0: presenter = Presenter() pred_viz_np = presenter.overlaid_image(images[0].data, mask_pred[0].data) labl_viz_np = presenter.overlaid_image(images[0].data, label_masks[0].data) comp = np.concatenate((pred_viz_np, labl_viz_np), axis=1) presenter.show_image(comp, "path_pred") if hasattr(self.sentence_embedding, "save_att_map"): self.sentence_embedding.save_att_map(self.get_iter(), i) total_loss = main_loss + 0.1 * total_class_loss + 0.001 * emb_loss + 0.1 * total_ground_loss total_loss = total_loss / (count + 1e-9) self.write_summaires("eval" if eval else "train", self.get_iter(), total_loss, main_loss, emb_loss, total_class_loss, total_ground_loss) self.inc_iter() return total_loss
class RolloutVisualizer: def __init__(self, resolution=512): self.presenter = Presenter() self.clear() self.current_rollout = {} self.current_rollout_name = None self.env_image = None self.current_timestep = None self.world_size_m = P.get_current_parameters()["Setup"]["world_size_m"] self.resolution = resolution def clear(self): self.current_rollout = {} self.current_rollout_name = None self.env_image = None self.current_timestep = None def _auto_contrast(self, image): import cv2 image_c = np.clip(image, 0.0, 1.0) hsv_image = cv2.cvtColor(image_c, cv2.COLOR_RGB2HSV) hsv_image[:, :, 1] *= 1.2 image_out = cv2.cvtColor(hsv_image, cv2.COLOR_HSV2RGB) image_out = np.clip(image_out, 0.0, 1.0) #print(image_out.min(), image_out.) print(image_out[:, :, 1].min(), image_out[:, :, 1].max()) return image def _integrate_mask(self, frames): frames_out = [frames[0]] for frame in frames[1:]: new_frame_out = np.maximum(frames_out[-1], frame) frames_out.append(new_frame_out) return frames_out def _draw_landmarks(self, image, env_id): lm_names, lm_idx, lm_pos = get_landmark_locations_airsim(env_id=env_id) image = self.presenter.draw_landmarks(image, lm_names, lm_pos, self.world_size_m) return image def load_video_clip(self, env_id, seg_idx, rollout, domain, cam_name, rollout_dir): video_path = os.path.join( rollout_dir, f"rollout_{cam_name}_{env_id}-0-{seg_idx}.mkv") try: #if os.path.getsize(video_path) > 1024 * 1024 * 30 print("Loading video: ", video_path) clip = mpy.VideoFileClip(video_path) except Exception as e: return None return clip def grab_frames(self, env_id, seg_idx, rollout, domain, frame_name, scale=1): frames = [] for sample in rollout: if frame_name == "image": frame = sample["state"].image elif frame_name == "action": action = sample["action"] bg = np.zeros((400, 400, 3)) frame = self.presenter.draw_action(bg, offset=(0, 0), action=action) elif frame_name == "v_dist_r_inner": frame_t = sample[frame_name][:3, :, :].transpose((1, 2, 0)) # TODO: These should come from params map_size = 64 crop_size = 16 gap = int((map_size - crop_size) / 2) crop_l = gap crop_r = map_size - gap frame_t = frame_t[crop_l:crop_r, crop_l:crop_r, :] frame_t[:, :, 0] /= (np.percentile(frame_t[:, :, 0], 99) + 1e-9) frame_t[:, :, 1] /= (np.percentile(frame_t[:, :, 1], 99) + 1e-9) frame_t = np.clip(frame_t, 0.0, 1.0) shp = list(frame_t.shape) shp[2] = 3 frame = np.zeros(shp) frame[:, :, :2] = frame_t frame = cv2.resize(frame, dsize=(self.resolution, self.resolution)) elif frame_name == "map_struct": frame_t = sample[frame_name][:3, :, :].transpose((1, 2, 0)) shp = list(frame_t.shape) shp[2] = 3 frame = np.zeros(shp) frame[:, :, :2] = frame_t frame = cv2.resize(frame, dsize=(self.resolution, self.resolution)) elif frame_name == "ego_obs_mask": frame_t = sample["map_struct"][:3, :, :].transpose((1, 2, 0)) shp = list(frame_t.shape) shp[2] = 3 canvas = np.zeros(shp) canvas[:, :, :] = 1 - frame_t[:, :, 0:1] canvas[:, :, :] -= frame_t[:, :, 1:2] canvas = np.clip(canvas, 0.0, 1.0) frame = cv2.resize(canvas, dsize=(self.resolution, self.resolution)) else: frame = sample[frame_name][0, :3, :, :].transpose((1, 2, 0)) if frame_name in ["image", "v_dist_r_inner"]: frame -= frame.min() frame = frame / (frame.max() + 1e-9) else: frame -= np.percentile(frame, 0) frame /= (np.percentile(frame, 95) + 1e-9) frame = np.clip(frame, 0.0, 1.0) if scale != 1: frame = self.presenter.scale_image(frame, scale) frames.append(frame) return frames def action_visualization(self, env_id, seg_idx, rollout, domain, frame_name="action"): frames = [] for sample in rollout: action = sample[frame_name] frame = np.ones((200, 200, 3), dtype=np.uint8) self.presenter.draw_action(frame, (1, 159), action) frames.append(frame) return frames def overlay_frames(self, under_frames, over_frames, strength=0.5): overlaid_frames = [ self.presenter.overlaid_image(u, o, strength=strength) for u, o in zip(under_frames, over_frames) ] return overlaid_frames def top_down_visualization(self, env_id, seg_idx, rollout, domain, params): fd = domain == "real" obl = domain in ["simulator", "sim"] print(domain, obl) if params["draw_topdown"]: bg_image = load_env_img( env_id, self.resolution, self.resolution, real_drone=True if domain == "real" else False, origin_bottom_left=obl, flipdiag=False, alpha=True) else: bg_image = np.zeros((self.resolution, self.resolution, 3)) if params["draw_landmarks"]: bg_image = self._draw_landmarks(bg_image, env_id) # Initialize stuff frames = [] poses_m = [] poses_px = [] for sample in rollout: sample_image = bg_image.copy() frames.append(sample_image) state = sample["state"] pose_m = state.get_drone_pose() pose_px = poses_m_to_px(pose_m, self.resolution, self.resolution, self.world_size_m, batch_dim=False) poses_px.append(pose_px) poses_m.append(pose_m) instruction = rollout[0]["instruction"] print("Instruction: ") print(instruction) # Draw visitation distributions if requested: if params["include_vdist"]: print("Drawing visitation distributions") if params["ego_vdist"]: inner_key = "v_dist_r_inner" outer_key = "v_dist_r_outer" else: inner_key = "v_dist_w_inner" outer_key = "v_dist_w_outer" for i, sample in enumerate(rollout): v_dist_w_inner = np.flipud(sample[inner_key].transpose( (2, 1, 0))) # Expand range of each channel separately so that stop entropy doesn't affect how trajectory looks v_dist_w_inner[:, :, 0] /= ( np.percentile(v_dist_w_inner[:, :, 0], 99.5) + 1e-9) v_dist_w_inner[:, :, 1] /= ( np.percentile(v_dist_w_inner[:, :, 1], 99.5) + 1e-9) v_dist_w_inner = np.clip(v_dist_w_inner, 0.0, 1.0) v_dist_w_outer = sample[outer_key] if bg_image.max() - bg_image.min() > 1e-9: f = self.presenter.blend_image(frames[i], v_dist_w_inner) else: f = self.presenter.overlaid_image(frames[i], v_dist_w_inner, strength=1.0) f = self.presenter.draw_prob_bars(f, v_dist_w_outer) frames[i] = f if params["include_layer"]: layer_name = params["include_layer"] print(f"Drawing first 3 channels of layer {layer_name}") accumulate = False invert = False gray = False if layer_name == "M_W_accum": accumulate = True layer_name = "M_W" if layer_name == "M_W_accum_inv": invert = True accumulate = True layer_name = "M_W" if layer_name.endswith("_Gray"): gray = True layer_name = layer_name[:-len("_Gray")] for i, sample in enumerate(rollout): layer = sample[layer_name] if len(layer.shape) == 4: layer = layer[0, :, :, :] layer = layer.transpose((2, 1, 0)) layer = np.flipud(layer) if layer_name in ["S_W", "F_W"]: layer = layer[:, :, :3] else: layer = layer[:, :, :3] if layer_name in ["S_W", "R_W", "F_W"]: if gray: layer -= np.percentile(layer, 1) layer /= (np.percentile(layer, 99) + 1e-9) else: layer /= (np.percentile(layer, 97) + 1e-9) layer = np.clip(layer, 0.0, 1.0) if layer_name in ["M_W"]: # Having a 0-1 mask does not encode properly with the codec. Add a bit of imperceptible gaussian noise. layer = layer.astype(np.float32) layer = np.tile(layer, (1, 1, 3)) if accumulate and i > 0: layer = np.maximum(layer, prev_layer) prev_layer = layer if invert: layer = 1 - layer if frames[i].max() > 0.01: frames[i] = self.presenter.blend_image( frames[i], layer[:, :, :3]) #frames[i] = self.presenter.overlaid_image(frames[i], layer[:, :, :3]) else: scale = (int(self.resolution / layer.shape[0]), int(self.resolution / layer.shape[1])) frames[i] = self.presenter.prep_image(layer[:, :, :3], scale=scale) if params["include_instr"]: print("Drawing instruction") for i, sample in enumerate(rollout): frames[i] = self.presenter.overlay_text( frames[i], sample["instruction"]) # Draw trajectory history if params["draw_trajectory"]: print("Drawing trajectory") for i, sample in enumerate(rollout): history = poses_px[:i + 1] position_history = [h.position for h in history] frames[i] = self.presenter.draw_trajectory( frames[i], position_history, self.world_size_m) # Draw drone if params["draw_drone"]: print("Drawing drone") for i, sample in enumerate(rollout): frames[i] = self.presenter.draw_drone(frames[i], poses_m[i], self.world_size_m) # Draw observability mask: if params["draw_fov"]: print("Drawing FOV") for i, sample in enumerate(rollout): frames[i] = self.presenter.draw_observability( frames[i], poses_m[i], self.world_size_m, 84) # Visualize if False: for i, sample in enumerate(rollout): self.presenter.show_image(frames[i], "sample_image", scale=1, waitkey=True) return frames def start_rollout(self, env_id, set_idx, seg_idx, domain, dataset, suffix=""): rollout_name = f"{env_id}:{set_idx}:{seg_idx}:{domain}:{dataset}:{suffix}" self.current_rollout = {"top-down": []} self.current_rollout_name = rollout_name self.env_image = load_env_img(512, 512, alpha=True) def start_timestep(self, timestep): self.current_timestep = timestep # Add top-down view image for the new timestep self.current_rollout["top-down"].append() self.current_rollout["top-down"].append(self.env_image.copy()) def set_drone_state(self, timestep, state): drone_pose = state.get_cam_pose() # Draw drone sprite on top_down image tdimg = self.current_rollout["top-down"][timestep] tdimg_n = self.presenter.draw_drone( tdimg, drone_pose, P.get_current_parameters()["Setup"]["world_size_m"]) self.current_rollout["top-down"][timestep] = tdimg_n