utilities/camera.py

import glm
import torch
import random

import numpy as np
import torchvision.transforms as transforms

from .resize_right import resize

blurs = [
    transforms.Compose([
        transforms.GaussianBlur(11, sigma=(5, 5))
    ]),
    transforms.Compose([
        transforms.GaussianBlur(11, sigma=(2, 2))
    ]),
    transforms.Compose([
        transforms.GaussianBlur(5, sigma=(5, 5))
    ]),
    transforms.Compose([
        transforms.GaussianBlur(5, sigma=(2, 2))
    ]),
]

def get_random_bg(h, w, rand_solid=False):

        p = torch.rand(1)

        if p > 0.66666:
            if rand_solid:
                background = torch.vstack([
                    torch.full( (1, h, w), torch.rand(1).item()),
                    torch.full( (1, h, w), torch.rand(1).item()),
                    torch.full( (1, h, w), torch.rand(1).item()),
                ]).unsqueeze(0) + torch.rand(1, 3, h, w)
                background = ((background - background.amin()) / (background.amax() - background.amin()))
                background = blurs[random.randint(0, 3)](background).permute(0, 2, 3, 1)
            else:
                background =  blurs[random.randint(0, 3)]( torch.rand((1, 3, h, w)) ).permute(0, 2, 3, 1)
        elif p > 0.333333:
            size = random.randint(5, 10)
            background = torch.vstack([
                torch.full( (1, size, size), torch.rand(1).item() / 2),
                torch.full( (1, size, size), torch.rand(1).item() / 2 ),
                torch.full( (1, size, size), torch.rand(1).item() / 2 ),
            ]).unsqueeze(0)

            second = torch.rand(3)

            background[:, 0, ::2, ::2] = second[0]
            background[:, 1, ::2, ::2] = second[1]
            background[:, 2, ::2, ::2] = second[2]

            background[:, 0, 1::2, 1::2] = second[0]
            background[:, 1, 1::2, 1::2] = second[1]
            background[:, 2, 1::2, 1::2] = second[2]

            background = blurs[random.randint(0, 3)]( resize(background, out_shape=(h, w)) )

            background = background.permute(0, 2, 3, 1)

        else:
            background = torch.vstack([
                torch.full( (1, h, w), torch.rand(1).item()),
                torch.full( (1, h, w), torch.rand(1).item()),
                torch.full( (1, h, w), torch.rand(1).item()),
            ]).unsqueeze(0).permute(0, 2, 3, 1)

        return background

def cosine_sample(N : np.ndarray) -> np.ndarray:
    """
    #----------------------------------------------------------------------------
    # Cosine sample around a vector N
    #----------------------------------------------------------------------------

    Copied from nvdiffmodelling

    """
    # construct local frame
    N = N/np.linalg.norm(N)

    dx0 = np.array([0, N[2], -N[1]])
    dx1 = np.array([-N[2], 0, N[0]])

    dx = dx0 if np.dot(dx0,dx0) > np.dot(dx1,dx1) else dx1
    dx = dx/np.linalg.norm(dx)
    dy = np.cross(N,dx)
    dy = dy/np.linalg.norm(dy)

    # cosine sampling in local frame
    phi = 2.0*np.pi*np.random.uniform()
    s = np.random.uniform()
    costheta = np.sqrt(s)
    sintheta = np.sqrt(1.0 - s)

    # cartesian vector in local space
    x = np.cos(phi)*sintheta
    y = np.sin(phi)*sintheta
    z = costheta

    # local to world
    return dx*x + dy*y + N*z

def persp_proj(fov_x=45, ar=1, near=1.0, far=50.0):
    """
    From https://github.com/rgl-epfl/large-steps-pytorch by @bathal1 (Baptiste Nicolet)

    Build a perspective projection matrix.
    Parameters
    ----------
    fov_x : float
        Horizontal field of view (in degrees).
    ar : float
        Aspect ratio (w/h).
    near : float
        Depth of the near plane relative to the camera.
    far : float
        Depth of the far plane relative to the camera.
    """
    fov_rad = np.deg2rad(fov_x)

    tanhalffov = np.tan( (fov_rad / 2) )
    max_y = tanhalffov * near
    min_y = -max_y
    max_x = max_y * ar
    min_x = -max_x

    z_sign = -1.0
    proj_mat = np.array([[0, 0, 0, 0],
                        [0, 0, 0, 0],
                        [0, 0, 0, 0],
                        [0, 0, 0, 0]])

    proj_mat[0, 0] = 2.0 * near / (max_x - min_x)
    proj_mat[1, 1] = 2.0 * near / (max_y - min_y)
    proj_mat[0, 2] = (max_x + min_x) / (max_x - min_x)
    proj_mat[1, 2] = (max_y + min_y) / (max_y - min_y)
    proj_mat[3, 2] = z_sign

    proj_mat[2, 2] = z_sign * far / (far - near)
    proj_mat[2, 3] = -(far * near) / (far - near)
    
    return proj_mat

def get_camera_params(elev_angle, azim_angle, distance, resolution, fov=60, look_at=[0, 0, 0], up=[0, -1, 0]):
    
    elev = np.radians( elev_angle )
    azim = np.radians( azim_angle ) 
    
    # Generate random view
    cam_z = distance * np.cos(elev) * np.sin(azim)
    cam_y = distance * np.sin(elev)
    cam_x = distance * np.cos(elev) * np.cos(azim)

    modl = glm.mat4()
    view  = glm.lookAt(
        glm.vec3(cam_x, cam_y, cam_z),
        glm.vec3(look_at[0], look_at[1], look_at[2]),
        glm.vec3(up[0], up[1], up[2]),
    )

    a_mv = view * modl
    a_mv = np.array(a_mv.to_list()).T
    proj_mtx = persp_proj(fov)
    
    a_mvp = np.matmul(proj_mtx, a_mv).astype(np.float32)[None, ...]
    
    a_lightpos = np.linalg.inv(a_mv)[None, :3, 3]
    a_campos = a_lightpos

    return {
        'mvp' : a_mvp,
        'lightpos' : a_lightpos,
        'campos' : a_campos,
        'resolution' : [resolution, resolution], 
        }

# Returns a batch of camera parameters
class CameraBatch(torch.utils.data.Dataset):
    def __init__(
        self,
        image_resolution,
        distances,
        azimuths,
        elevation_params,
        fovs,
        aug_loc, 
        aug_light,
        aug_bkg,
        bs,
        look_at=[0, 0, 0], up=[0, -1, 0],
        rand_solid=False
    ):

        self.res = image_resolution

        self.dist_min = distances[0]
        self.dist_max = distances[1]

        self.azim_min = azimuths[0]
        self.azim_max = azimuths[1]

        self.fov_min = fovs[0]
        self.fov_max = fovs[1]
        
        self.elev_alpha = elevation_params[0]
        self.elev_beta  = elevation_params[1]
        self.elev_max   = elevation_params[2]

        self.aug_loc   = aug_loc
        self.aug_light = aug_light
        self.aug_bkg   = aug_bkg

        self.look_at = look_at
        self.up = up

        self.batch_size = bs
        self.rand_solid = rand_solid

    def __len__(self):
        return self.batch_size
        
    def __getitem__(self, index):

        elev = np.radians( np.random.beta( self.elev_alpha, self.elev_beta ) * self.elev_max )
        azim = np.radians( np.random.uniform( self.azim_min, self.azim_max+1.0 ) )
        dist = np.random.uniform( self.dist_min, self.dist_max )
        fov = np.random.uniform( self.fov_min, self.fov_max )
        
        proj_mtx = persp_proj(fov)
        
        # Generate random view
        cam_z = dist * np.cos(elev) * np.sin(azim)
        cam_y = dist * np.sin(elev)
        cam_x = dist * np.cos(elev) * np.cos(azim)
        
        if self.aug_loc:

            # Random offset
            limit  = self.dist_min // 2
            rand_x = np.random.uniform( -limit, limit )
            rand_y = np.random.uniform( -limit, limit )

            modl = glm.translate(glm.mat4(), glm.vec3(rand_x, rand_y, 0))

        else:
        
            modl = glm.mat4()
            
        view  = glm.lookAt(
            glm.vec3(cam_x, cam_y, cam_z),
            glm.vec3(self.look_at[0], self.look_at[1], self.look_at[2]),
            glm.vec3(self.up[0], self.up[1], self.up[2]),
        )

        r_mv = view * modl
        r_mv = np.array(r_mv.to_list()).T

        mvp     = np.matmul(proj_mtx, r_mv).astype(np.float32)
        campos  = np.linalg.inv(r_mv)[:3, 3]

        if self.aug_light:
            lightpos = cosine_sample(campos)*dist
        else:
            lightpos = campos*dist

        if self.aug_bkg:
            bkgs = get_random_bg(self.res, self.res, self.rand_solid).squeeze(0)
        else:
            bkgs = torch.ones(self.res, self.res, 3)

        return {
            'mvp': torch.from_numpy( mvp ).float(),
            'lightpos': torch.from_numpy( lightpos ).float(),
            'campos': torch.from_numpy( campos ).float(),
            'bkgs': bkgs,
            'azim': torch.tensor(azim).float(),
            'elev': torch.tensor(elev).float(),
        }

class ListCameraBatch(torch.utils.data.Dataset):
    def __init__(self, datasets, bs, weights=None):
        self.datasets = datasets
        self.batch_size = bs
        self.weights = weights
    def __len__(self):
        return self.batch_size
    def __getitem__(self, index):
        d = random.choices(self.datasets, weights=self.weights)[0]
        return d[index]