<filename>lib/data/blender.py
import os
import json
import math

import torch

from PIL import Image

from torchvision.transforms.functional import to_tensor

from . import BaseData


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

rot_phi = lambda phi : torch.Tensor([[1,             0,              0, 0],
                                     [0, math.cos(phi), -math.sin(phi), 0],
                                     [0, math.sin(phi),  math.cos(phi), 0],
                                     [0,             0,              0, 1]])

rot_theta = lambda th : torch.Tensor([[math.cos(th), 0, -math.sin(th), 0],
                                      [           0, 1,             0, 0],
                                      [math.sin(th), 0,  math.cos(th), 0],
                                      [           0, 0,             0, 1]])


def pose_spherical(theta, phi, radius):
    permute = torch.Tensor([[-1, 0, 0, 0],
                            [ 0, 0, 1, 0],
                            [ 0, 1, 0, 0],
                            [ 0, 0, 0, 1]])
    camera2world = permute @ rot_theta(theta/180*math.pi) @ rot_phi(phi/180*math.pi) @ trans_t(radius)
    return camera2world


def load_blender_data(root, split='train', skip=1, white_bkgd=True, factor=None):
    """
    Args:
        root: str
            path to dataset
        split: str
            'train', 'val', 'test'
        skip: int
            sampling factor
        white_bkgd: bool
            if True, replace alpha channel with white
        factor: int
            scaling factor

    Outputs:
        images: torch.Tensor
            tensor of shape (n_images, height, width, 3)
        poses: torch.Tensor
            tensor of shape (n_images, 4, 4)
        render_poses: torch.Tensor
            tensor of shape (40, 4, 4)
        height, width: int
            image height and width
        focal_length: float
    """
    with open(os.path.join(root, f'transforms_{split}.json'), 'r') as f:
        meta = json.load(f)
    
    images = []
    poses = []
    for frame in meta['frames'][::skip]:
        images.append(to_tensor(Image.open(os.path.join(root, frame['file_path'] + '.png'))))
        poses.append(torch.Tensor(frame['transform_matrix']))

    images = torch.stack(images)
    poses = torch.stack(poses)
    if factor is not None:
        images = torch.nn.functional.interpolate(images, scale_factor=1/factor, mode='area')
    height, width = images.shape[-2:]
    images = images.permute(0, 2, 3, 1) # (N, 3, H, W) -> (N, H, W, 4)
    camera_angle_x = float(meta['camera_angle_x'])
    focal_length = 0.5 * width / math.tan(0.5 * camera_angle_x)

    render_poses = torch.stack([pose_spherical(angle, -30, 4)
                                for angle in torch.linspace(-180, 180, 41)[:-1]], 0)

    if white_bkgd: # alpha channel as white
        images = images[..., :3] * images[..., -1:] + (1 - images[..., -1:])
    else: # alpha channel as black
        images = images[..., :3]

    return images, poses, render_poses, (height, width, focal_length)


class Blender(BaseData):
    def __init__(
            self,
            root,
            split='train',
            skip=1,
            near=None,
            far=None,
            pixel_as_sample=False,
            factor=None,
            ndc=False,
            white_bkgd=True):
        images, camera2world, render_poses, meta = load_blender_data(root, split, skip, white_bkgd, factor)

        self.render_poses = render_poses # used to render novel view

        super().__init__(images, camera2world, meta, pixel_as_sample)

        if ndc:
            self.near = 0
            self.far = 1
            self.ray_d, self.ray_o = self.ndc_ray(*meta, self.ray_o, self.ray_d)
        else:
            self.near = near or 2
            self.far = far or 6

<filename>lib/data/blender.py
import os
import json
import math

import torch

from PIL import Image

from torchvision.transforms.functional import to_tensor

from . import BaseData


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

rot_phi = lambda phi : torch.Tensor([[1,             0,              0, 0],
                                     [0, math.cos(phi), -math.sin(phi), 0],
                                     [0, math.sin(phi),  math.cos(phi), 0],
                                     [0,             0,              0, 1]])

rot_theta = lambda th : torch.Tensor([[math.cos(th), 0, -math.sin(th), 0],
                                      [           0, 1,             0, 0],
                                      [math.sin(th), 0,  math.cos(th), 0],
                                      [           0, 0,             0, 1]])


def pose_spherical(theta, phi, radius):
    permute = torch.Tensor([[-1, 0, 0, 0],
                            [ 0, 0, 1, 0],
                            [ 0, 1, 0, 0],
                            [ 0, 0, 0, 1]])
    camera2world = permute @ rot_theta(theta/180*math.pi) @ rot_phi(phi/180*math.pi) @ trans_t(radius)
    return camera2world


def load_blender_data(root, split='train', skip=1, white_bkgd=True, factor=None):
    """
    Args:
        root: str
            path to dataset
        split: str
            'train', 'val', 'test'
        skip: int
            sampling factor
        white_bkgd: bool
            if True, replace alpha channel with white
        factor: int
            scaling factor

    Outputs:
        images: torch.Tensor
            tensor of shape (n_images, height, width, 3)
        poses: torch.Tensor
            tensor of shape (n_images, 4, 4)
        render_poses: torch.Tensor
            tensor of shape (40, 4, 4)
        height, width: int
            image height and width
        focal_length: float
    """
    with open(os.path.join(root, f'transforms_{split}.json'), 'r') as f:
        meta = json.load(f)
    
    images = []
    poses = []
    for frame in meta['frames'][::skip]:
        images.append(to_tensor(Image.open(os.path.join(root, frame['file_path'] + '.png'))))
        poses.append(torch.Tensor(frame['transform_matrix']))

    images = torch.stack(images)
    poses = torch.stack(poses)
    if factor is not None:
        images = torch.nn.functional.interpolate(images, scale_factor=1/factor, mode='area')
    height, width = images.shape[-2:]
    images = images.permute(0, 2, 3, 1) # (N, 3, H, W) -> (N, H, W, 4)
    camera_angle_x = float(meta['camera_angle_x'])
    focal_length = 0.5 * width / math.tan(0.5 * camera_angle_x)

    render_poses = torch.stack([pose_spherical(angle, -30, 4)
                                for angle in torch.linspace(-180, 180, 41)[:-1]], 0)

    if white_bkgd: # alpha channel as white
        images = images[..., :3] * images[..., -1:] + (1 - images[..., -1:])
    else: # alpha channel as black
        images = images[..., :3]

    return images, poses, render_poses, (height, width, focal_length)


class Blender(BaseData):
    def __init__(
            self,
            root,
            split='train',
            skip=1,
            near=None,
            far=None,
            pixel_as_sample=False,
            factor=None,
            ndc=False,
            white_bkgd=True):
        images, camera2world, render_poses, meta = load_blender_data(root, split, skip, white_bkgd, factor)

        self.render_poses = render_poses # used to render novel view

        super().__init__(images, camera2world, meta, pixel_as_sample)

        if ndc:
            self.near = 0
            self.far = 1
            self.ray_d, self.ray_o = self.ndc_ray(*meta, self.ray_o, self.ray_d)
        else:
            self.near = near or 2
            self.far = far or 6

Args:         root: str             path to dataset         split: str             'train', 'val', 'test'         skip: int             sampling factor         white_bkgd: bool             if True, replace alpha channel with white         factor: int             scaling factor      Outputs:         images: torch.Tensor             tensor of shape (n_images, height, width, 3)         poses: torch.Tensor             tensor of shape (n_images, 4, 4)         render_poses: torch.Tensor             tensor of shape (40, 4, 4)         height, width: int             image height and width         focal_length: float # (N, 3, H, W) -> (N, H, W, 4) # alpha channel as white # alpha channel as black # used to render novel view