Spaces:
Runtime error
Runtime error
# Author: Bingxin Ke | |
# Last modified: 2024-02-19 | |
import argparse | |
import os | |
import cv2 | |
import h5py | |
import numpy as np | |
import pandas as pd | |
from hypersim_util import dist_2_depth, tone_map | |
from tqdm import tqdm | |
IMG_WIDTH = 1024 | |
IMG_HEIGHT = 768 | |
FOCAL_LENGTH = 886.81 | |
if "__main__" == __name__: | |
parser = argparse.ArgumentParser() | |
parser.add_argument( | |
"--split_csv", | |
type=str, | |
default="data/Hypersim/metadata_images_split_scene_v1.csv", | |
) | |
parser.add_argument("--dataset_dir", type=str, default="data/Hypersim/raw_data") | |
parser.add_argument("--output_dir", type=str, default="data/Hypersim/processed") | |
args = parser.parse_args() | |
split_csv = args.split_csv | |
dataset_dir = args.dataset_dir | |
output_dir = args.output_dir | |
# %% | |
raw_meta_df = pd.read_csv(split_csv) | |
meta_df = raw_meta_df[raw_meta_df.included_in_public_release].copy() | |
# %% | |
for split in ["train", "val", "test"]: | |
split_output_dir = os.path.join(output_dir, split) | |
os.makedirs(split_output_dir) | |
split_meta_df = meta_df[meta_df.split_partition_name == split].copy() | |
split_meta_df["rgb_path"] = None | |
split_meta_df["rgb_mean"] = np.nan | |
split_meta_df["rgb_std"] = np.nan | |
split_meta_df["rgb_min"] = np.nan | |
split_meta_df["rgb_max"] = np.nan | |
split_meta_df["depth_path"] = None | |
split_meta_df["depth_mean"] = np.nan | |
split_meta_df["depth_std"] = np.nan | |
split_meta_df["depth_min"] = np.nan | |
split_meta_df["depth_max"] = np.nan | |
split_meta_df["invalid_ratio"] = np.nan | |
for i, row in tqdm(split_meta_df.iterrows(), total=len(split_meta_df)): | |
# Load data | |
rgb_path = os.path.join( | |
row.scene_name, | |
"images", | |
f"scene_{row.camera_name}_final_hdf5", | |
f"frame.{row.frame_id:04d}.color.hdf5", | |
) | |
dist_path = os.path.join( | |
row.scene_name, | |
"images", | |
f"scene_{row.camera_name}_geometry_hdf5", | |
f"frame.{row.frame_id:04d}.depth_meters.hdf5", | |
) | |
render_entity_id_path = os.path.join( | |
row.scene_name, | |
"images", | |
f"scene_{row.camera_name}_geometry_hdf5", | |
f"frame.{row.frame_id:04d}.render_entity_id.hdf5", | |
) | |
assert os.path.exists(os.path.join(dataset_dir, rgb_path)) | |
assert os.path.exists(os.path.join(dataset_dir, dist_path)) | |
with h5py.File(os.path.join(dataset_dir, rgb_path), "r") as f: | |
rgb = np.array(f["dataset"]).astype(float) | |
with h5py.File(os.path.join(dataset_dir, dist_path), "r") as f: | |
dist_from_center = np.array(f["dataset"]).astype(float) | |
with h5py.File(os.path.join(dataset_dir, render_entity_id_path), "r") as f: | |
render_entity_id = np.array(f["dataset"]).astype(int) | |
# Tone map | |
rgb_color_tm = tone_map(rgb, render_entity_id) | |
rgb_int = (rgb_color_tm * 255).astype(np.uint8) # [H, W, RGB] | |
# Distance -> depth | |
plane_depth = dist_2_depth( | |
IMG_WIDTH, IMG_HEIGHT, FOCAL_LENGTH, dist_from_center | |
) | |
valid_mask = render_entity_id != -1 | |
# Record invalid ratio | |
invalid_ratio = (np.prod(valid_mask.shape) - valid_mask.sum()) / np.prod( | |
valid_mask.shape | |
) | |
plane_depth[~valid_mask] = 0 | |
# Save as png | |
scene_path = row.scene_name | |
if not os.path.exists(os.path.join(split_output_dir, row.scene_name)): | |
os.makedirs(os.path.join(split_output_dir, row.scene_name)) | |
rgb_name = f"rgb_{row.camera_name}_fr{row.frame_id:04d}.png" | |
rgb_path = os.path.join(scene_path, rgb_name) | |
cv2.imwrite( | |
os.path.join(split_output_dir, rgb_path), | |
cv2.cvtColor(rgb_int, cv2.COLOR_RGB2BGR), | |
) | |
plane_depth *= 1000.0 | |
plane_depth = plane_depth.astype(np.uint16) | |
depth_name = f"depth_plane_{row.camera_name}_fr{row.frame_id:04d}.png" | |
depth_path = os.path.join(scene_path, depth_name) | |
cv2.imwrite(os.path.join(split_output_dir, depth_path), plane_depth) | |
# Meta data | |
split_meta_df.at[i, "rgb_path"] = rgb_path | |
split_meta_df.at[i, "rgb_mean"] = np.mean(rgb_int) | |
split_meta_df.at[i, "rgb_std"] = np.std(rgb_int) | |
split_meta_df.at[i, "rgb_min"] = np.min(rgb_int) | |
split_meta_df.at[i, "rgb_max"] = np.max(rgb_int) | |
split_meta_df.at[i, "depth_path"] = depth_path | |
restored_depth = plane_depth / 1000.0 | |
split_meta_df.at[i, "depth_mean"] = np.mean(restored_depth) | |
split_meta_df.at[i, "depth_std"] = np.std(restored_depth) | |
split_meta_df.at[i, "depth_min"] = np.min(restored_depth) | |
split_meta_df.at[i, "depth_max"] = np.max(restored_depth) | |
split_meta_df.at[i, "invalid_ratio"] = invalid_ratio | |
with open( | |
os.path.join(split_output_dir, f"filename_list_{split}.txt"), "w+" | |
) as f: | |
lines = split_meta_df.apply( | |
lambda r: f"{r['rgb_path']} {r['depth_path']}", axis=1 | |
).tolist() | |
f.writelines("\n".join(lines)) | |
with open( | |
os.path.join(split_output_dir, f"filename_meta_{split}.csv"), "w+" | |
) as f: | |
split_meta_df.to_csv(f, header=True) | |
print("Preprocess finished") | |