# Copyright (C) 2022-present Naver Corporation. All rights reserved. # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). import os from tqdm import tqdm import argparse import PIL.Image import numpy as np import json from datasets.habitat_sim.multiview_habitat_sim_generator import ( MultiviewHabitatSimGenerator, NoNaviguableSpaceError, ) from datasets.habitat_sim.paths import list_scenes_available import cv2 import quaternion import shutil def generate_multiview_images_for_scene( scene_dataset_config_file, scene, navmesh, output_dir, views_count, size, exist_ok=False, generate_depth=False, **kwargs, ): """ Generate tuples of overlapping views for a given scene. generate_depth: generate depth images and camera parameters. """ if os.path.exists(output_dir) and not exist_ok: print(f"Scene {scene}: data already generated. Ignoring generation.") return try: print(f"Scene {scene}: {size} multiview acquisitions to generate...") os.makedirs(output_dir, exist_ok=exist_ok) metadata_filename = os.path.join(output_dir, "metadata.json") metadata_template = dict( scene_dataset_config_file=scene_dataset_config_file, scene=scene, navmesh=navmesh, views_count=views_count, size=size, generate_depth=generate_depth, **kwargs, ) metadata_template["multiviews"] = dict() if os.path.exists(metadata_filename): print("Metadata file already exists:", metadata_filename) print("Loading already generated metadata file...") with open(metadata_filename, "r") as f: metadata = json.load(f) for key in metadata_template.keys(): if key != "multiviews": assert ( metadata_template[key] == metadata[key] ), f"existing file is inconsistent with the input parameters:\nKey: {key}\nmetadata: {metadata[key]}\ntemplate: {metadata_template[key]}." else: print("No temporary file found. Starting generation from scratch...") metadata = metadata_template starting_id = len(metadata["multiviews"]) print(f"Starting generation from index {starting_id}/{size}...") if starting_id >= size: print("Generation already done.") return generator = MultiviewHabitatSimGenerator( scene_dataset_config_file=scene_dataset_config_file, scene=scene, navmesh=navmesh, views_count=views_count, size=size, **kwargs, ) for idx in tqdm(range(starting_id, size)): # Generate / re-generate the observations try: data = generator[idx] observations = data["observations"] positions = data["positions"] orientations = data["orientations"] idx_label = f"{idx:08}" for oidx, observation in enumerate(observations): 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) metadata["multiviews"][idx_label] = { "positions": positions.tolist(), "orientations": orientations.tolist(), "covisibility_ratios": data["covisibility_ratios"].tolist(), "valid_fractions": data["valid_fractions"].tolist(), "pairwise_visibility_ratios": data[ "pairwise_visibility_ratios" ].tolist(), } except RecursionError: print( "Recursion error: unable to sample observations for this scene. We will stop there." ) break # Regularly save a temporary metadata file, in case we need to restart the generation if idx % 10 == 0: with open(metadata_filename, "w") as f: json.dump(metadata, f) # Save metadata with open(metadata_filename, "w") as f: json.dump(metadata, f) generator.close() except NoNaviguableSpaceError: pass def create_commandline(scene_data, generate_depth, exist_ok=False): """ Create a commandline string to generate a scene. """ def my_formatting(val): if val is None or val == "": return '""' else: return val commandline = f"""python {__file__} --scene {my_formatting(scene_data.scene)} --scene_dataset_config_file {my_formatting(scene_data.scene_dataset_config_file)} --navmesh {my_formatting(scene_data.navmesh)} --output_dir {my_formatting(scene_data.output_dir)} --generate_depth {int(generate_depth)} --exist_ok {int(exist_ok)} """ commandline = " ".join(commandline.split()) return commandline if __name__ == "__main__": os.umask(2) parser = argparse.ArgumentParser( description="""Example of use -- listing commands to generate data for scenes available: > python datasets/habitat_sim/generate_multiview_habitat_images.py --list_commands """ ) parser.add_argument("--output_dir", type=str, required=True) parser.add_argument( "--list_commands", action="store_true", help="list commandlines to run if true" ) parser.add_argument("--scene", type=str, default="") parser.add_argument("--scene_dataset_config_file", type=str, default="") parser.add_argument("--navmesh", type=str, default="") parser.add_argument("--generate_depth", type=int, default=1) parser.add_argument("--exist_ok", type=int, default=0) kwargs = dict(resolution=(256, 256), hfov=60, views_count=2, size=1000) args = parser.parse_args() generate_depth = bool(args.generate_depth) exist_ok = bool(args.exist_ok) if args.list_commands: # Listing scenes available... scenes_data = list_scenes_available(base_output_dir=args.output_dir) for scene_data in scenes_data: print( create_commandline( scene_data, generate_depth=generate_depth, exist_ok=exist_ok ) ) else: if args.scene == "" or args.output_dir == "": print("Missing scene or output dir argument!") print(parser.format_help()) else: generate_multiview_images_for_scene( scene=args.scene, scene_dataset_config_file=args.scene_dataset_config_file, navmesh=args.navmesh, output_dir=args.output_dir, exist_ok=exist_ok, generate_depth=generate_depth, **kwargs, )