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