ascarlettvfx's picture
Upload 99 files
660acc7 verified
# 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")