import os import glob import json import cv2 import numpy as np import argparse import torch import imageio.v2 as imageio import imageio.v3 as imageio_v3 import nvdiffrast.torch as dr # Enable OpenEXR support in OpenCV os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' def swap_yz_in_extrinsic_matrix(matrix): assert matrix.shape == (4, 4), "Input must be a 4x4 matrix" new_matrix = matrix.copy() new_matrix[1, :], new_matrix[2, :] = new_matrix[2, :].copy(), new_matrix[1, :].copy() new_matrix[:, 1], new_matrix[:, 2] = new_matrix[:, 2].copy(), new_matrix[:, 1].copy() return new_matrix def swap_yz_output_in_extrinsic_matrix(matrix): assert matrix.shape == (4, 4), "Input must be a 4x4 matrix" new_matrix = matrix.copy() new_matrix[1, :], new_matrix[2, :] = new_matrix[2, :].copy(), new_matrix[1, :].copy() return new_matrix def euler_to_rotation_matrix(euler_angles, inverse_y=True, y_bias=0): alpha, gamma, beta = euler_angles beta += y_bias if inverse_y: beta = -beta R_x = np.array([[1, 0, 0], [0, np.cos(alpha), -np.sin(alpha)], [0, np.sin(alpha), np.cos(alpha)]]) R_y = np.array([[np.cos(beta), 0, np.sin(beta)], [0, 1, 0], [-np.sin(beta), 0, np.cos(beta)]]) R_z = np.array([[np.cos(gamma), -np.sin(gamma), 0], [np.sin(gamma), np.cos(gamma), 0], [0, 0, 1]]) R = np.dot(R_z, np.dot(R_y, R_x)) return R def remove_yaw_rotation(c2w_list): c2w_0 = c2w_list[0] rotation_0 = c2w_0[:3, :3] yaw_0 = np.arctan2(rotation_0[2, 0], rotation_0[0, 0]) yaw_rotation_matrix = np.array([ [np.cos(-yaw_0), 0, np.sin(-yaw_0), 0], [0, 1, 0, 0], [-np.sin(-yaw_0), 0, np.cos(-yaw_0), 0], [0, 0, 0, 1] ]) c2w_list_adjusted = [yaw_rotation_matrix @ c2w_i for c2w_i in c2w_list] return c2w_list_adjusted def adjust_yaw_rotation(c2w_list, yaw_0): yaw_rotation_matrix = np.array([ [np.cos(-yaw_0), 0, np.sin(-yaw_0), 0], [0, 1, 0, 0], [-np.sin(-yaw_0), 0, np.cos(-yaw_0), 0], [0, 0, 0, 1] ]) c2w_list_adjusted = [yaw_rotation_matrix @ c2w_i for c2w_i in c2w_list] return c2w_list_adjusted def reverse_yaw_rotation(c2w_list): c2w_list_reversed = [] for c2w in c2w_list: rotation = c2w[:3, :3] yaw = np.arctan2(rotation[2, 0], rotation[0, 0]) reversed_yaw = 2 * yaw reverse_yaw_rotation_matrix = np.array([ [np.cos(reversed_yaw), 0, np.sin(reversed_yaw), 0], [0, 1, 0, 0], [-np.sin(reversed_yaw), 0, np.cos(reversed_yaw), 0], [0, 0, 0, 1] ]) c2w_reversed = reverse_yaw_rotation_matrix @ c2w c2w_list_reversed.append(c2w_reversed) return c2w_list_reversed def prepare_camera_poses(num_frames, fixed_pose, pose_file, pose_offset, pose_reset, device, ign_camera_pose=True, swap_type=0, load_w2c=False, remove_y_rotation=False, reverse_y_rotation=False, yaw_0=0, pose_list=None, rotation_euler=None): """Prepare camera poses based on the provided arguments.""" if pose_list is not None: c2w_list = pose_list for frame_idx, transform_matrix in enumerate(c2w_list): if swap_type == 0: new_transform_matrix = transform_matrix elif swap_type == 1: new_transform_matrix = swap_yz_in_extrinsic_matrix(transform_matrix) elif swap_type == 2: new_transform_matrix = swap_yz_output_in_extrinsic_matrix(transform_matrix) if load_w2c: new_transform_matrix = np.linalg.inv(new_transform_matrix) c2w_list[frame_idx] = new_transform_matrix if pose_reset: w2c_0 = np.linalg.inv(c2w_list[0]) c2w_list = [w2c_0 @ c2w_i for c2w_i in c2w_list] if remove_y_rotation: c2w_list = remove_yaw_rotation(c2w_list) if reverse_y_rotation: c2w_list = reverse_yaw_rotation(c2w_list) if rotation_euler is not None: for frame_idx, transform_matrix in enumerate(c2w_list): rotation_matrix = euler_to_rotation_matrix(rotation_euler, inverse_y=True, y_bias=np.pi/2) if ign_camera_pose: transform_matrix = c2w_list[0] else: transform_matrix = c2w rotation_matrix_4x4 = np.eye(4) rotation_matrix_4x4[:3, :3] = rotation_matrix new_transform_matrix = np.dot(rotation_matrix_4x4, transform_matrix) c2w_list[frame_idx] = new_transform_matrix return c2w_list elif fixed_pose or pose_file is None: return [np.eye(4) for _ in range(num_frames)] with open(pose_file, 'r') as f: meta = json.load(f) frames = meta['frames'][pose_offset:pose_offset + num_frames] for frame_idx, data in enumerate(frames): transform_matrix = np.array(data["transform_matrix"]) if swap_type == 0: new_transform_matrix = transform_matrix elif swap_type == 1: new_transform_matrix = swap_yz_in_extrinsic_matrix(transform_matrix) if load_w2c: new_transform_matrix = np.linalg.inv(new_transform_matrix) data["transform_matrix"] = new_transform_matrix.tolist() frames[frame_idx] = data if ign_camera_pose: c2w_list = [np.array(frames[0]["transform_matrix"]) for frame in frames] else: c2w_list = [np.array(frame['transform_matrix']) for frame in frames] if pose_reset: w2c_0 = np.linalg.inv(c2w_list[0]) c2w_list = [w2c_0 @ c2w_i for c2w_i in c2w_list] # compute c2c0 if remove_y_rotation: c2w_list = remove_yaw_rotation(c2w_list) if reverse_y_rotation: c2w_list = reverse_yaw_rotation(c2w_list) if yaw_0 != 0: c2w_list = adjust_yaw_rotation(c2w_list, yaw_0) for frame_idx, (data, c2w) in enumerate(zip(frames, c2w_list)): if "hdri_euler" in data.keys(): rotation_matrix = euler_to_rotation_matrix(data["hdri_euler"], inverse_y=False) if ign_camera_pose: transform_matrix = c2w_list[0] else: transform_matrix = c2w rotation_matrix_4x4 = np.eye(4) rotation_matrix_4x4[:3, :3] = rotation_matrix new_transform_matrix = np.dot(rotation_matrix_4x4, transform_matrix) c2w_list[frame_idx] = new_transform_matrix else: print(f"Warning: 'hdri_euler' not found in frame {frame_idx} of {pose_file}. Using original transform matrix.") break return c2w_list def latlong_vec(res, device=None): gy, gx = torch.meshgrid(torch.linspace( 0.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device=device), torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device=device), indexing='ij') sintheta, costheta = torch.sin(gy*np.pi), torch.cos(gy*np.pi) sinphi, cosphi = torch.sin(gx*np.pi), torch.cos(gx*np.pi) dir_vec = torch.stack(( sintheta*sinphi, costheta, -sintheta*cosphi ), dim=-1) # return dr.texture(cubemap[None, ...], dir_vec[None, ...].contiguous(), filter_mode='linear', boundary_mode='cube')[0] return dir_vec #[H, W, 3] def envmap_vec(res, device=None): return -latlong_vec(res, device).flip(0).flip(1) #[H, W, 3] def dot(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return torch.sum(x*y, -1, keepdim=True) def length(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor: return torch.sqrt(torch.clamp(dot(x,x), min=eps)) # Clamp to avoid nan gradients because grad(sqrt(0)) = NaN def safe_normalize(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor: return x / length(x, eps) def cube_to_dir(s, x, y): if s == 0: rx, ry, rz = torch.ones_like(x), -y, -x elif s == 1: rx, ry, rz = -torch.ones_like(x), -y, x elif s == 2: rx, ry, rz = x, torch.ones_like(x), y elif s == 3: rx, ry, rz = x, -torch.ones_like(x), -y elif s == 4: rx, ry, rz = x, -y, torch.ones_like(x) elif s == 5: rx, ry, rz = -x, -y, -torch.ones_like(x) return torch.stack((rx, ry, rz), dim=-1) def latlong_to_cubemap(latlong_map, res): cubemap = torch.zeros(6, res[0], res[1], latlong_map.shape[-1], dtype=torch.float32, device='cuda') for s in range(6): gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device='cuda'), torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device='cuda'), indexing='ij') v = safe_normalize(cube_to_dir(s, gx, gy)) tu = torch.atan2(v[..., 0:1], -v[..., 2:3]) / (2 * np.pi) + 0.5 tv = torch.acos(torch.clamp(v[..., 1:2], min=-1, max=1)) / np.pi texcoord = torch.cat((tu, tv), dim=-1) cubemap[s, ...] = dr.texture(latlong_map[None, ...], texcoord[None, ...], filter_mode='linear')[0] return cubemap def load_and_preprocess_hdr(hdr_dir, env_strength, env_flip, env_rot, device, rotation180=False, inverse_env=False, flip_env=False): """Load and preprocess the HDR environment map.""" if hdr_dir.endswith('.hdr') or hdr_dir.endswith('.exr'): latlong_img = imageio_v3.imread(hdr_dir, flags=cv2.IMREAD_UNCHANGED, plugin='opencv') elif hdr_dir.endswith('.jpg') or hdr_dir.endswith('.png'): import skimage latlong_img = skimage.io.imread(hdr_dir)[..., :3] latlong_img = skimage.img_as_float(latlong_img) latlong_img = np.power(latlong_img, 2.4).astype(np.float32) latlong_img *= 2 # for mit dataset if rotation180: height, width, channels = latlong_img.shape shift_amount = width // 2 shifted_hdr = np.zeros_like(latlong_img) shifted_hdr[:, -shift_amount:, :] = latlong_img[:, :shift_amount, :] shifted_hdr[:, :-shift_amount, :] = latlong_img[:, shift_amount:, :] latlong_img = shifted_hdr if inverse_env: latlong_img = latlong_img[:, ::-1, :] height, width, channels = latlong_img.shape shift_amount = width // 2 shifted_hdr = np.zeros_like(latlong_img) shifted_hdr[:, -shift_amount:, :] = latlong_img[:, :shift_amount, :] shifted_hdr[:, :-shift_amount, :] = latlong_img[:, shift_amount:, :] latlong_img = shifted_hdr if flip_env: latlong_img = latlong_img[:, ::-1, :].copy() latlong_img = torch.tensor(latlong_img, dtype=torch.float32, device=device) latlong_img *= env_strength # Cleanup NaNs and Infs latlong_img = torch.nan_to_num(latlong_img, nan=0.0, posinf=65504.0, neginf=0.0) latlong_img = latlong_img.clamp(0.0, 65504.0) if env_flip: latlong_img = torch.flip(latlong_img, dims=[1]) if env_rot != 0: lat_h, lat_w = latlong_img.shape[:2] pixel_rot = int(lat_w * env_rot / 360) latlong_img = torch.roll(latlong_img, shifts=pixel_rot, dims=1) # Convert to cubemap cubemap = latlong_to_cubemap(latlong_img, [512, 512]) return cubemap def prepare_metadata(hdr_dir, env_rot, env_flip, env_strength, fixed_pose, rotate_envlight, save_dir, prefix): """Prepare metadata about the environment map processing.""" env_meta = { 'envmap': os.path.basename(hdr_dir), 'envmap_rot': env_rot, 'envmap_flip': env_flip, 'envmap_strength': env_strength, 'fixed_pose': fixed_pose, 'rotate_envlight': rotate_envlight, } if save_dir: os.makedirs(save_dir, exist_ok=True) meta_path = os.path.join(save_dir, f'{prefix}.meta.json') with open(meta_path, 'w') as f: json.dump(env_meta, f, indent=4) return env_meta def rotate_y(a, device=None): s, c = np.sin(a), np.cos(a) return torch.tensor([[ c, 0, s, 0], [ 0, 1, 0, 0], [-s, 0, c, 0], [ 0, 0, 0, 1]], dtype=torch.float32, device=device) def process_projected_envmap(cubemap, vec, c2w, y_rot, H, W): """Process the camera-oriented projected environment map.""" vec_cam = vec.view(-1, 3) @ c2w[:3, :3].T vec_query = (vec_cam @ y_rot[:3, :3].T).view(1, H, W, 3) env_proj = dr.texture(cubemap.unsqueeze(0), -vec_query.contiguous(), filter_mode='linear', boundary_mode='cube')[0] env_proj = torch.flip(env_proj, dims=[0, 1]) return env_proj def rgb2srgb(rgb): return torch.where(rgb <= 0.0031308, 12.92 * rgb, 1.055 * rgb**(1/2.4) - 0.055) def reinhard(x, max_point=16): y_rein = x * (1 + x / (max_point ** 2)) / (1 + x) return y_rein def hdr_mapping(env_hdr, log_scale): """Map HDR environment maps to LDR and logarithmic representations.""" env_ev0 = rgb2srgb(reinhard(env_hdr, max_point=16).clamp(0, 1)) env_log = rgb2srgb(torch.log1p(env_hdr) / np.log1p(log_scale)).clamp(0, 1) return { 'env_hdr': env_hdr, # Original HDR image 'env_ev0': env_ev0, # LDR image after tone mapping 'env_log': env_log, # Logarithmic scaling } def process_environment_map( hdr_dir, resolution=(512, 512), num_frames=1, fixed_pose=True, pose_file=None, pose_list=None, rotation_euler=None, pose_offset=0, pose_reset=False, rotate_envlight=False, env_format=['proj'], log_scale=10000, env_strength=1.0, env_flip=True, env_rot=180.0, save_dir=None, prefix='0000', device=None, rotation180=False, ign_camera_pose=True, inverse_env=False, swap_type=0, load_w2c=False, flip_env=False, remove_y_rotation=False, reverse_y_rotation=False, yaw_0=0, ): """ Preprocess HDR environment maps for rendering. FIXME: Note that this function bakes in a flip and rotate operation for the environment light. Set to env_flip=True and env_rot=180 is considered as loading the original environment map. Args: hdr_dir (str): Path to the HDR environment map file. resolution (tuple of int): Resolution of the output images (H, W). num_frames (int): Number of frames to process. fixed_pose (bool): Use a fixed camera pose (identity matrix) if True. pose_file (str): Path to the camera pose file (JSON). pose_offset (int): Offset for the pose frames in the pose file. pose_reset (bool): Reset camera poses to be relative to the first frame. rotate_envlight (bool): Rotate the environment light over frames if True. env_format (list of str): Formats of the environment maps to generate ('proj', 'fixed', 'ball'). log_scale (int): Log scale factor for HDR mapping. env_strength (float): Strength multiplier for the environment map. env_flip (bool): Flip the environment map horizontally if True. env_rot (float): Rotation angle for the environment map in degrees. save_dir (str): Directory to save the processed images (optional). prefix (str): Prefix for the output files (used if saving images). Returns: dict: A dictionary containing the processed environment maps and metadata. { 'metadata': env_meta, 'fixed': mapping_results_for_fixed_envmap, # Only if 'fixed' in env_format 'env_ldr': stacked_tensor_of_proj_env_ldr, # Only if 'proj' in env_format 'env_log': stacked_tensor_of_proj_env_log, # Only if 'proj' in env_format 'ball_env_ldr': stacked_tensor_of_ball_env_ldr, # Only if 'ball' in env_format 'ball_env_log': stacked_tensor_of_ball_env_log, # Only if 'ball' in env_format } Tensors are with shape (T, H, W, 3) in [0, 1] """ H, W = resolution # (704, 1280) if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu' vec = latlong_vec((H, W), device=device) # Prepare camera poses poses = prepare_camera_poses( num_frames=num_frames, fixed_pose=fixed_pose, pose_file=pose_file, pose_offset=pose_offset, pose_reset=pose_reset, device=device, ign_camera_pose=ign_camera_pose, swap_type=swap_type, load_w2c=load_w2c, remove_y_rotation=remove_y_rotation, reverse_y_rotation=reverse_y_rotation, yaw_0=yaw_0, pose_list=pose_list, rotation_euler=rotation_euler, ) # Prepare rotations for the environment light # 57 * rot rots = np.linspace(0, 2 * np.pi, num_frames) if rotate_envlight else [0] * num_frames # Load and preprocess the HDR environment map cubemap = load_and_preprocess_hdr( hdr_dir=hdr_dir, env_strength=env_strength, env_flip=env_flip, env_rot=env_rot, device=device, rotation180=rotation180, inverse_env=inverse_env, flip_env=flip_env, ) # Prepare metadata env_meta = prepare_metadata( hdr_dir=hdr_dir, env_rot=env_rot, env_flip=env_flip, env_strength=env_strength, fixed_pose=fixed_pose, rotate_envlight=rotate_envlight, save_dir=save_dir, prefix=prefix ) # Initialize result dictionary results = { 'metadata': env_meta, } # Prepare lists to collect per-frame tensors if 'proj' in env_format: proj_env_ldr_list = [] proj_env_log_list = [] # Process per-frame environment maps for i in range(num_frames): c2w = torch.from_numpy(poses[i]).float().to(device) y_rot = rotate_y(rots[i], device=device) if 'proj' in env_format: env_proj = process_projected_envmap(cubemap, vec, c2w, y_rot, H, W) mapping_results = hdr_mapping(env_proj, log_scale=log_scale) proj_env_ldr_list.append(mapping_results['env_ev0']) proj_env_log_list.append(mapping_results['env_log']) if 'proj' in env_format: results['env_ldr'] = torch.stack(proj_env_ldr_list, dim=0) results['env_log'] = torch.stack(proj_env_log_list, dim=0) return results def save_array_as_video(video_array, output_path: str, fps: int = 24): """ video_array: t h w c, np.array or tensors """ if isinstance(video_array, torch.Tensor): video_array = video_array.cpu().numpy() if video_array.dtype != np.uint8: print("float 2 uint8") # If the data range is [-1, 1] if video_array.min() < 0: video_array = ((video_array + 1) * 127.5).clip(0, 255).astype(np.uint8) # If the data range is [0, 1] else: video_array = (video_array * 255).clip(0, 255).astype(np.uint8) try: if not os.path.isfile(output_path): imageio.mimsave( output_path, [frame for frame in video_array], fps=fps, codec='libx264' ) print("succeed to save vide") print(f"video already exists in {output_path}") except Exception as e: print(f"fail to save video: {e}") def process_hdr(hdr_path, save_path, env_strength=1.0, inverse_env=False): hdr_path = glob.glob(f'{hdr_path}*')[0] if '.hdr' in hdr_path: env_strength = env_strength / 3.0 num_of_frames = 57 ldr_list = [] hdr_log_list = [] env_dir_list = [] envlight_dict = process_environment_map( hdr_dir=hdr_path, resolution=(320, 576), num_frames=num_of_frames, # 1 for mit dataset fixed_pose=True, rotate_envlight=False, env_format=['proj', ], device='cuda', rotation180=False, inverse_env=inverse_env, # True for mit dataset, False for others log_scale=60000, env_strength=env_strength, # 1.0 for mit dataset ) # Tensors are with shape (T, H, W, 3) in [0, 1] ldr_list = (envlight_dict['env_ldr'].cpu().numpy() * 255).astype(np.uint8) hdr_log_list = (envlight_dict['env_log'].cpu().numpy() * 255).astype(np.uint8) env_nrm = ((envmap_vec([320, 576], device='cpu').cpu().numpy()*0.5 + 0.5) * 255).astype(np.uint8) for _ in range(num_of_frames): env_dir_list.append(env_nrm) os.makedirs(save_path, exist_ok=True) ldr_video_path = os.path.join(save_path, "ldr_video_fix_first_frame.mp4") hdr_log_video_path = os.path.join(save_path, "hdr_log_video_fix_first_frame.mp4") env_dir_video_path = os.path.join(save_path, "env_dir_video_fix_first_frame.mp4") if os.path.exists(ldr_video_path): os.remove(ldr_video_path) if os.path.exists(hdr_log_video_path): os.remove(hdr_log_video_path) if os.path.exists(env_dir_video_path): os.remove(env_dir_video_path) save_array_as_video(np.array(ldr_list),ldr_video_path) save_array_as_video(np.array(hdr_log_list),hdr_log_video_path) save_array_as_video(np.array(env_dir_list),env_dir_video_path) def parse_arguments() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Env maps processing script") # Specific arguments parser.add_argument( "--env_dir", type=str, default=None, help="Path to the directory containing the environment map." ) parser.add_argument( "--save_path", type=str, default=None, help="Path to the directory where the processed environment maps will be saved." ) parser.add_argument( "--env_strength", type=float, default=1.0, help="Strength of the environment map." ) return parser.parse_args() if __name__ == "__main__": args = parse_arguments() process_hdr(args.env_dir, save_path=args.save_path, env_strength=args.env_strength)