tools/process_env_maps.py

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)