import os, sys import imageio import random import time import torch import torch.nn as nn import torch.nn.functional as F import numpy as np def run_network(inputs, viewdirs, fn, embed_fn, embeddirs_fn,label, netchunk=1024*64): """Prepares inputs and applies network 'fn'. """ inputs_flat = torch.reshape(inputs, [inputs.shape[0],-1, inputs.shape[-1]]) embedded = embed_fn(inputs_flat) if viewdirs is not None: input_dirs = viewdirs[:,:,None].expand(inputs.shape) input_dirs_flat = torch.reshape(input_dirs, [inputs.shape[0],-1, input_dirs.shape[-1]]) embedded_dirs = embeddirs_fn(input_dirs_flat) #embedded = torch.cat([embedded, embedded_dirs], -1) input_all=torch.cat([inputs_flat,embedded_dirs],-1) outputs_flat = fn(input_all,label) outputs = torch.reshape(outputs_flat, list(inputs.shape[:-1]) + [outputs_flat.shape[-1]]) return outputs import torch # torch.autograd.set_detect_anomaly(True) import torch.nn as nn import torch.nn.functional as F import numpy as np # Misc img2mse = lambda x, y : torch.mean((x - y) ** 2) mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.]).to(x.device)) to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8) # 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, i=0): if i == -1: return nn.Identity(), 3 embed_kwargs = { 'include_input' : True, 'input_dims' : 3, '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 class Triplane(nn.Module): def __init__( self, ): super().__init__() self.plane_axis=self.generate_planes() def generate_planes(self): """ Defines planes by the three vectors that form the "axes" of the plane. Should work with arbitrary number of planes and planes of arbitrary orientation. """ return torch.tensor([[[1, 0, 0], [0, 1, 0], [0, 0, 1]], [[1, 0, 0], [0, 0, 1], [0, 1, 0]], [[0, 0, 1], [0, 1, 0], [1, 0, 0]]], dtype=torch.float32) def project_onto_planes(self,planes, coordinates): """ Does a projection of a 3D point onto a batch of 2D planes, returning 2D plane coordinates. Takes plane axes of shape n_planes, 3, 3 # Takes coordinates of shape N, M, 3 # returns projections of shape N*n_planes, M, 2 """ N, M, C = coordinates.shape n_planes, _, _ = planes.shape coordinates = coordinates.unsqueeze(1).expand(-1, n_planes, -1, -1).reshape(N*n_planes, M, 3) inv_planes = torch.linalg.inv(planes).unsqueeze(0).expand(N, -1, -1, -1).reshape(N*n_planes, 3, 3).to(device=coordinates.device) projections = torch.bmm(coordinates, inv_planes) return projections[..., :2] def sample_from_planes(self,plane_axes, plane_features, coordinates, mode='bilinear', padding_mode='zeros', box_warp=None): assert padding_mode == 'zeros' N, n_planes, C, H, W = plane_features.shape _, M, _ = coordinates.shape plane_features = plane_features.view(N*n_planes, C, H, W) coordinates = (2/box_warp) * coordinates # TODO: add specific box bounds #ipdb.set_trace() coordinates = self.project_onto_planes(plane_axes, coordinates).unsqueeze(1) output_features = torch.nn.functional.grid_sample(plane_features, coordinates.float(), mode=mode, padding_mode=padding_mode, align_corners=False).permute(0, 3, 2, 1).reshape(N, n_planes, M, C) # xy,xz,zy return output_features def forward(self, planes, sample_coordinates,box=1): #ipdb.set_trace() return self.sample_from_planes(self.plane_axis, planes, sample_coordinates, padding_mode='zeros', box_warp=box) def positional_encoding(positions, freqs): freq_bands = (2**torch.arange(freqs).float()).to(positions.device) # (F,) pts = (positions[..., None] * freq_bands).reshape( positions.shape[:-1] + (freqs * positions.shape[-1], )) # (..., DF) pts = torch.cat([torch.sin(pts), torch.cos(pts)], dim=-1) return pts def exists(val): return val is not None def resize_image_to( image, target_image_size, clamp_range = None, mode = 'nearest' ): orig_image_size = image.shape[-1] if orig_image_size == target_image_size: return image out = F.interpolate(image, target_image_size, mode = mode) if exists(clamp_range): out = out.clamp(*clamp_range) return out class NeRF(nn.Module): def __init__(self, D=8, W=256, input_ch=3, size=256,input_ch_views=3, output_ch=4, skips=[4], use_viewdirs=False,num_instance=1): """ """ super(NeRF, self).__init__() self.D = D self.W = W self.input_ch = input_ch//3 self.input_ch_views = input_ch_views self.skips = skips self.use_viewdirs = use_viewdirs self.hidden_dim=W self.triplane=Triplane() #ipdb.set_trace() self.tri_planes = nn.Parameter(torch.randn(num_instance, input_ch, size, size)) nn.init.normal_(self.tri_planes, mean=0, std=0.1) #self.weight=nn.Parameter(torch.ones(1,3,1,input_ch)) #ipdb.set_trace() self.pts_linears = nn.ModuleList( [nn.Linear(self.input_ch, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W) for i in range(D-1)]) ### 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)]) self.softplus=nn.Softplus() self.feature_linear = nn.Linear(W, W) self.alpha_linear = nn.Linear(W, 1) self.rgb_linear = nn.Linear(W, 3) # for m in self.children(): # if isinstance(m, nn.Linear): # nn.init.normal_(m.weight, std=0.01) def forward(self, x, label): #ipdb.set_trace() input_pts, input_views = torch.split(x, [int(x.shape[-1]-self.input_ch_views), self.input_ch_views], dim=-1) B,N,M=input_views.shape #ipdb.set_trace() # eal=resize_image_to(self.tri_planes[label],256) # eal=resize_image_to(eal,256) norm=torch.abs(self.tri_planes[label]).max(2)[0].max(2)[0].unsqueeze(-1).unsqueeze(-1) sample_triplane=(self.tri_planes[label]/norm).view(1,3,self.tri_planes.shape[-3]//3,self.tri_planes.shape[-2],self.tri_planes.shape[-1]).repeat(B,1,1,1,1) #ipdb.set_trace() input_pts=(self.triplane(sample_triplane,input_pts,2.5)).mean(1).view(-1,self.tri_planes.shape[-3]//3) #ipdb.set_trace() h = input_pts for i, l in enumerate(self.pts_linears): #ipdb.set_trace() h = self.pts_linears[i](h) h = F.relu(h) if i in self.skips: # ipdb.set_trace() h = torch.cat([input_pts, h], -1) alpha = self.alpha_linear(h) feature = self.feature_linear(h) h = torch.cat([feature, input_views.view(B*N,M)], -1) for i, l in enumerate(self.views_linears): h = self.views_linears[i](h) h = F.relu(h) rgb = self.rgb_linear(h) outputs = torch.cat([rgb.view(B,N,3), alpha.view(B,N,1)], -1) #ipdb.set_trace() return outputs class NeRF11(nn.Module): def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, output_ch=4, skips=[4], use_viewdirs=False,num_instance=1): """ """ super(NeRF11, self).__init__() self.D = D self.W = W self.input_ch = input_ch//3 self.input_ch_views = input_ch_views self.skips = skips self.use_viewdirs = use_viewdirs self.hidden_dim=W self.triplane=Triplane() self.weight = nn.Parameter(torch.zeros(1, 1, 256)) self.tri_planes = nn.Parameter(torch.randn(num_instance, input_ch, 256, 256)) #nn.init.normal_(self.tri_planes, mean=0, std=0.1) #self.weight=nn.Parameter(torch.ones(1,3,1,input_ch)) #ipdb.set_trace() self.label_emb = nn.Embedding(num_instance, W) self.pts_linears = nn.ModuleList( [nn.Linear(self.input_ch, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W) for i in range(D-1)]) ### 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)]) self.softplus=nn.Softplus() #self.label_feature = nn.Linear(W, W) self.feature_linear = nn.Linear(W, W) self.alpha_linear = nn.Linear(W, 1) self.rgb_linear = nn.Linear(W, 3) def forward(self, x,label): #ipdb.set_trace() input_pts, input_views = torch.split(x, [int(x.shape[-1]-self.input_ch_views), self.input_ch_views], dim=-1) B,N,M=input_views.shape sample_triplane=self.tri_planes[label].view(B,3,self.tri_planes.shape[-3]//3,self.tri_planes.shape[-2],self.tri_planes.shape[-1]) input_pts=(self.triplane(sample_triplane,input_pts,4)).mean(1).view(B,-1,self.tri_planes.shape[-3]//3) #ipdb.set_trace() label_emb=(self.weight*self.label_emb(label).unsqueeze(1)).expand(-1,N,-1) h = input_pts for i, l in enumerate(self.pts_linears): #ipdb.set_trace() h = self.pts_linears[i](h) h=h+label_emb h = F.relu(h) if i in self.skips: # ipdb.set_trace() h = torch.cat([input_pts, h], -1) alpha = self.alpha_linear(h) feature = self.feature_linear(h) h = torch.cat([feature, input_views.view(B,N,M)], -1) for i, l in enumerate(self.views_linears): h = self.views_linears[i](h) h = F.relu(h) rgb = self.rgb_linear(h) outputs = torch.cat([rgb.view(B,N,3), alpha.view(B,N,1)], -1) return outputs class NeRF0(nn.Module): def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, output_ch=4, skips=[4], use_viewdirs=False,num_instance=1): """ """ super(NeRF0, self).__init__() self.D = D self.W = W self.input_ch = input_ch//3 self.input_ch_views = input_ch_views self.skips = skips self.use_viewdirs = use_viewdirs self.hidden_dim=W self.triplane=Triplane() self.tri_planes = nn.Parameter(torch.randn(num_instance, input_ch, 256, 256)) #nn.init.normal_(self.tri_planes, mean=0, std=0.1) #self.weight=nn.Parameter(torch.ones(1,3,1,input_ch)) #ipdb.set_trace() self.pts_linears = nn.ModuleList( [nn.Linear(self.input_ch, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W) for i in range(D-1)]) ### 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)]) self.softplus=nn.Softplus() self.feature_linear = nn.Linear(W, W) self.alpha_linear = nn.Linear(W, 1) self.rgb_linear = nn.Linear(W, 3) def forward(self, x,label): #ipdb.set_trace() input_pts, input_views = torch.split(x, [int(x.shape[-1]-self.input_ch_views), self.input_ch_views], dim=-1) B,N,M=input_views.shape sample_triplane=self.tri_planes[label].view(B,3,self.tri_planes.shape[-3]//3,self.tri_planes.shape[-2],self.tri_planes.shape[-1]) #ipdb.set_trace() input_pts=(self.triplane(sample_triplane,input_pts,8)).mean(1).view(-1,self.tri_planes.shape[-3]//3) h = input_pts for i, l in enumerate(self.pts_linears): #ipdb.set_trace() h = self.pts_linears[i](h) h = F.relu(h) if i in self.skips: # ipdb.set_trace() h = torch.cat([input_pts, h], -1) alpha = self.alpha_linear(h) feature = self.feature_linear(h) h = torch.cat([feature, input_views.view(B*N,M)], -1) for i, l in enumerate(self.views_linears): h = self.views_linears[i](h) h = F.relu(h) rgb = self.rgb_linear(h) outputs = torch.cat([rgb.view(B,N,3), alpha.view(B,N,1)], -1) return outputs class NeRF1(nn.Module): def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, output_ch=4, skips=[4], use_viewdirs=False,num_instance=1): """ """ super(NeRF1, self).__init__() self.D = D self.W = W self.input_ch = input_ch//3*9 self.input_ch_views = input_ch_views self.skips = skips self.use_viewdirs = use_viewdirs self.hidden_dim=W self.triplane=Triplane() self.tri_planes = nn.Parameter(torch.randn(num_instance, input_ch, 256, 256)) #nn.init.normal_(self.tri_planes, mean=0, std=0.1) #self.weight=nn.Parameter(torch.ones(1,3,1,input_ch)) #ipdb.set_trace() self.pts_linears = nn.ModuleList( [nn.Linear(self.input_ch, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W) for i in range(D-1)]) ### 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)]) self.softplus=nn.Softplus() self.feature_linear = nn.Linear(W, W) self.alpha_linear = nn.Linear(W, 1) self.rgb_linear = nn.Linear(W//2, 3) def forward(self, x,label): #ipdb.set_trace() input_pts, input_views = torch.split(x, [int(x.shape[-1]-self.input_ch_views), self.input_ch_views], dim=-1) B,N,M=input_views.shape sample_triplane=self.tri_planes[label].view(B,3,self.tri_planes.shape[-3]//3,self.tri_planes.shape[-2],self.tri_planes.shape[-1]) #ipdb.set_trace() input_pts=(self.triplane(sample_triplane,input_pts,4)).mean(1).view(-1,self.tri_planes.shape[-3]//3) h = torch.cat((input_pts,positional_encoding(input_pts,4)),-1) for i, l in enumerate(self.pts_linears): #ipdb.set_trace() h = self.pts_linears[i](h) h = F.relu(h) if i in self.skips: # ipdb.set_trace() h = torch.cat([input_pts, h], -1) alpha = self.alpha_linear(h) feature = self.feature_linear(h) h = torch.cat([feature, input_views.view(B*N,M)], -1) for i, l in enumerate(self.views_linears): h = self.views_linears[i](h) h = F.relu(h) rgb = self.rgb_linear(h) outputs = torch.cat([rgb.view(B,N,3), alpha.view(B,N,1)], -1) return outputs class NeRF_dual(nn.Module): def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, output_ch=4, skips=[4], use_viewdirs=False,num_instance=1): """ """ super(NeRF_dual, self).__init__() self.D = D self.W = W self.input_ch = input_ch//3*9+input_ch_views self.input_ch2= input_ch//3*5 self.input_ch_views = input_ch_views self.skips = skips self.use_viewdirs = use_viewdirs self.hidden_dim=W self.triplane=Triplane() self.tri_planes1 = nn.Parameter(torch.randn(num_instance, input_ch, 256, 256)) self.tri_planes2 = nn.Parameter(torch.randn(num_instance, input_ch, 256, 256)) #nn.init.normal_(self.tri_planes, mean=0, std=0.1) #self.weight=nn.Parameter(torch.ones(1,3,1,input_ch)) #ipdb.set_trace() self.pts_linears = nn.ModuleList( [nn.Linear(self.input_ch, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W) for i in range(D-1)]) self.pts_linears2 = nn.ModuleList( [nn.Linear(self.input_ch2, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W) for i in range(D-1)]) ### 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(W, W//2)]) self.softplus=nn.Softplus() self.feature_linear = nn.Linear(W, W) self.alpha_linear = nn.Linear(W, 1) self.rgb_linear = nn.Linear(W//2, 3) def forward(self, x,label): #ipdb.set_trace() input_pts, input_views = torch.split(x, [int(x.shape[-1]-self.input_ch_views), self.input_ch_views], dim=-1) B,N,M=input_views.shape sample_triplane1=self.tri_planes1[label].view(B,3,self.tri_planes1.shape[-3]//3,self.tri_planes1.shape[-2],self.tri_planes1.shape[-1]) #ipdb.set_trace() input_pts1=(self.triplane(sample_triplane1,input_pts,8)).mean(1).view(B,-1,self.tri_planes1.shape[-3]//3) sample_triplane2=self.tri_planes2[label].view(B,3,self.tri_planes2.shape[-3]//3,self.tri_planes2.shape[-2],self.tri_planes2.shape[-1]) #ipdb.set_trace() input_pts2=(self.triplane(sample_triplane2,input_pts,8)).mean(1).view(B,-1,self.tri_planes2.shape[-3]//3) #ipdb.set_trace() h = torch.cat((input_pts1,positional_encoding(input_pts1,4),input_views),-1) for i, l in enumerate(self.pts_linears): #ipdb.set_trace() h = self.pts_linears[i](h) h = F.relu(h) if i in self.skips: # ipdb.set_trace() h = torch.cat([input_pts, h], -1) h = self.feature_linear(h) for i, l in enumerate(self.views_linears): h = self.views_linears[i](h) h = F.relu(h) rgb = self.rgb_linear(h) h = torch.cat((input_pts2,positional_encoding(input_pts2,2)),-1) for i, l in enumerate(self.pts_linears2): #ipdb.set_trace() h = self.pts_linears2[i](h) h = F.relu(h) alpha = self.alpha_linear(h) outputs = torch.cat([rgb.view(B,N,3), alpha.view(B,N,1)], -1) return outputs # Ray helpers def get_rays(H, W, K, c2w): i, j = torch.meshgrid(torch.linspace(0, W-1, W), torch.linspace(0, H-1, H)) # pytorch's meshgrid has indexing='ij' i = i.t() j = j.t() dirs = torch.stack([(i-K[0][2])/K[0][0], -(j-K[1][2])/K[1][1], -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_rays_np(H, W, K, c2w): i, j = np.meshgrid(np.arange(W, dtype=np.float32), np.arange(H, dtype=np.float32), indexing='xy') dirs = np.stack([(i-K[0][2])/K[0][0], -(j-K[1][2])/K[1][1], -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 = torch.searchsorted(cdf, u, right=True) 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 def render_path1(batch_rays, chunk, render_kwargs, gt_imgs=None, savedir=None,savedir1=None,near=None,far=None,label=None): rgbs = [] disps = [] t = time.time() #render(chunk=newargs.chunk, rays=batch_rays,near=near,far=far,label=label, retraw=True, **render_kwargs_train) rgbs, disps, acc, _ = render(chunk=chunk, rays=batch_rays,near=near,far=far, label=label,**render_kwargs) #ipdb.set_trace() reso=int(rgbs.shape[-2]**0.5) rgbs=rgbs.view(-1,reso,reso,3) disps=disps.view(-1,reso,reso,1) acc=acc.view(-1,reso,reso,1) #ipdb.set_trace() gt_imgs=gt_imgs.view(-1,reso,reso,3) mask=(gt_imgs.mean(-1)<0.9999) #ipdb.set_trace() if savedir is not None: for i in range(len(rgbs)): rgb8 =to8b(rgbs[i].cpu().numpy())#np.fliplr(np.rot90(to8b(rgbs[i]),-1)) filename = os.path.join(savedir) imageio.imwrite(savedir, rgb8) imageio.imwrite(savedir1, np.uint8(gt_imgs[i].cpu().numpy()*255)) #ipdb.set_trace() print('psnr:' ,mse2psnr(img2mse(torch.Tensor(rgb8/255).to(device=gt_imgs.device)[mask[i]],(gt_imgs[i])[mask[i]]))) print('psnr_all:' ,mse2psnr(img2mse(torch.Tensor(rgb8/255).to(device=gt_imgs.device),(gt_imgs[i])))) psnr_list = [] for i in range(len(rgbs)): rgb8 = to8b(rgbs[i].cpu().numpy()) psnr = mse2psnr(img2mse(torch.Tensor(rgb8/255).to(device=gt_imgs.device),(gt_imgs[i]))) psnr_list.append(psnr) #ipdb.set_trace() return rgbs, disps, acc, psnr_list def render(chunk=1024*32, rays=None, c2w=None, ndc=True,label=None, near=0., far=1., use_viewdirs=False, c2w_staticcam=None, **kwargs): """Render rays Args: H: int. Height of image in pixels. W: int. Width of image in pixels. focal: float. Focal length of pinhole camera. chunk: int. Maximum number of rays to process simultaneously. Used to control maximum memory usage. Does not affect final results. rays: array of shape [2, batch_size, 3]. Ray origin and direction for each example in batch. c2w: array of shape [3, 4]. Camera-to-world transformation matrix. ndc: bool. If True, represent ray origin, direction in NDC coordinates. near: float or array of shape [batch_size]. Nearest distance for a ray. far: float or array of shape [batch_size]. Farthest distance for a ray. use_viewdirs: bool. If True, use viewing direction of a point in space in model. c2w_staticcam: array of shape [3, 4]. If not None, use this transformation matrix for camera while using other c2w argument for viewing directions. Returns: rgb_map: [batch_size, 3]. Predicted RGB values for rays. disp_map: [batch_size]. Disparity map. Inverse of depth. acc_map: [batch_size]. Accumulated opacity (alpha) along a ray. extras: dict with everything returned by render_rays(). """ #ipdb.set_trace() rays_o, rays_d = rays[:,0,...], rays[:,1,...] viewdirs = rays_d viewdirs = viewdirs / torch.norm(viewdirs, dim=-1, keepdim=True) viewdirs = torch.reshape(viewdirs, [rays_d.shape[0],-1,3]).float() sh = rays_d.shape # [..., 3] # Create ray batch rays_o = torch.reshape(rays_o, [sh[0],-1,3]).float() rays_d = torch.reshape(rays_d, [sh[0],-1,3]).float() #ipdb.set_trace() near, far = near[:,None,:] * torch.ones_like(rays_d[...,:1]), far[:,None,:] * torch.ones_like(rays_d[...,:1]) rays = torch.cat([rays_o, rays_d, near, far], -1) if use_viewdirs: rays = torch.cat([rays, viewdirs], -1) #ipdb.set_trace() # Render and reshape all_ret = batchify_rays(rays, label,chunk,**kwargs) for k in all_ret: k_sh = list(sh[:-1]) + list(all_ret[k].shape[2:]) all_ret[k] = torch.reshape(all_ret[k], k_sh) k_extract = ['rgb_map', 'disp_map', 'acc_map'] ret_list = [all_ret[k] for k in k_extract] ret_dict = {k : all_ret[k] for k in all_ret if k not in k_extract} return ret_list + [ret_dict] def batchify_rays(rays_flat,label, chunk=1024*32, **kwargs): """Render rays in smaller minibatches to avoid OOM. """ all_ret = {} for i in range(0, rays_flat.shape[1], chunk): #ipdb.set_trace() ret = render_rays(rays_flat[:,i:i+chunk],label=label, **kwargs) for k in ret: if k not in all_ret: all_ret[k] = [] all_ret[k].append(ret[k]) #ipdb.set_trace() all_ret = {k : torch.cat(all_ret[k], 1) for k in all_ret} return all_ret def render_rays(ray_batch, network_fn, network_query_fn, N_samples, retraw=False, lindisp=False, perturb=0., N_importance=0, network_fine=None, white_bkgd=False, raw_noise_std=0., label=None, verbose=False, pytest=False): """Volumetric rendering. Args: ray_batch: array of shape [batch_size, ...]. All information necessary for sampling along a ray, including: ray origin, ray direction, min dist, max dist, and unit-magnitude viewing direction. network_fn: function. Model for predicting RGB and density at each point in space. network_query_fn: function used for passing queries to network_fn. N_samples: int. Number of different times to sample along each ray. retraw: bool. If True, include model's raw, unprocessed predictions. lindisp: bool. If True, sample linearly in inverse depth rather than in depth. perturb: float, 0 or 1. If non-zero, each ray is sampled at stratified random points in time. N_importance: int. Number of additional times to sample along each ray. These samples are only passed to network_fine. network_fine: "fine" network with same spec as network_fn. white_bkgd: bool. If True, assume a white background. raw_noise_std: ... verbose: bool. If True, print more debugging info. Returns: rgb_map: [num_rays, 3]. Estimated RGB color of a ray. Comes from fine model. disp_map: [num_rays]. Disparity map. 1 / depth. acc_map: [num_rays]. Accumulated opacity along each ray. Comes from fine model. raw: [num_rays, num_samples, 4]. Raw predictions from model. rgb0: See rgb_map. Output for coarse model. disp0: See disp_map. Output for coarse model. acc0: See acc_map. Output for coarse model. z_std: [num_rays]. Standard deviation of distances along ray for each sample. """ B,N_rays,_ = ray_batch.shape rays_o, rays_d = ray_batch[:,:,0:3], ray_batch[:,:,3:6] # [N_rays, 3] each viewdirs = ray_batch[:,:,-3:] bounds = torch.reshape(ray_batch[...,6:8], [B,-1,1,2]) near, far = bounds[...,0], bounds[...,1] # [-1,1] t_vals = torch.linspace(0., 1., steps=N_samples).to(near.device) z_vals = near * (1.-t_vals) + far * (t_vals) #z_vals = z_vals.expand([N_rays, N_samples]) if perturb > 0.: # get intervals between samples mids = .5 * (z_vals[...,1:] + z_vals[...,:-1]) upper = torch.cat([mids, z_vals[...,-1:]], -1) lower = torch.cat([z_vals[...,:1], mids], -1) # stratified samples in those intervals t_rand = torch.rand(z_vals.shape) # Pytest, overwrite u with numpy's fixed random numbers if pytest: np.random.seed(0) t_rand = np.random.rand(*list(z_vals.shape)) t_rand = torch.Tensor(t_rand) z_vals = lower + (upper - lower) * t_rand pts = rays_o[...,None,:] + rays_d[...,None,:] * z_vals[...,:,None] # [N_rays, N_samples, 3] # raw = run_network(pts) raw = network_query_fn(pts, viewdirs, label,network_fn) rgb_map, disp_map, acc_map, weights, depth_map = raw2outputs(raw, z_vals, rays_d, raw_noise_std, white_bkgd, pytest=pytest) ret = {'rgb_map' : rgb_map, 'disp_map' : disp_map, 'acc_map' : acc_map} if retraw: ret['raw'] = raw if N_importance > 0: ret['rgb0'] = rgb_map_0 ret['disp0'] = disp_map_0 ret['acc0'] = acc_map_0 ret['z_std'] = torch.std(z_samples, dim=-1, unbiased=False) # [N_rays] for k in ret: if (torch.isnan(ret[k]).any() or torch.isinf(ret[k]).any()): print(f"! [Numerical Error] {k} contains nan or inf.") return ret def raw2outputs(raw, z_vals, rays_d, raw_noise_std=0, white_bkgd=False, pytest=False): """Transforms model's predictions to semantically meaningful values. Args: raw: [num_rays, num_samples along ray, 4]. Prediction from model. z_vals: [num_rays, num_samples along ray]. Integration time. rays_d: [num_rays, 3]. Direction of each ray. Returns: rgb_map: [num_rays, 3]. Estimated RGB color of a ray. disp_map: [num_rays]. Disparity map. Inverse of depth map. acc_map: [num_rays]. Sum of weights along each ray. weights: [num_rays, num_samples]. Weights assigned to each sampled color. depth_map: [num_rays]. Estimated distance to object. """ #ipdb.set_trace() act_ff=nn.Softplus() raw2alpha = lambda raw, dists, act_fn=act_ff: 1.-torch.exp(-act_fn(raw)*dists) dists = z_vals[...,1:] - z_vals[...,:-1] dists = torch.cat([dists, torch.Tensor([1e10]).to(dists.device).expand(dists[...,:1].shape)], -1) # [N_rays, N_samples] dists = dists * torch.norm(rays_d[...,None,:], dim=-1) rgb = torch.sigmoid(raw[...,:3]) # [N_rays, N_samples, 3] noise = 0. if raw_noise_std > 0.: noise = torch.randn(raw[...,3].shape) * raw_noise_std # Overwrite randomly sampled data if pytest if pytest: np.random.seed(0) noise = np.random.rand(*list(raw[...,3].shape)) * raw_noise_std noise = torch.Tensor(noise) #ipdb.set_trace() alpha = raw2alpha(raw[...,3] + noise, dists) # [N_rays, N_samples] #ipdb.set_trace() weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0],alpha.shape[1], 1)).to(alpha.device), 1.-alpha + 1e-10], -1), -1)[:,:, :-1] rgb_map = torch.sum(weights[...,None] * rgb, -2) # [N_rays, 3] depth_map = torch.sum(weights * z_vals, -1) disp_map = 1./torch.max(1e-10 * torch.ones_like(depth_map), depth_map / torch.sum(weights, -1)) acc_map = torch.sum(weights, -1) if white_bkgd: rgb_map = rgb_map + (1.-acc_map[...,None]) return rgb_map, disp_map, acc_map, weights, depth_map def get_rays(H, W, K, c2w): i, j = torch.meshgrid(torch.linspace(0, W-1, W), torch.linspace(0, H-1, H)) # pytorch's meshgrid has indexing='ij' i = i.t() j = j.t() dirs = torch.stack([(i-K[0][2])/K[0][0], (j-K[1][2])/K[1][1], 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