Spaces:
Running
Running
File size: 5,602 Bytes
f90241e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
#!/usr/bin/env python3
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# main executable for preprocessing habitat
# export METADATA_DIR="/path/to/habitat/5views_v1_512x512_metadata"
# export SCENES_DIR="/path/to/habitat/data/scene_datasets/"
# export OUTPUT_DIR="data/habitat_processed"
# export PYTHONPATH=$(pwd)
# python preprocess_habitat.py --scenes_dir=$SCENES_DIR --metadata_dir=$METADATA_DIR --output_dir=$OUTPUT_DIR | parallel -j 16
# --------------------------------------------------------
import os
import glob
import json
import os
import PIL.Image
import json
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" # noqa
import cv2
from habitat_renderer import multiview_crop_generator
from tqdm import tqdm
def preprocess_metadata(metadata_filename,
scenes_dir,
output_dir,
crop_resolution=[512, 512],
equirectangular_resolution=None,
fix_existing_dataset=False):
# Load data
with open(metadata_filename, "r") as f:
metadata = json.load(f)
if metadata["scene_dataset_config_file"] == "":
scene = os.path.join(scenes_dir, metadata["scene"])
scene_dataset_config_file = ""
else:
scene = metadata["scene"]
scene_dataset_config_file = os.path.join(scenes_dir, metadata["scene_dataset_config_file"])
navmesh = None
# Use 4 times the crop size as resolution for rendering the environment map.
max_res = max(crop_resolution)
if equirectangular_resolution == None:
# Use 4 times the crop size as resolution for rendering the environment map.
max_res = max(crop_resolution)
equirectangular_resolution = (4*max_res, 8*max_res)
print("equirectangular_resolution:", equirectangular_resolution)
if os.path.exists(output_dir) and not fix_existing_dataset:
raise FileExistsError(output_dir)
# Lazy initialization
highres_dataset = None
for batch_label, batch in tqdm(metadata["view_batches"].items()):
for view_label, view_params in batch.items():
assert view_params["size"] == crop_resolution
label = f"{batch_label}_{view_label}"
output_camera_params_filename = os.path.join(output_dir, f"{label}_camera_params.json")
if fix_existing_dataset and os.path.isfile(output_camera_params_filename):
# Skip generation if we are fixing a dataset and the corresponding output file already exists
continue
# Lazy initialization
if highres_dataset is None:
highres_dataset = multiview_crop_generator.HabitatMultiviewCrops(scene=scene,
navmesh=navmesh,
scene_dataset_config_file=scene_dataset_config_file,
equirectangular_resolution=equirectangular_resolution,
crop_resolution=crop_resolution,)
os.makedirs(output_dir, exist_ok=bool(fix_existing_dataset))
# Generate a higher resolution crop
original_projection, position = multiview_crop_generator.dict_to_perspective_projection(view_params)
# Render an envmap at the given position
viewpoint_data = highres_dataset.render_viewpoint_data(position)
projection = original_projection
colormap, depthmap, pointmap, _ = highres_dataset.extract_cropped_camera(
projection, viewpoint_data.colormap, viewpoint_data.distancemap, viewpoint_data.pointmap)
camera_params = multiview_crop_generator.perspective_projection_to_dict(projection, position)
# Color image
PIL.Image.fromarray(colormap).save(os.path.join(output_dir, f"{label}.jpeg"))
# Depth image
cv2.imwrite(os.path.join(output_dir, f"{label}_depth.exr"),
depthmap, [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF])
# Camera parameters
with open(output_camera_params_filename, "w") as f:
json.dump(camera_params, f)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--metadata_dir", required=True)
parser.add_argument("--scenes_dir", required=True)
parser.add_argument("--output_dir", required=True)
parser.add_argument("--metadata_filename", default="")
args = parser.parse_args()
if args.metadata_filename == "":
# Walk through the metadata dir to generate commandlines
for filename in glob.iglob(os.path.join(args.metadata_dir, "**/metadata.json"), recursive=True):
output_dir = os.path.join(args.output_dir, os.path.relpath(os.path.dirname(filename), args.metadata_dir))
if not os.path.exists(output_dir):
commandline = f"python {__file__} --metadata_filename={filename} --metadata_dir={args.metadata_dir} --scenes_dir={args.scenes_dir} --output_dir={output_dir}"
print(commandline)
else:
preprocess_metadata(metadata_filename=args.metadata_filename,
scenes_dir=args.scenes_dir,
output_dir=args.output_dir)
|