blender.py

# ------------------------------------------------------------------------------------
# NeRF-Factory
# Copyright (c) 2022 POSTECH, KAIST, Kakao Brain Corp. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------------------
# ------------------------------------------------------------------------------------
# Modified from NeRF (https://github.com/bmild/nerf)
# Copyright (c) 2020 Google LLC. All Rights Reserved.
# ------------------------------------------------------------------------------------

import json
import os

import gdown
import imageio
import numpy as np
import torch

trans_t = lambda t: torch.tensor(
    [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, t], [0, 0, 0, 1]]
).float()

rot_phi = lambda phi: torch.tensor(
    [
        [1, 0, 0, 0],
        [0, np.cos(phi), -np.sin(phi), 0],
        [0, np.sin(phi), np.cos(phi), 0],
        [0, 0, 0, 1],
    ]
).float()

rot_theta = lambda th: torch.tensor(
    [
        [np.cos(th), 0, -np.sin(th), 0],
        [0, 1, 0, 0],
        [np.sin(th), 0, np.cos(th), 0],
        [0, 0, 0, 1],
    ]
).float()


def pose_spherical(theta, phi, radius):
    c2w = trans_t(radius)
    c2w = rot_phi(phi / 180.0 * np.pi) @ c2w
    c2w = rot_theta(theta / 180.0 * np.pi) @ c2w
    c2w = (
        torch.tensor([[-1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]).float()
        @ c2w
    )
    return c2w


def load_blender_data(
    datadir: str,
    scene_name: str,
    compare_method: str,
    train_skip: int,
    val_skip: int,
    test_skip: int,
    cam_scale_factor: float,
    white_bkgd: bool,
):
    basedir = os.path.join(datadir, scene_name)
    cam_trans = np.diag(np.array([1, -1, -1, 1], dtype=np.float32))
    
    splits = ["train", "val", "test"]
    metas = {}
    # for s in splits:
    #     with open(os.path.join(basedir, compare_method, "transforms_{}_gt.json".format(s)), "r") as fp:
    #         metas[s] = json.load(fp)
    
    #print('compare_method', compare_method)
    fp_train = open(os.path.join(basedir, compare_method, "transforms_train.json"), "r") # training file
    fp_val = open(os.path.join(basedir, "transforms_val.json"), "r") # evaluation file
    fp_test = open(os.path.join(basedir, "transforms_test.json"), "r")  # testing file
    metas["train"], metas["val"], metas["test"] = json.load(fp_train), json.load(fp_val), json.load(fp_test)

    images = []
    extrinsics = []
    counts = [0]

    for s in splits:    # train/ val/ test
        meta = metas[s]
        imgs = []
        poses = []

        if s == "train":
            skip = train_skip
        elif s == "val":
            skip = val_skip
        elif s == "test":
            skip = test_skip

        for frame in meta["frames"][::skip]:
            fname = os.path.join(basedir, frame["file_path"]) 
            imgs.append(imageio.imread(fname))
            poses.append(np.array(frame["transform_matrix"]))
        imgs = (np.array(imgs) / 255.0).astype(np.float32)  # keep all 4 channels (RGBA)
        poses = np.array(poses).astype(np.float32)
        counts.append(counts[-1] + imgs.shape[0])
        images.append(imgs)
        extrinsics.append(poses)
    
    i_split = [np.arange(counts[i], counts[i + 1]) for i in range(3)]
    
    images = np.concatenate(images, 0)
    extrinsics = np.concatenate(extrinsics, 0)

    extrinsics[:, :3, 3] *= cam_scale_factor
    extrinsics = extrinsics @ cam_trans

    h, w = imgs[0].shape[:2]
    num_frame = len(extrinsics)
    i_split += [np.arange(num_frame)]

    camera_angle_x = float(meta["camera_angle_x"])
    focal = 0.5 * w / np.tan(0.5 * camera_angle_x)
    intrinsics = np.array(
        [
            [[focal, 0.0, 0.5 * w], [0.0, focal, 0.5 * h], [0.0, 0.0, 1.0]]
            for _ in range(num_frame)
        ]
    )
    image_sizes = np.array([[h, w] for _ in range(num_frame)])

    # rendering poses in LOM dataset
    render_poses = torch.stack(
        [
            pose_spherical(angle, -10.0, 4.0) @ cam_trans
            for angle in np.linspace(70, 110, 30 + 1)[:-1]
        ],
        0,
    ) 

    render_poses2 = torch.stack([
            pose_spherical(110, angle2, 4.0) @ cam_trans
            for angle2 in np.linspace(-10, 10, 20 + 1)[:-1]
        ],
        0,)

    render_poses3 = torch.stack([
            pose_spherical(110*i, 10, 4.0*i) @ cam_trans
            for i in np.linspace(1, 0.8, 20 + 1)[:-1]
        ],
        0,)

    render_poses = torch.cat((render_poses,render_poses2,render_poses3), 0)

    render_poses[:, :3, 3] *= cam_scale_factor
    near = 2.0
    far = 6.0

    if white_bkgd:
        images = images[..., :3] * images[..., -1:] + (1.0 - images[..., -1:])
    else:
        images = images[..., :3]
        
    return (
        images,
        intrinsics,
        extrinsics,
        image_sizes,
        near,
        far,
        (-1, -1),
        i_split,
        render_poses,
    )