Spaces:
Running
on
Zero
Running
on
Zero
| ####################################################################### | |
| # Name: env.py | |
| # | |
| # - Reads and processes training and test maps (E.g. DungeonMaps) | |
| # - Processes rewards, new frontiers given action | |
| # - Updates a graph representation of environment for input into network | |
| ####################################################################### | |
| import sys | |
| import cv2 | |
| from matplotlib.colors import LogNorm, PowerNorm | |
| if sys.modules['TRAINING']: | |
| from parameter import * | |
| else: | |
| from test_parameter import * | |
| import copy | |
| import pandas as pd | |
| import rasterio | |
| from skimage import io | |
| import matplotlib.pyplot as plt | |
| import os | |
| from skimage.measure import block_reduce | |
| from sensor import * | |
| from graph_generator import * | |
| from node import * | |
| from scipy.ndimage import label, find_objects | |
| import matplotlib.image as mpimg | |
| from matplotlib.colors import Normalize | |
| # import matplotlib | |
| # matplotlib.use("Agg") # <-- key line to avoid tkinter dependency | |
| class Env(): | |
| def __init__(self, map_index, n_agent, k_size=20, plot=False, test=False, mask_index=None): | |
| self.n_agent = n_agent | |
| self.test = test | |
| self.map_dir = GRIDMAP_SET_DIR | |
| # Import environment gridmap | |
| self.map_list = os.listdir(self.map_dir) | |
| self.map_list.sort(reverse=True) | |
| # NEW: Import segmentation utility map | |
| self.seg_dir = MASK_SET_DIR | |
| self.segmentation_mask, self.target_positions, self.target_found_idxs = None, [], [] | |
| self.segmentation_mask_list = os.listdir(self.seg_dir) | |
| self.segmentation_mask_list.sort(reverse=True) | |
| # Import target maps (if relevant) | |
| if TARGETS_SET_DIR != "": | |
| self.targets_map_list = os.listdir(TARGETS_SET_DIR) | |
| self.targets_map_list.sort(reverse=True) | |
| # # NEW: Find common files in both directories | |
| # if TARGETS_SET_DIR != "": | |
| # common_files = [file for file in self.map_list if file in self.segmentation_mask_list and file in self.targets_map_list] | |
| # else: | |
| # common_files = [file for file in self.map_list if file in self.segmentation_mask_list] | |
| self.map_index = map_index % len(self.map_list) | |
| if mask_index is not None: | |
| self.mask_index = mask_index % len(self.segmentation_mask_list) | |
| else: | |
| self.mask_index = map_index % len(self.segmentation_mask_list) | |
| # self.common_map_file = common_files[self.map_index] | |
| # print("self.common_map_file: ", self.common_map_file) | |
| # Import ground truth and segmentation mask | |
| self.ground_truth, self.map_start_position = self.import_ground_truth( | |
| os.path.join(self.map_dir, self.map_list[self.map_index]))# self.common_map_file)) | |
| self.ground_truth_size = np.shape(self.ground_truth) # (480, 640) | |
| self.robot_belief = np.ones(self.ground_truth_size) * 127 # unexplored 127 | |
| self.downsampled_belief = None | |
| self.old_robot_belief = copy.deepcopy(self.robot_belief) | |
| self.coverage_belief = np.ones(self.ground_truth_size) * 127 # unexplored 127 | |
| # Import segmentation mask | |
| mask_filename = self.segmentation_mask_list[self.mask_index] | |
| self.segmentation_mask = self.import_segmentation_mask( | |
| os.path.join(self.seg_dir, mask_filename))# self.common_map_file)) | |
| # print("mask_filename: ", mask_filename) | |
| # Overwrite target positions if directory specified | |
| self.gt_segmentation_mask = None | |
| if self.test and TARGETS_SET_DIR != "": | |
| self.gt_segmentation_mask = self.import_segmentation_mask( | |
| os.path.join(TARGETS_SET_DIR, self.map_list[self.map_index])) # UNUSED - self.common_map_file)) | |
| # print("target_positions: ", self.target_positions) | |
| # print("np.unique(self.segmentation_mask): ", np.unique(self.segmentation_mask)) | |
| self.segmentation_info_mask = None | |
| self.gt_segmentation_info_mask = None | |
| self.segmentation_info_mask_unnormalized = None | |
| self.filtered_seg_info_mask = None | |
| self.num_targets_found = 0 | |
| self.num_new_targets_found = 0 | |
| # # Link score masks to raw image files | |
| # csv_file = pd.read_csv(RAW_IMG_PATH_DICT, header=None) | |
| # img_score_paths = csv_file.iloc[:,2].tolist() | |
| # raw_img_paths = csv_file.iloc[:,0].tolist() | |
| # self.score_to_img_dict = {os.path.basename(img_score_path): raw_img_path for img_score_path, raw_img_path in zip(img_score_paths, raw_img_paths)} | |
| self.resolution = 4 | |
| self.sensor_range = SENSOR_RANGE | |
| self.explored_rate = 0 | |
| self.targets_found_rate = 0 | |
| self.info_gain = 0 | |
| self.total_info = 0 | |
| self.graph_generator = Graph_generator(map_size=self.ground_truth_size, sensor_range=self.sensor_range, k_size=k_size, plot=plot) | |
| self.node_coords, self.graph, self.node_utility, self.guidepost = None, None, None, None | |
| self.frontiers = None | |
| self.start_positions = [] | |
| self.begin(self.map_start_position) | |
| self.plot = plot | |
| self.frame_files = [] | |
| def find_index_from_coords(self, position): | |
| index = np.argmin(np.linalg.norm(self.node_coords - position, axis=1)) | |
| return index | |
| def begin(self, start_position): | |
| # self.robot_belief = self.update_robot_belief(robot_position, self.sensor_range, self.robot_belief, | |
| # self.ground_truth) | |
| self.robot_belief = self.ground_truth | |
| self.downsampled_belief = block_reduce(self.robot_belief.copy(), block_size=(self.resolution, self.resolution), | |
| func=np.min) | |
| self.frontiers = self.find_frontier() | |
| self.old_robot_belief = copy.deepcopy(self.robot_belief) | |
| self.node_coords, self.graph, self.node_utility, self.guidepost = self.graph_generator.generate_graph( | |
| self.robot_belief, self.frontiers) | |
| # Find non-conflicting start positions | |
| if FIX_START_POSITION: | |
| coords_res_row = int(self.robot_belief.shape[0]/NUM_COORDS_HEIGHT) | |
| coords_res_col = int(self.robot_belief.shape[1]/NUM_COORDS_WIDTH) | |
| self.start_positions = [(int(self.robot_belief.shape[1]/2)-coords_res_col/2,int(self.robot_belief.shape[0]/2)-coords_res_row/2) for _ in range(self.n_agent)] # bottom-left corner | |
| else: | |
| nearby_coords = self.graph_generator.get_neighbors_grid_coords(start_position) | |
| itr = 0 | |
| for i in range(self.n_agent): | |
| if i == 0 or len(nearby_coords) == 0: | |
| self.start_positions.append(start_position) | |
| else: | |
| idx = min(itr, len(nearby_coords)-1) | |
| self.start_positions.append(nearby_coords[idx]) | |
| itr += 1 | |
| for i in range(len(self.start_positions)): | |
| self.start_positions[i] = self.node_coords[self.find_index_from_coords(self.start_positions[i])] | |
| self.coverage_belief = self.update_robot_belief(self.start_positions[i], self.sensor_range, self.coverage_belief, | |
| self.ground_truth) | |
| for start_position in self.start_positions: | |
| self.graph_generator.route_node.append(start_position) | |
| # Info map from ground truth | |
| rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH) | |
| rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT) | |
| self.segmentation_info_mask = np.zeros((len(self.node_coords), 1)) | |
| self.gt_segmentation_info_mask = np.zeros((len(self.node_coords), 1)) | |
| for i, node_coord in enumerate(self.node_coords): | |
| max_x = min(node_coord[0] + int(math.ceil(rng_x)), self.ground_truth.shape[1]) | |
| min_x = max(node_coord[0] - int(math.ceil(rng_x)), 0) | |
| max_y = min(node_coord[1] + int(math.ceil(rng_y)), self.ground_truth.shape[0]) | |
| min_y = max(node_coord[1] - int(math.ceil(rng_y)), 0) | |
| # if np.any(self.segmentation_mask[min_y:max_y, min_x:max_x] == 255): | |
| # self.segmentation_info_mask[i] = 1.0 | |
| # else: | |
| # self.segmentation_info_mask[i] = 0.0 | |
| # self.segmentation_info_mask[i] = np.mean(self.segmentation_mask[min_y:max_y, min_x:max_x]) | |
| # self.segmentation_info_mask[i] = np.max(self.segmentation_mask[min_y:max_y, min_x:max_x]) | |
| if TARGETS_SET_DIR == "": # If targets combined with segmentation mask | |
| exclude = {208} # Exclude target positions | |
| else: | |
| exclude = {} | |
| self.segmentation_info_mask[i] = max(x for x in self.segmentation_mask[min_y:max_y, min_x:max_x].flatten() if x not in exclude) / 100.0 | |
| if self.gt_segmentation_mask is not None: | |
| self.gt_segmentation_info_mask[i] = max(x for x in self.gt_segmentation_mask[min_y:max_y, min_x:max_x].flatten() if x not in exclude) / 100.0 | |
| # print("np.unique(self.segmentation_info_mask): ", np.unique(self.segmentation_info_mask)) | |
| self.filtered_seg_info_mask = copy.deepcopy(self.segmentation_info_mask) | |
| # In case targets found at beginning... | |
| done, num_targets_found = self.check_done() | |
| self.num_targets_found = num_targets_found | |
| def multi_robot_step(self, next_position_list, dist_list, travel_dist_list): | |
| temp_frontiers = copy.deepcopy(self.frontiers) | |
| reward_list = [] | |
| for dist, robot_position in zip(dist_list, next_position_list): | |
| self.graph_generator.route_node.append(robot_position) | |
| next_node_index = self.find_index_from_coords(robot_position) | |
| self.graph_generator.nodes_list[next_node_index].set_visited() | |
| # self.robot_belief = self.update_robot_belief(robot_position, self.sensor_range, self.robot_belief, | |
| # self.ground_truth) | |
| self.coverage_belief = self.update_robot_belief(robot_position, self.sensor_range, self.coverage_belief, | |
| self.ground_truth) | |
| self.robot_belief = self.ground_truth | |
| self.downsampled_belief = block_reduce(self.robot_belief.copy(), | |
| block_size=(self.resolution, self.resolution), | |
| func=np.min) | |
| frontiers = self.find_frontier() | |
| #num_observed_frontiers = self.calculate_num_observed_frontiers(temp_frontiers, frontiers) | |
| #temp_frontiers = frontiers | |
| num_observed_frontiers = self.node_utility[next_node_index] | |
| # individual_reward = num_observed_frontiers / 50 - dist / 64 | |
| individual_reward = -dist / 32 # 64 | |
| info_gain_reward = 0 | |
| robot_position_idx = self.find_index_from_coords(robot_position) | |
| # if self.segmentation_info_mask[robot_position_idx] == 1.0 and self.guidepost[robot_position_idx] == 0.0: | |
| # # print("High Info (Unvisited)") | |
| # info_gain_reward = (HIGH_INFO_REWARD_RATIO * 1.5) | |
| # elif self.segmentation_info_mask[robot_position_idx] == 0.0 and self.guidepost[robot_position_idx] == 0.0: | |
| # # print("Low Info (Unvisited)") | |
| # info_gain_reward = ((1-HIGH_INFO_REWARD_RATIO) * 1.5) | |
| info_gain_reward = self.filtered_seg_info_mask[robot_position_idx][0] * 1.5 | |
| if self.guidepost[robot_position_idx] == 0.0: | |
| info_gain_reward += 0.2 | |
| # print("info_gain_reward: ", info_gain_reward) | |
| individual_reward += info_gain_reward | |
| # print("dist / 64: ", dist / 64) | |
| # print("info gain reward: ", info_gain_reward) | |
| reward_list.append(individual_reward) | |
| self.node_coords, self.graph, self.node_utility, self.guidepost = self.graph_generator.update_graph(self.robot_belief, self.old_robot_belief, frontiers, self.frontiers) | |
| self.old_robot_belief = copy.deepcopy(self.robot_belief) | |
| self.filtered_seg_info_mask = [info[0] if self.guidepost[i] == 0.0 else 0.0 for i, info in enumerate(self.segmentation_info_mask)] | |
| self.filtered_seg_info_mask = np.expand_dims(np.array(self.filtered_seg_info_mask), axis=1) | |
| num_observed_frontiers = self.calculate_num_observed_frontiers(self.frontiers, frontiers) | |
| self.frontiers = frontiers | |
| self.explored_rate = self.evaluate_exploration_rate() | |
| done, num_targets_found = self.check_done() | |
| self.num_new_targets_found = num_targets_found - self.num_targets_found | |
| # #team_reward = sum(reward_list) / len(reward_list) | |
| # # team_reward = num_observed_frontiers / 50 | |
| # team_reward = self.num_new_targets_found * 5.0 | |
| team_reward = 0.0 | |
| # # print("target found reward: ", self.num_new_targets_found * 5.0) | |
| self.num_targets_found = num_targets_found | |
| self.targets_found_rate = self.evaluate_targets_found_rate() | |
| self.info_gain, self.total_info = self.evaluate_info_gain() | |
| if done: | |
| # team_reward += np.sum(self.robot_belief == 255) / sum(travel_dist_list) | |
| team_reward += 40 # 20 | |
| for i in range(len(reward_list)): | |
| reward_list[i] += team_reward | |
| return reward_list, done | |
| def import_ground_truth(self, map_index): | |
| # occupied 1, free 255, unexplored 127 | |
| try: | |
| # ground_truth = (io.imread(map_index, 1) * 255).astype(int) | |
| ground_truth = (io.imread(map_index, 1)).astype(int) | |
| if np.all(ground_truth == 0): | |
| ground_truth = (io.imread(map_index, 1) * 255).astype(int) | |
| except: | |
| new_map_index = self.map_dir + '/' + self.map_list[0] | |
| ground_truth = (io.imread(new_map_index, 1)).astype(int) | |
| print('could not read the map_path ({}), hence skipping it and using ({}).'.format(map_index, new_map_index)) | |
| robot_location = np.nonzero(ground_truth == 208) | |
| # print("robot_location: ", robot_location) | |
| # print("np.array(robot_location)[1, 127]: ", np.array(robot_location)[1, 127]) | |
| robot_location = np.array([np.array(robot_location)[1, 127], np.array(robot_location)[0, 127]]) | |
| ground_truth = (ground_truth > 150) | |
| ground_truth = ground_truth * 254 + 1 | |
| return ground_truth, robot_location | |
| def import_segmentation_mask(self, map_index): | |
| # occupied 1, free 255, unexplored 127 | |
| # mask = (io.imread(map_index, 1) * 255).astype(int) # NOTE: Cannot work well with seg mask self-generated | |
| mask = cv2.imread(map_index).astype(int) | |
| # print("np.unique(segmentation_mask): ", np.unique(mask)) | |
| # NOTE: Could contain mutiple start positions | |
| # target_position = np.nonzero(mask == 208) | |
| # target_positions = self.find_target_locations(mask) | |
| # target_position = np.array([np.array(target_position)[1, 127], np.array(target_position)[0, 127]]) | |
| return mask #, target_positions | |
| def find_target_locations(self, image_array, grey_value=208): | |
| # Load the image | |
| # image = Image.open(image_path) | |
| # image_array = np.array(image) | |
| # Identify pixels equal to the grey value | |
| grey_pixels = np.where(image_array == grey_value) | |
| # Create a binary array where grey pixels are marked as True | |
| binary_array = np.zeros_like(image_array, dtype=bool) | |
| binary_array[grey_pixels] = True | |
| # Label connected components | |
| labeled_array, num_features = label(binary_array) | |
| # Find objects returns slices for each connected component | |
| slices = find_objects(labeled_array) | |
| # Calculate the center of each box | |
| centers = [] | |
| for slice in slices: | |
| row_center = (slice[0].start + slice[0].stop - 1) // 2 | |
| col_center = (slice[1].start + slice[1].stop - 1) // 2 | |
| centers.append((col_center, row_center)) # (y,x) | |
| return centers | |
| def free_cells(self): | |
| index = np.where(self.ground_truth == 255) | |
| free = np.asarray([index[1], index[0]]).T | |
| return free | |
| def update_robot_belief(self, robot_position, sensor_range, robot_belief, ground_truth): | |
| robot_belief = sensor_work(robot_position, sensor_range, robot_belief, ground_truth) | |
| return robot_belief | |
| def check_done(self, robot_id=0): | |
| """ All agnets to have explored most of the env map """ | |
| done = False | |
| # for idx in range(self.n_agent): | |
| # if np.sum(self.ground_truth == 255) - np.sum(self.all_robot_belief[idx][idx] == 255) > 40: | |
| # done = False | |
| # NEW: ADDITIONAL VLM SEARCH CRITERIA | |
| num_targets_found = 0 | |
| self.target_found_idxs = [] | |
| for i, target in enumerate(self.target_positions): | |
| if self.coverage_belief[target[1], target[0]] == 255: # 255: | |
| num_targets_found += 1 | |
| self.target_found_idxs.append(i) | |
| # free_cells_mask = self.all_robot_belief[robot_id][robot_id] == 255 | |
| # filtered_segmentation_mask = np.where(free_cells_mask, self.segmentation_mask, 0) | |
| # targets = self.find_target_locations(filtered_segmentation_mask) | |
| # print("num_targets_found: ", num_targets_found) | |
| if TERMINATE_ON_TGTS_FOUND and num_targets_found >= len(self.target_positions): | |
| done = True | |
| if not TERMINATE_ON_TGTS_FOUND and np.sum(self.coverage_belief == 255) / np.sum(self.ground_truth == 255) >= 0.99: | |
| done = True | |
| return done, num_targets_found | |
| def calculate_num_observed_frontiers(self, old_frontiers, frontiers): | |
| frontiers_to_check = frontiers[:, 0] + frontiers[:, 1] * 1j | |
| pre_frontiers_to_check = old_frontiers[:, 0] + old_frontiers[:, 1] * 1j | |
| frontiers_num = np.intersect1d(frontiers_to_check, pre_frontiers_to_check).shape[0] | |
| pre_frontiers_num = pre_frontiers_to_check.shape[0] | |
| delta_num = pre_frontiers_num - frontiers_num | |
| return delta_num | |
| def calculate_reward(self, dist, frontiers): | |
| reward = 0 | |
| reward -= dist / 64 | |
| frontiers_to_check = frontiers[:, 0] + frontiers[:, 1] * 1j | |
| pre_frontiers_to_check = self.frontiers[:, 0] + self.frontiers[:, 1] * 1j | |
| frontiers_num = np.intersect1d(frontiers_to_check, pre_frontiers_to_check).shape[0] | |
| pre_frontiers_num = pre_frontiers_to_check.shape[0] | |
| delta_num = pre_frontiers_num - frontiers_num | |
| reward += delta_num / 50 | |
| return reward | |
| def evaluate_exploration_rate(self): | |
| # rate = np.sum(self.robot_belief == 255) / np.sum(self.ground_truth == 255) | |
| rate = np.sum(self.coverage_belief == 255) / np.sum(self.ground_truth == 255) | |
| return rate | |
| def evaluate_targets_found_rate(self): | |
| if len(self.target_positions) == 0: | |
| return 0 | |
| else: | |
| rate = self.num_targets_found / len(self.target_positions) | |
| return rate | |
| def evaluate_info_gain(self): | |
| # print("self.segmentation_mask.shape: ", self.segmentation_mask.shape) | |
| # coverage_belief = (self.coverage_belief == 255) | |
| # print("coverage_belief.shape: ", coverage_belief.shape) | |
| # print("np.unique(coverage_belief): ", np.unique(coverage_belief)) | |
| # print("np.count_nonzero(coverage_belief): ", np.count_nonzero(coverage_belief)) | |
| # print("np.count_zero(coverage_belief): ", coverage_belief.size - np.count_nonzero(coverage_belief)) | |
| # print("self.segmentation_mask[self.coverage_belief == 255].shape: ", self.segmentation_mask[self.coverage_belief == 255].shape) | |
| if self.test and TARGETS_SET_DIR != "": | |
| info_gained = np.sum(self.gt_segmentation_mask[self.coverage_belief == 255]) / 100.0 | |
| total_info = np.sum(self.gt_segmentation_mask) / 100.0 | |
| else: | |
| info_gained = np.sum(self.segmentation_mask[self.coverage_belief == 255]) / 100.0 | |
| total_info = np.sum(self.segmentation_mask) / 100.0 | |
| return info_gained, total_info | |
| def calculate_new_free_area(self): | |
| old_free_area = self.old_robot_belief == 255 | |
| current_free_area = self.robot_belief == 255 | |
| new_free_area = (current_free_area.astype(np.int) - old_free_area.astype(np.int)) * 255 | |
| return new_free_area, np.sum(old_free_area) | |
| def calculate_dist_path(self, path): | |
| dist = 0 | |
| start = path[0] | |
| end = path[-1] | |
| for index in path: | |
| if index == end: | |
| break | |
| dist += np.linalg.norm(self.node_coords[start] - self.node_coords[index]) | |
| start = index | |
| return dist | |
| def find_frontier(self): | |
| y_len = self.downsampled_belief.shape[0] | |
| x_len = self.downsampled_belief.shape[1] | |
| mapping = self.downsampled_belief.copy() | |
| belief = self.downsampled_belief.copy() | |
| # 0-1 unknown area map | |
| mapping = (mapping == 127) * 1 | |
| mapping = np.lib.pad(mapping, ((1, 1), (1, 1)), 'constant', constant_values=0) | |
| fro_map = mapping[2:][:, 1:x_len + 1] + mapping[:y_len][:, 1:x_len + 1] + mapping[1:y_len + 1][:, 2:] + \ | |
| mapping[1:y_len + 1][:, :x_len] + mapping[:y_len][:, 2:] + mapping[2:][:, :x_len] + mapping[2:][:, | |
| 2:] + \ | |
| mapping[:y_len][:, :x_len] | |
| ind_free = np.where(belief.ravel(order='F') == 255)[0] | |
| ind_fron_1 = np.where(1 < fro_map.ravel(order='F'))[0] | |
| ind_fron_2 = np.where(fro_map.ravel(order='F') < 8)[0] | |
| ind_fron = np.intersect1d(ind_fron_1, ind_fron_2) | |
| ind_to = np.intersect1d(ind_free, ind_fron) | |
| map_x = x_len | |
| map_y = y_len | |
| x = np.linspace(0, map_x - 1, map_x) | |
| y = np.linspace(0, map_y - 1, map_y) | |
| t1, t2 = np.meshgrid(x, y) | |
| points = np.vstack([t1.T.ravel(), t2.T.ravel()]).T | |
| f = points[ind_to] | |
| f = f.astype(int) | |
| f = f * self.resolution | |
| return f | |
| def plot_env(self, n, path, step, travel_dist, robots_route, img_path_override=None, sat_path_override=None, msk_name_override=None, sound_id_override=None, colormap_mid_val=None): | |
| # # TEMP | |
| # if TAXABIND_TTA: | |
| # # Save self.segmentation_info_mask as .npy file in gifs_path | |
| # side_dim = int(np.sqrt(self.segmentation_info_mask.shape[0])) | |
| # mask_viz = self.segmentation_info_mask.squeeze().reshape((side_dim, side_dim)).T | |
| # np.save(os.path.join(path, f"seg_mask_step{step}.npy"), mask_viz) | |
| plt.switch_backend('agg') | |
| # plt.ion() | |
| plt.cla() | |
| color_list = ["r", "g", "c", "m", "y", "k"] | |
| if TARGETS_SET_DIR == "" and not TAXABIND_TTA: | |
| fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5)) | |
| else: | |
| fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, figsize=(12, 8)) | |
| ### Fig1: Environment ### | |
| msk_name = "" | |
| if TAXABIND_TTA: | |
| image = mpimg.imread(sat_path_override) | |
| msk_name = msk_name_override | |
| # else: | |
| # plt.imshow(self.robot_belief, cmap='gray') | |
| # ax1.imshow(self.coverage_belief, cmap='gray') | |
| # image = mpimg.imread("Maps/real_maps/real/4259_masked_img_0.jpg") | |
| # msk_name = self.map_list[self.map_index] | |
| # raw_img_path = self.score_to_img_dict[msk_name] | |
| # if "flair" in raw_img_path: | |
| # with rasterio.open(raw_img_path) as src_img: | |
| # image = src_img.read([1,2,3]) | |
| # image = np.transpose(image, (1, 2, 0)) | |
| # else: | |
| # image = mpimg.imread(raw_img_path) | |
| ### Fig1: Environment ### | |
| ax = ax1 # if TAXABIND_TTA else ax1 | |
| ax.imshow(image) | |
| ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0)) | |
| ax.set_title("Image") | |
| # if VIZ_GRAPH_EDGES: | |
| # for i in range(len(self.graph_generator.x)): | |
| # ax.plot(self.graph_generator.x[i], self.graph_generator.y[i], 'tan', zorder=1) | |
| # # ax.scatter(self.node_coords[:, 0], self.node_coords[:, 1], c=self.node_utility, zorder=5) | |
| # ax.scatter(self.node_coords[:, 0], self.node_coords[:, 1], c=self.segmentation_info_mask, zorder=5) | |
| # ax.scatter(self.frontiers[:, 0], self.frontiers[:, 1], c='r', s=2, zorder=3) | |
| for i, route in enumerate(robots_route): | |
| robot_marker_color = color_list[i % len(color_list)] | |
| xPoints = route[0] | |
| yPoints = route[1] | |
| ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2) | |
| # ax.plot(xPoints[-1], yPoints[-1], 'mo', markersize=8, zorder=10) | |
| ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black") | |
| ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5) | |
| # Sensor range | |
| rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH) | |
| rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT) | |
| max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1]) | |
| min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0) | |
| max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0]) | |
| min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0) | |
| ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1) | |
| ### Fig2: Graph ### | |
| ax = ax4 if TAXABIND_TTA else ax1 | |
| # ax.imshow(image) | |
| ax.imshow(self.coverage_belief, cmap='gray') | |
| ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0)) | |
| ax.set_title("Information Graph") | |
| if VIZ_GRAPH_EDGES: | |
| for i in range(len(self.graph_generator.x)): | |
| ax.plot(self.graph_generator.x[i], self.graph_generator.y[i], 'tan', zorder=1) | |
| # ax.scatter(self.node_coords[:, 0], self.node_coords[:, 1], c=self.node_utility, zorder=5) | |
| # ax.scatter(self.node_coords[:, 0], self.node_coords[:, 1], c=self.segmentation_info_mask, zorder=5) | |
| # filtered_seg_info_mask = [info[0] if self.guidepost[i] == 0.0 else 0.0 for i, info in enumerate(self.segmentation_info_mask)] | |
| ax.scatter(self.node_coords[:, 0], self.node_coords[:, 1], c=self.filtered_seg_info_mask, zorder=5, s=8) | |
| # ax.scatter(self.frontiers[:, 0], self.frontiers[:, 1], c='r', s=2, zorder=3) | |
| for i, route in enumerate(robots_route): | |
| robot_marker_color = color_list[i % len(color_list)] | |
| xPoints = route[0] | |
| yPoints = route[1] | |
| ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2) | |
| # ax.plot(xPoints[-1], yPoints[-1], 'mo', markersize=8, zorder=10) | |
| ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black") | |
| ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5) | |
| # Sensor range | |
| rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH) | |
| rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT) | |
| max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1]) | |
| min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0) | |
| max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0]) | |
| min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0) | |
| ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1) | |
| # Plot target positions | |
| for target in self.target_positions: | |
| if self.coverage_belief[target[1], target[0]] == 255: | |
| # ax.plot(target[0], target[1], 'go', markersize=8, zorder=99) | |
| ax.plot(target[0], target[1], color='g', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99) | |
| else: | |
| # ax.plot(target[0], target[1], 'ro', markersize=8, zorder=99) | |
| ax.plot(target[0], target[1], color='r', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99) | |
| # ax.pause(0.1) | |
| ### Fig3: Segmentation Mask ### | |
| ax = ax5 if TAXABIND_TTA else ax2 | |
| if TAXABIND_TTA and USE_CLIP_PREDS: | |
| side_dim = int(np.sqrt(self.segmentation_info_mask.shape[0])) | |
| mask_viz = self.segmentation_info_mask.squeeze().reshape((side_dim, side_dim)).T | |
| scale_y = math.ceil(self.ground_truth_size[1] / side_dim) | |
| scale_x = math.ceil(self.ground_truth_size[0] / side_dim) | |
| upscaled_mask_viz = np.kron(mask_viz, np.ones((scale_y, scale_x))) # Integer scaling only | |
| upscaled_mask_viz = upscaled_mask_viz[:self.ground_truth_size[1], :self.ground_truth_size[0]] | |
| im = ax.imshow(upscaled_mask_viz, cmap="viridis") | |
| ax.axis("off") | |
| else: | |
| im = ax.imshow(self.segmentation_mask.mean(axis=-1), cmap='viridis', vmin=0, vmax=100) # cmap='gray' | |
| ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0)) | |
| ax.set_title(f"Predicted Mask (Normalized)") | |
| for i, route in enumerate(robots_route): | |
| robot_marker_color = color_list[i % len(color_list)] | |
| xPoints = route[0] | |
| yPoints = route[1] | |
| ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2) | |
| # ax.plot(xPoints[-1], yPoints[-1], 'mo', markersize=8, zorder=10) | |
| ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black") | |
| ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5) | |
| # Sensor range | |
| rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH) | |
| rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT) | |
| max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1]) | |
| min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0) | |
| max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0]) | |
| min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0) | |
| ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1) | |
| # Add a colorbar | |
| cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) | |
| cbar.set_label("Normalized Probs") | |
| # ax.pause(0.1) | |
| ### Fig4: Segmentation Mask ### | |
| if TAXABIND_TTA and USE_CLIP_PREDS: | |
| ax = ax6 | |
| side_dim = int(np.sqrt(self.segmentation_info_mask_unnormalized.shape[0])) | |
| mask_viz = self.segmentation_info_mask_unnormalized.squeeze().reshape((side_dim, side_dim)).T | |
| scale_y = math.ceil(self.ground_truth_size[1] / side_dim) | |
| scale_x = math.ceil(self.ground_truth_size[0] / side_dim) | |
| upscaled_mask_viz = np.kron(mask_viz, np.ones((scale_y, scale_x))) # Integer scaling only | |
| upscaled_mask_viz = upscaled_mask_viz[:self.ground_truth_size[1], :self.ground_truth_size[0]] | |
| max_val = 0.15 # TO CHANGE | |
| mid_val = colormap_mid_val if colormap_mid_val is not None else 0.05 | |
| # mid_val = np.max(self.segmentation_info_mask_unnormalized) | |
| norm = CustomNorm(vmin=0.0, vmax=max_val, mid=mid_val, lower_portion=0.8) | |
| im = ax.imshow(upscaled_mask_viz, cmap="viridis", norm=norm) # norm=LogNorm(vmin=0.01, vmax=0.1)) | |
| # norm = PowerNorm(gamma=0.25, vmin=0.01, vmax=0.2) | |
| # norm=LogNorm(vmin=0.01, vmax=0.2) | |
| im = ax.imshow(upscaled_mask_viz, cmap="viridis", norm=norm) # norm=LogNorm(vmin=0.01, vmax=0.1)) | |
| ax.axis("off") | |
| # else: | |
| # im = ax.imshow(self.segmentation_mask.mean(axis=-1), cmap='viridis', vmin=0, vmax=100) # cmap='gray' | |
| # ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0)) | |
| ax.set_title(f"Predicted Mask (Unnormalized)") | |
| for i, route in enumerate(robots_route): | |
| robot_marker_color = color_list[i % len(color_list)] | |
| xPoints = route[0] | |
| yPoints = route[1] | |
| ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2) | |
| # ax.plot(xPoints[-1], yPoints[-1], 'mo', markersize=8, zorder=10) | |
| ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black") | |
| ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5) | |
| # Sensor range | |
| rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH) | |
| rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT) | |
| max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1]) | |
| min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0) | |
| max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0]) | |
| min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0) | |
| ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1) | |
| # Add a colorbar | |
| cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) | |
| if TAXABIND_TTA and USE_CLIP_PREDS: | |
| cbar.set_ticks([0.0, mid_val, max_val]) | |
| cbar.set_label("Probs (Scaled by expectation)") | |
| # Fog5: GT Mask | |
| if TARGETS_SET_DIR != "": | |
| ax = ax2 | |
| im = ax.imshow(self.gt_segmentation_mask.mean(axis=-1), cmap='viridis', vmin=0, vmax=100) # cmap='gray' | |
| ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0)) | |
| ax.set_title(f"Ground Truth Mask") | |
| for i, route in enumerate(robots_route): | |
| robot_marker_color = color_list[i % len(color_list)] | |
| xPoints = route[0] | |
| yPoints = route[1] | |
| ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2) | |
| # ax.plot(xPoints[-1], yPoints[-1], 'mo', markersize=8, zorder=10) | |
| ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black") | |
| ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5) | |
| # Sensor range | |
| rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH) | |
| rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT) | |
| max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1]) | |
| min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0) | |
| max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0]) | |
| min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0) | |
| ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1) | |
| # Add a colorbar | |
| cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) | |
| cbar.set_label("Normalized Mask Value") | |
| # ax4.pause(0.1) | |
| ### Fig6: Segmentation Mask (GT) ### | |
| if TAXABIND_TTA: | |
| ax = ax3 | |
| image = mpimg.imread(img_path_override) | |
| ax.imshow(image) | |
| ax.set_title("Ground Image") | |
| ax.axis("off") | |
| sound_id = sound_id_override if sound_id_override is not None else "-1" | |
| plt.suptitle('Targets Found: {}/{} Coverage ratio: {:.4g} Travel Dist: {:.4g} Info Gain: {:.4g}% \n ({}) \n (Sound ID: {})'.format(self.num_targets_found, \ | |
| len(self.target_positions), self.explored_rate, travel_dist, (100*self.info_gain/self.total_info), msk_name, | |
| sound_id)) | |
| plt.tight_layout() | |
| plt.savefig('{}/{}_{}_samples.png'.format(path, n, step, dpi=100)) | |
| # plt.show() | |
| frame = '{}/{}_{}_samples.png'.format(path, n, step) | |
| self.frame_files.append(frame) | |
| plt.close() | |
| #################### | |
| # ADDED: For app.py | |
| #################### | |
| def plot_heatmap(self, save_dir, step, travel_dist, robots_route=None): | |
| """Plot only the segmentation heatmap and save it as ``{step}.png`` in | |
| ``save_dir``. This lightweight helper is meant for asynchronous | |
| streaming in the Gradio demo when full `plot_env` is too heavy. | |
| Parameters | |
| ---------- | |
| save_dir : str | |
| Directory to save the generated PNG file. | |
| step : int | |
| Current timestep; becomes the filename ``{step}.png``. | |
| robots_route : list | None | |
| Optional list of routes (xPoints, yPoints) to overlay. | |
| Returns | |
| ------- | |
| str | |
| Full path to the generated PNG file. | |
| """ | |
| import os | |
| plt.switch_backend('agg') | |
| # Do not clear the global figure state in case it interferes with | |
| # the current figure. Each call creates its own Figure object that | |
| # we close explicitly at the end, so a global clear is unnecessary | |
| # and may break concurrent drawing. | |
| # plt.cla() | |
| color_list = ["r", "g", "c", "m", "y", "k"] | |
| fig, ax = plt.subplots(1, 1, figsize=(6, 6)) | |
| # Select the mask to visualise | |
| # if TAXABIND_TTA and USE_CLIP_PREDS: | |
| side_dim = int(np.sqrt(self.segmentation_info_mask.shape[0])) | |
| mask_viz = self.segmentation_info_mask.squeeze().reshape((side_dim, side_dim)).T | |
| # Properly map image to pixel coordinates and keep limits fixed | |
| H, W = self.ground_truth_size # rows (y), cols (x) | |
| im = ax.imshow( | |
| mask_viz, | |
| cmap="viridis", | |
| origin="upper", | |
| extent=[0, W, H, 0], # x: 0..W, y: H..0 (origin at top-left) | |
| interpolation="nearest", # keep cell edges sharp & aligned | |
| zorder=0, | |
| ) | |
| ax.set_xlim(0, W) | |
| ax.set_ylim(H, 0) | |
| ax.set_axis_off() # hide ticks but keep limits | |
| # else: | |
| # im = ax.imshow(self.segmentation_mask.mean(axis=-1), cmap='viridis', vmin=0, vmax=100) | |
| # ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0)) | |
| # Optionally overlay robot paths | |
| if robots_route is not None: | |
| for i, route in enumerate(robots_route): | |
| robot_marker_color = color_list[i % len(color_list)] | |
| xPoints, yPoints = route | |
| ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2) | |
| ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black") | |
| ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5) | |
| # Plot target positions | |
| for target in self.target_positions: | |
| if self.coverage_belief[target[1], target[0]] == 255: | |
| # ax.plot(target[0], target[1], 'go', markersize=8, zorder=99) | |
| ax.plot(target[0], target[1], color='g', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99) | |
| else: | |
| # ax.plot(target[0], target[1], 'ro', markersize=8, zorder=99) | |
| ax.plot(target[0], target[1], color='r', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99) | |
| # Sensor range | |
| rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH) | |
| rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT) | |
| max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1]) | |
| min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0) | |
| max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0]) | |
| min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0) | |
| ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1) | |
| ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1) | |
| # Color bar | |
| cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) | |
| cbar.set_label("Normalized Probs") | |
| # Change coverage to 1dp | |
| plt.suptitle('Targets Found: {}/{} Coverage: {:.1f}% Steps: {}/{}'.format( | |
| self.num_targets_found, \ | |
| len(self.target_positions), | |
| self.explored_rate*100, | |
| step+1, | |
| NUM_EPS_STEPS), | |
| y=0.94, # Closer to plot | |
| ) | |
| plt.tight_layout() | |
| os.makedirs(save_dir, exist_ok=True) | |
| out_path = os.path.join(save_dir, f"{step}.png") | |
| # Save atomically: write to temp file then move into place so the poller never sees a partial file. | |
| tmp_path = out_path + ".tmp" | |
| fig.savefig(tmp_path, dpi=100, format='png') | |
| os.replace(tmp_path, out_path) # atomic on same filesystem | |
| plt.close(fig) | |
| return out_path | |
| #################### | |
| class CustomNorm(Normalize): | |
| """ | |
| A custom normalization that allocates a larger fraction of the colormap | |
| to the lower data range [vmin, mid] than to [mid, vmax]. | |
| Parameters | |
| ---------- | |
| vmin : float | |
| Minimum data value | |
| vmax : float | |
| Maximum data value | |
| mid : float | |
| Midpoint in data where we switch from 'lower' to 'upper' mapping | |
| lower_portion : float | |
| Fraction of the colormap to allocate for [vmin, mid]. | |
| For example, 0.8 => 80% of colors for [vmin, mid], 20% for [mid, vmax]. | |
| clip : bool | |
| Whether to clip data outside [vmin, vmax]. | |
| """ | |
| def __init__(self, vmin=None, vmax=None, mid=0.05, lower_portion=0.8, clip=False): | |
| self.mid = mid | |
| self.lower_portion = lower_portion | |
| super().__init__(vmin, vmax, clip) | |
| def __call__(self, value, clip=None): | |
| """Forward transform: data -> [0..1] color space.""" | |
| vmin, vmax, mid = self.vmin, self.vmax, self.mid | |
| lp = self.lower_portion | |
| value = np.asarray(value, dtype=np.float64) | |
| # Piecewise linear mapping: | |
| # [vmin..mid] => [0..lp] | |
| # [mid..vmax] => [lp..1] | |
| normed = np.where( | |
| value <= mid, | |
| lp * (value - vmin) / (mid - vmin), | |
| lp + (value - mid) / (vmax - mid) * (1 - lp) | |
| ) | |
| return np.clip(normed, 0, 1) | |
| def inverse(self, value): | |
| """ | |
| Inverse transform: [0..1] color space -> data space. | |
| Matplotlib's colorbar calls this to place ticks correctly. | |
| """ | |
| vmin, vmax, mid = self.vmin, self.vmax, self.mid | |
| lp = self.lower_portion | |
| value = np.asarray(value, dtype=np.float64) | |
| # For color space [0..lp], invert to [vmin..mid] | |
| # For color space [lp..1], invert to [mid..vmax] | |
| below = (value <= lp) | |
| above = ~below | |
| # Allocate array for results | |
| data = np.zeros_like(value, dtype=np.float64) | |
| # Invert lower segment | |
| data[below] = vmin + (value[below] / lp) * (mid - vmin) | |
| # Invert upper segment | |
| data[above] = mid + ((value[above] - lp) / (1 - lp)) * (vmax - mid) | |
| return data | |