run_evdnerf_helpers.py

import torch
torch.autograd.set_detect_anomaly(True)
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch import searchsorted
import glob, os
import cv2

SMALL_EPS = 1e-10
EPS = 1e-5

# Misc
img2mse = lambda x, y : torch.mean((x - y) ** 2)
img2mse_masked = lambda x, y, mask : torch.mean((torch.masked_select(x, torch.unsqueeze(mask, 1)) - torch.masked_select(y, torch.unsqueeze(mask, 1))) ** 2)
mse2psnr = lambda x, range=1.0 : -10. * torch.log10(x / torch.Tensor([range])**2)
to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)

# RGB -> B intensity -> L Gamma-corrected log image conversion
img2I = lambda Im : torch.sum(Im, 1)/3.
img2L = lambda Im, g=1.0 : 1/g * torch.log( Im )

compute_pred_ev = lambda im1, im0, g=1.0 : 1/g * ( torch.log(im1+SMALL_EPS) - torch.log(im0+SMALL_EPS) ).squeeze()
evim_to8b = lambda evim, max_val : (np.abs(evim/max_val) * 255).astype(np.uint8)

# compute a piecewise loss function that:
# - is 0 for the range in which the correct number of events would be calculated from the predicted difference in log intensity (for a single positive event, the range [pos_thresh, pos_thresh+pos_thresh))
# - is squared error outside the range, (relative to center of the range?)
# inputs:
# - pred: predicted difference in log intensity (continuous signal)
# - gt: ground truth eventstream, binned by the threshold values
def ev_threshold_loss(pred, gt, pos_thresh=0.2, neg_thresh=0.2):
    neg_thresh *= -1 # make neg_thresh negative to aid in arithmetic
    # print(pos_thresh, neg_thresh)
    loss = torch.zeros_like(pred)
    pos_inwindow = torch.logical_and(gt >= 0.0, torch.logical_and(pred >= gt, pred < gt + pos_thresh))
    # print(pos_inwindow)
    neg_inwindow = torch.logical_and(gt < 0.0, torch.logical_and(pred <= gt, pred > gt + neg_thresh))
    # print(neg_inwindow)
    loss[torch.logical_and(pos_inwindow, neg_inwindow)] = 0.0
    # make center of in_window the target
    target = gt.clone()
    target[gt > 0] += pos_thresh/2
    target[gt < 0] += neg_thresh/2
    diff = pred - target
    loss[torch.logical_and(~pos_inwindow, ~neg_inwindow)] = (diff[torch.logical_and(~pos_inwindow, ~neg_inwindow)])**2
    return loss.mean()

# dist loss from mip nerf 360
def compute_dist_loss(weights, zvals):
    mid_zs = 0.5 * (zvals[...,1:] + zvals[...,:-1])
    dist_loss = ((weights[...,:-1].unsqueeze(1) * weights[...,:-1].unsqueeze(2)) * torch.abs(mid_zs.unsqueeze(1) - mid_zs.unsqueeze(2))).sum((1,2)) + 1/3 * ((weights[...,:-1]**2) * (zvals[...,1:] - zvals[...,:-1])).sum(1)
    dist_loss = dist_loss.mean()    
    return dist_loss

# affine transformation to adjust predicted intensity images to match ground truth white balance
# TODO this should be done in log-space but the transformation isn't working
def correct_prediction_image(pred, target_rgb, path_to_gt_ims):
    if target_rgb is None or target_rgb.max() < .01:
        # use image(s) in path_to_gt_ims
        # load images in path_to_gt_ims as grayscale
        gt = np.stack([cv2.resize(cv2.imread(imfile).mean(2)/255.0, (pred.shape[1], pred.shape[0]), interpolation=cv2.INTER_AREA) for imfile in glob.glob(os.path.join(path_to_gt_ims, '*.png'))])
    else:
        gt = target_rgb

    X = torch.log(gt.reshape(-1, 1) + 0.1)
    Y = torch.log(pred.reshape(-1, 1) + 0.1)
    # calulate affine transformation in log-space
    A = (torch.mean(X*Y, 0)-torch.mean(X, 0)*torch.mean(Y, 0))/(torch.mean(Y*Y, 0)-torch.mean(Y, 0)**2)
    B = torch.mean(X, 0)-A*torch.mean(Y, 0)
    Y = torch.exp( Y * A + B ) - 0.1

    return (torch.maximum(Y, torch.zeros_like(Y))).reshape(pred.shape)

# Positional encoding (section 5.1)
class Embedder:
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self.create_embedding_fn()
        
    def create_embedding_fn(self):
        embed_fns = []
        d = self.kwargs['input_dims']
        out_dim = 0
        if self.kwargs['include_input']:
            embed_fns.append(lambda x : x)
            out_dim += d
            
        max_freq = self.kwargs['max_freq_log2']
        N_freqs = self.kwargs['num_freqs']
        
        if self.kwargs['log_sampling']:
            freq_bands = 2.**torch.linspace(0., max_freq, steps=N_freqs)
        else:
            freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=N_freqs)
            
        for freq in freq_bands:
            for p_fn in self.kwargs['periodic_fns']:
                embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq))
                out_dim += d
                    
        self.embed_fns = embed_fns
        self.out_dim = out_dim
        
    def embed(self, inputs):
        return torch.cat([fn(inputs) for fn in self.embed_fns], -1)


def get_embedder(multires, input_dims, i=0):
    if i == -1:
        return nn.Identity(), input_dims
    
    embed_kwargs = {
                'include_input' : True,
                'input_dims' : input_dims,
                'max_freq_log2' : multires-1,
                'num_freqs' : multires,
                'log_sampling' : True,
                'periodic_fns' : [torch.sin, torch.cos],
    }
    
    embedder_obj = Embedder(**embed_kwargs)
    embed = lambda x, eo=embedder_obj : eo.embed(x)
    return embed, embedder_obj.out_dim


# Model
class DirectTemporalNeRF(nn.Module):
    def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, input_ch_time=1, output_ch=4, skips=[4],
                 use_viewdirs=False, memory=[], embed_fn=None, zero_canonical=True, output_color_ch=3, init_method=None):
        super(DirectTemporalNeRF, self).__init__()
        self.D = D
        self.W = W
        self.input_ch = input_ch
        self.input_ch_views = input_ch_views
        self.input_ch_time = input_ch_time
        self.skips = skips
        self.use_viewdirs = use_viewdirs
        self.memory = memory
        self.embed_fn = embed_fn
        self.zero_canonical = zero_canonical
        self.init_method = init_method

        self._occ = NeRFOriginal(D=D, W=W, input_ch=input_ch, input_ch_views=input_ch_views,
                                 input_ch_time=input_ch_time, output_ch=output_ch, skips=skips,
                                 use_viewdirs=use_viewdirs, memory=memory, embed_fn=embed_fn, output_color_ch=output_color_ch, init_method=init_method)
        self._time, self._time_out = self.create_time_net()

    def create_time_net(self):
        layers = [nn.Linear(self.input_ch + self.input_ch_time, self.W)]
        for i in range(self.D - 1):
            if i in self.memory:
                raise NotImplementedError
            else:
                layer = nn.Linear

            in_channels = self.W
            if i in self.skips:
                in_channels += self.input_ch

            layers += [layer(in_channels, self.W)]
            
        densenetwork = nn.ModuleList(layers)
        final = nn.Linear(self.W, 3)
        if self.init_method is not None:
            for layer in densenetwork:
                self.init_method(layer.weight)
                self.init_method(final.weight)

        return densenetwork, final

    def query_time(self, new_pts, t, net, net_final):
        h = torch.cat([new_pts, t], dim=-1)
        for i, l in enumerate(net):
            h = net[i](h)
            h = torch.relu(h)
            if i in self.skips:
                h = torch.cat([new_pts, h], -1)

        return net_final(h)

    def forward(self, x, ts):
        input_pts, input_views = torch.split(x, [self.input_ch, self.input_ch_views], dim=-1)
        t = ts[0]

        # assert len(torch.unique(t[:, :1])) == 1, "Only accepts all points from same time"
        cur_time = t[0, 0]
        # if cur_time == 0. and self.zero_canonical:
        #     # t=0 is by definition without deformation
        #     dx = torch.zeros_like(input_pts[:, :3])
        # else:

        nz_t = (t[:,0]!=0.)

        # run deformation network
        dx = self.query_time(input_pts[nz_t], t[nz_t], self._time, self._time_out)
        dx_full = torch.zeros_like(input_pts[:, :3], device=input_pts.device)
        dx_full[nz_t] = dx
        input_pts_orig = input_pts[:, :3]
        input_pts = self.embed_fn(input_pts_orig + dx_full)

        # run canonical network
        out, _ = self._occ(torch.cat([input_pts, input_views], dim=-1), t)
        return out, dx_full


class NeRF:
    @staticmethod
    def get_by_name(type,  *args, **kwargs):
        print ("NeRF type selected: %s" % type)

        if type == "original":
            model = NeRFOriginal(*args, **kwargs)
        elif type == "direct_temporal":
            model = DirectTemporalNeRF(*args, **kwargs)
        else:
            raise ValueError("Type %s not recognized." % type)
        return model

class NeRFOriginal(nn.Module):
    def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, input_ch_time=1, output_ch=4, skips=[4],
                 use_viewdirs=False, memory=[], embed_fn=None, output_color_ch=3, zero_canonical=True, init_method=None):
        super(NeRFOriginal, self).__init__()
        self.D = D
        self.W = W
        self.input_ch = input_ch
        self.input_ch_views = input_ch_views
        self.skips = skips
        self.use_viewdirs = use_viewdirs

        # self.pts_linears = nn.ModuleList(
        #     [nn.Linear(input_ch, W)] +
        #     [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + input_ch, W) for i in range(D-1)])

        layers = [nn.Linear(input_ch, W)]
        for i in range(D - 1):
            if i in memory:
                raise NotImplementedError
            else:
                layer = nn.Linear

            in_channels = W
            if i in self.skips:
                in_channels += input_ch

            layers += [layer(in_channels, W)]

        self.pts_linears = nn.ModuleList(layers)

        ### Implementation according to the official code release (https://github.com/bmild/nerf/blob/master/run_nerf_helpers.py#L104-L105)
        self.views_linears = nn.ModuleList([nn.Linear(input_ch_views + W, W//2)])

        ### Implementation according to the paper
        # self.views_linears = nn.ModuleList(
        #     [nn.Linear(input_ch_views + W, W//2)] + [nn.Linear(W//2, W//2) for i in range(D//2)])

        if use_viewdirs:
            self.feature_linear = nn.Linear(W, W)
            self.alpha_linear = nn.Linear(W, 1)
            self.rgb_linear = nn.Linear(W//2, output_color_ch)
        else:
            self.output_linear = nn.Linear(W, output_ch)

        if init_method is not None:
            for layer in self.pts_linears:
                init_method(layer.weight)
            for layer in self.views_linears:
                init_method(layer.weight)
            init_method(self.feature_linear.weight)
            init_method(self.alpha_linear.weight)
            init_method(self.rgb_linear.weight)

    def forward(self, x, ts):
        input_pts, input_views = torch.split(x, [self.input_ch, self.input_ch_views], dim=-1)
        h = input_pts
        for i, l in enumerate(self.pts_linears):
            h = self.pts_linears[i](h)
            h = torch.relu(h)
            if i in self.skips:
                h = torch.cat([input_pts, h], -1)

        if self.use_viewdirs:
            alpha = self.alpha_linear(h)
            feature = self.feature_linear(h)
            h = torch.cat([feature, input_views], -1)

            for i, l in enumerate(self.views_linears):
                h = self.views_linears[i](h)
                h = torch.relu(h)

            rgb = self.rgb_linear(h)
            outputs = torch.cat([rgb, alpha], -1)
        else:
            outputs = self.output_linear(h)

        return outputs, torch.zeros_like(input_pts[:, :3])

    def load_weights_from_keras(self, weights):
        assert self.use_viewdirs, "Not implemented if use_viewdirs=False"

        # Load pts_linears
        for i in range(self.D):
            idx_pts_linears = 2 * i
            self.pts_linears[i].weight.data = torch.from_numpy(np.transpose(weights[idx_pts_linears]))
            self.pts_linears[i].bias.data = torch.from_numpy(np.transpose(weights[idx_pts_linears+1]))

        # Load feature_linear
        idx_feature_linear = 2 * self.D
        self.feature_linear.weight.data = torch.from_numpy(np.transpose(weights[idx_feature_linear]))
        self.feature_linear.bias.data = torch.from_numpy(np.transpose(weights[idx_feature_linear+1]))

        # Load views_linears
        idx_views_linears = 2 * self.D + 2
        self.views_linears[0].weight.data = torch.from_numpy(np.transpose(weights[idx_views_linears]))
        self.views_linears[0].bias.data = torch.from_numpy(np.transpose(weights[idx_views_linears+1]))

        # Load rgb_linear
        idx_rbg_linear = 2 * self.D + 4
        self.rgb_linear.weight.data = torch.from_numpy(np.transpose(weights[idx_rbg_linear]))
        self.rgb_linear.bias.data = torch.from_numpy(np.transpose(weights[idx_rbg_linear+1]))

        # Load alpha_linear
        idx_alpha_linear = 2 * self.D + 6
        self.alpha_linear.weight.data = torch.from_numpy(np.transpose(weights[idx_alpha_linear]))
        self.alpha_linear.bias.data = torch.from_numpy(np.transpose(weights[idx_alpha_linear+1]))


def hsv_to_rgb(h, s, v):
    '''
    h,s,v in range [0,1]
    '''
    hi = torch.floor(h * 6)
    f = h * 6. - hi
    p = v * (1. - s)
    q = v * (1. - f * s)
    t = v * (1. - (1. - f) * s)

    rgb = torch.cat([hi, hi, hi], -1) % 6
    rgb[rgb == 0] = torch.cat((v, t, p), -1)[rgb == 0]
    rgb[rgb == 1] = torch.cat((q, v, p), -1)[rgb == 1]
    rgb[rgb == 2] = torch.cat((p, v, t), -1)[rgb == 2]
    rgb[rgb == 3] = torch.cat((p, q, v), -1)[rgb == 3]
    rgb[rgb == 4] = torch.cat((t, p, v), -1)[rgb == 4]
    rgb[rgb == 5] = torch.cat((v, p, q), -1)[rgb == 5]
    return rgb


# Ray helpers
def get_rays(H, W, focal, c2w, ray_jitter=False):
    i, j = torch.meshgrid(torch.linspace(0, W-1, W)+0.5, torch.linspace(0, H-1, H)+0.5)  # pytorch's meshgrid has indexing='ij'
    i = i.t()
    j = j.t()
    if ray_jitter:
        i += torch.rand_like(i) - 0.5
        j += torch.rand_like(j) - 0.5
    dirs = torch.stack([(i-W*.5)/focal, -(j-H*.5)/focal, -torch.ones_like(i)], -1)
    # Rotate ray directions from camera frame to the world frame
    rays_d = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1)  # dot product, equals to: [c2w.dot(dir) for dir in dirs]
    # Translate camera frame's origin to the world frame. It is the origin of all rays.
    rays_o = c2w[:3,-1].expand(rays_d.shape)
    return rays_o, rays_d

def get_batched_rays(H, W, focal, c2w, x, y, ray_jitter=False):
    x = x.type(torch.float32) + 0.5
    y = y.type(torch.float32) + 0.5
    if ray_jitter:
        x += torch.rand_like(x) - 0.5
        y += torch.rand_like(y) - 0.5
    dirs = torch.stack([(x-W*.5)/focal, -(y-H*.5)/focal, -torch.ones_like(x)], -1)
    rays_d = c2w[..., :3, :3] @ dirs.reshape(-1, 3, 1)
    rays_o = c2w[..., :3, -1]
    return rays_o, rays_d.squeeze()

def get_rays_np(H, W, focal, c2w):
    i, j = np.meshgrid(np.arange(W, dtype=np.float32), np.arange(H, dtype=np.float32), indexing='xy')
    dirs = np.stack([(i-W*.5)/focal, -(j-H*.5)/focal, -np.ones_like(i)], -1)
    # Rotate ray directions from camera frame to the world frame
    rays_d = np.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1)  # dot product, equals to: [c2w.dot(dir) for dir in dirs]
    # Translate camera frame's origin to the world frame. It is the origin of all rays.
    rays_o = np.broadcast_to(c2w[:3,-1], np.shape(rays_d))
    return rays_o, rays_d

def ndc_rays(H, W, focal, near, rays_o, rays_d):
    # Shift ray origins to near plane
    t = -(near + rays_o[...,2]) / rays_d[...,2]
    rays_o = rays_o + t[...,None] * rays_d
    
    # Projection
    o0 = -1./(W/(2.*focal)) * rays_o[...,0] / rays_o[...,2]
    o1 = -1./(H/(2.*focal)) * rays_o[...,1] / rays_o[...,2]
    o2 = 1. + 2. * near / rays_o[...,2]

    d0 = -1./(W/(2.*focal)) * (rays_d[...,0]/rays_d[...,2] - rays_o[...,0]/rays_o[...,2])
    d1 = -1./(H/(2.*focal)) * (rays_d[...,1]/rays_d[...,2] - rays_o[...,1]/rays_o[...,2])
    d2 = -2. * near / rays_o[...,2]
    
    rays_o = torch.stack([o0,o1,o2], -1)
    rays_d = torch.stack([d0,d1,d2], -1)
    
    return rays_o, rays_d


# Hierarchical sampling (section 5.2)
def sample_pdf(bins, weights, N_samples, det=False, pytest=False):
    # Get pdf
    weights = weights + 1e-5 # prevent nans
    pdf = weights / torch.sum(weights, -1, keepdim=True)
    cdf = torch.cumsum(pdf, -1)
    cdf = torch.cat([torch.zeros_like(cdf[...,:1]), cdf], -1)  # (batch, len(bins))

    # Take uniform samples
    if det:
        u = torch.linspace(0., 1., steps=N_samples)
        u = u.expand(list(cdf.shape[:-1]) + [N_samples])
    else:
        u = torch.rand(list(cdf.shape[:-1]) + [N_samples])

    # Pytest, overwrite u with numpy's fixed random numbers
    if pytest:
        np.random.seed(0)
        new_shape = list(cdf.shape[:-1]) + [N_samples]
        if det:
            u = np.linspace(0., 1., N_samples)
            u = np.broadcast_to(u, new_shape)
        else:
            u = np.random.rand(*new_shape)
        u = torch.Tensor(u)

    # Invert CDF
    u = u.contiguous()
    inds = searchsorted(cdf, u, right=True)# side='right')
    below = torch.max(torch.zeros_like(inds-1), inds-1)
    above = torch.min((cdf.shape[-1]-1) * torch.ones_like(inds), inds)
    inds_g = torch.stack([below, above], -1)  # (batch, N_samples, 2)

    # cdf_g = tf.gather(cdf, inds_g, axis=-1, batch_dims=len(inds_g.shape)-2)
    # bins_g = tf.gather(bins, inds_g, axis=-1, batch_dims=len(inds_g.shape)-2)
    matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
    cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
    bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)

    denom = (cdf_g[...,1]-cdf_g[...,0])
    denom = torch.where(denom<1e-5, torch.ones_like(denom), denom)
    t = (u-cdf_g[...,0])/denom
    samples = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0])

    return samples