croco/datasets/habitat_sim/generate_from_metadata.py

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