# 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)