croco/datasets/habitat_sim/generate_multiview_images.py

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