Spaces:
Sleeping
Sleeping
| # Copyright (C) 2022-present Naver Corporation. All rights reserved. | |
| # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). | |
| """ | |
| Script to generate image pairs for a given scene reproducing poses provided in a metadata file. | |
| """ | |
| import os | |
| from datasets.habitat_sim.multiview_habitat_sim_generator import MultiviewHabitatSimGenerator | |
| from datasets.habitat_sim.paths import SCENES_DATASET | |
| import argparse | |
| import quaternion | |
| import PIL.Image | |
| import cv2 | |
| import json | |
| from tqdm import tqdm | |
| def generate_multiview_images_from_metadata(metadata_filename, | |
| output_dir, | |
| overload_params = dict(), | |
| scene_datasets_paths=None, | |
| exist_ok=False): | |
| """ | |
| Generate images from a metadata file for reproducibility purposes. | |
| """ | |
| # Reorder paths by decreasing label length, to avoid collisions when testing if a string by such label | |
| if scene_datasets_paths is not None: | |
| scene_datasets_paths = dict(sorted(scene_datasets_paths.items(), key= lambda x: len(x[0]), reverse=True)) | |
| with open(metadata_filename, 'r') as f: | |
| input_metadata = json.load(f) | |
| metadata = dict() | |
| for key, value in input_metadata.items(): | |
| # Optionally replace some paths | |
| if key in ("scene_dataset_config_file", "scene", "navmesh") and value != "": | |
| if scene_datasets_paths is not None: | |
| for dataset_label, dataset_path in scene_datasets_paths.items(): | |
| if value.startswith(dataset_label): | |
| value = os.path.normpath(os.path.join(dataset_path, os.path.relpath(value, dataset_label))) | |
| break | |
| metadata[key] = value | |
| # Overload some parameters | |
| for key, value in overload_params.items(): | |
| metadata[key] = value | |
| generation_entries = dict([(key, value) for key, value in metadata.items() if not (key in ('multiviews', 'output_dir', 'generate_depth'))]) | |
| generate_depth = metadata["generate_depth"] | |
| os.makedirs(output_dir, exist_ok=exist_ok) | |
| generator = MultiviewHabitatSimGenerator(**generation_entries) | |
| # Generate views | |
| for idx_label, data in tqdm(metadata['multiviews'].items()): | |
| positions = data["positions"] | |
| orientations = data["orientations"] | |
| n = len(positions) | |
| for oidx in range(n): | |
| observation = generator.render_viewpoint(positions[oidx], quaternion.from_float_array(orientations[oidx])) | |
| observation_label = f"{oidx + 1}" # Leonid is indexing starting from 1 | |
| # Color image saved using PIL | |
| img = PIL.Image.fromarray(observation['color'][:,:,:3]) | |
| filename = os.path.join(output_dir, f"{idx_label}_{observation_label}.jpeg") | |
| img.save(filename) | |
| if generate_depth: | |
| # Depth image as EXR file | |
| filename = os.path.join(output_dir, f"{idx_label}_{observation_label}_depth.exr") | |
| cv2.imwrite(filename, observation['depth'], [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF]) | |
| # Camera parameters | |
| camera_params = dict([(key, observation[key].tolist()) for key in ("camera_intrinsics", "R_cam2world", "t_cam2world")]) | |
| filename = os.path.join(output_dir, f"{idx_label}_{observation_label}_camera_params.json") | |
| with open(filename, "w") as f: | |
| json.dump(camera_params, f) | |
| # Save metadata | |
| with open(os.path.join(output_dir, "metadata.json"), "w") as f: | |
| json.dump(metadata, f) | |
| generator.close() | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--metadata_filename", required=True) | |
| parser.add_argument("--output_dir", required=True) | |
| args = parser.parse_args() | |
| generate_multiview_images_from_metadata(metadata_filename=args.metadata_filename, | |
| output_dir=args.output_dir, | |
| scene_datasets_paths=SCENES_DATASET, | |
| overload_params=dict(), | |
| exist_ok=True) | |