models/SAP/utils.py

# THE CODE WAS TAKEN AND ADAPTED FROM https://pengsongyou.github.io/sap
# @inproceedings{Peng2021SAP,
# author = {Peng, Songyou and Jiang, Chiyu "Max" and Liao, Yiyi and Niemeyer, Michael and Pollefeys, Marc and Geiger, Andreas},
# title = {Shape As Points: A Differentiable Poisson Solver},
# booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
# year = {2021}
# }

import torch
import io, os, logging, urllib
import yaml
import trimesh
import imageio
import numbers
import math
import numpy as np
from collections import OrderedDict
from plyfile import PlyData
from torch import nn
from torch.nn import functional as F
from torch.utils import model_zoo
from skimage import measure, img_as_float32
from igl import adjacency_matrix, connected_components

##################################################
# Below are functions for DPSR

def fftfreqs(res, dtype=torch.float32, exact=True):
 """
 Helper function to return frequency tensors
 :param res: n_dims int tuple of number of frequency modes
 :return:
 """

 n_dims = len(res)
 freqs = []
 for dim in range(n_dims - 1):
 r_ = res[dim]
 freq = np.fft.fftfreq(r_, d=1/r_)
 freqs.append(torch.tensor(freq, dtype=dtype))
 r_ = res[-1]
 if exact:
 freqs.append(torch.tensor(np.fft.rfftfreq(r_, d=1/r_), dtype=dtype))
 else:
 freqs.append(torch.tensor(np.fft.rfftfreq(r_, d=1/r_)[:-1], dtype=dtype))
 omega = torch.meshgrid(freqs)
 omega = list(omega)
 omega = torch.stack(omega, dim=-1)

 return omega

def img(x, deg=1): # imaginary of tensor (assume last dim: real/imag)
 """
 multiply tensor x by i ** deg
 """
 deg %= 4
 if deg == 0:
 res = x
 elif deg == 1:
 res = x[..., [1, 0]]
 res[..., 0] = -res[..., 0]
 elif deg == 2:
 res = -x
 elif deg == 3:
 res = x[..., [1, 0]]
 res[..., 1] = -res[..., 1]
 return res

def spec_gaussian_filter(res, sig):
 omega = fftfreqs(res, dtype=torch.float64) # [dim0, dim1, dim2, d]
 dis = torch.sqrt(torch.sum(omega ** 2, dim=-1))
 filter_ = torch.exp(-0.5*((sig*2*dis/res[0])**2)).unsqueeze(-1).unsqueeze(-1)
 filter_.requires_grad = False

 return filter_

def grid_interp(grid, pts, batched=True):
 """
 :param grid: tensor of shape (batch, *size, in_features)
 :param pts: tensor of shape (batch, num_points, dim) within range (0, 1)
 :return values at query points
 """
 if not batched:
 grid = grid.unsqueeze(0)
 pts = pts.unsqueeze(0)
 dim = pts.shape[-1]
 bs = grid.shape[0]
 size = torch.tensor(grid.shape[1:-1]).to(grid.device).type(pts.dtype)
 cubesize = 1.0 / size
 
 ind0 = torch.floor(pts / cubesize).long() # (batch, num_points, dim)
 ind1 = torch.fmod(torch.ceil(pts / cubesize), size).long() # periodic wrap-around
 ind01 = torch.stack((ind0, ind1), dim=0) # (2, batch, num_points, dim)
 tmp = torch.tensor([0,1],dtype=torch.long)
 com_ = torch.stack(torch.meshgrid(tuple([tmp] * dim)), dim=-1).view(-1, dim)
 dim_ = torch.arange(dim).repeat(com_.shape[0], 1) # (2**dim, dim)
 ind_ = ind01[com_, ..., dim_] # (2**dim, dim, batch, num_points)
 ind_n = ind_.permute(2, 3, 0, 1) # (batch, num_points, 2**dim, dim)
 ind_b = torch.arange(bs).expand(ind_n.shape[1], ind_n.shape[2], bs).permute(2, 0, 1) # (batch, num_points, 2**dim)
 # latent code on neighbor nodes
 if dim == 2:
 lat = grid.clone()[ind_b, ind_n[..., 0], ind_n[..., 1]] # (batch, num_points, 2**dim, in_features)
 else:
 lat = grid.clone()[ind_b, ind_n[..., 0], ind_n[..., 1], ind_n[..., 2]] # (batch, num_points, 2**dim, in_features)

 # weights of neighboring nodes
 xyz0 = ind0.type(cubesize.dtype) * cubesize # (batch, num_points, dim)
 xyz1 = (ind0.type(cubesize.dtype) + 1) * cubesize # (batch, num_points, dim)
 xyz01 = torch.stack((xyz0, xyz1), dim=0) # (2, batch, num_points, dim)
 pos = xyz01[com_, ..., dim_].permute(2,3,0,1) # (batch, num_points, 2**dim, dim)
 pos_ = xyz01[1-com_, ..., dim_].permute(2,3,0,1) # (batch, num_points, 2**dim, dim)
 pos_ = pos_.type(pts.dtype)
 dxyz_ = torch.abs(pts.unsqueeze(-2) - pos_) / cubesize # (batch, num_points, 2**dim, dim)
 weights = torch.prod(dxyz_, dim=-1, keepdim=False) # (batch, num_points, 2**dim)
 query_values = torch.sum(lat * weights.unsqueeze(-1), dim=-2) # (batch, num_points, in_features)
 if not batched:
 query_values = query_values.squeeze(0)
 
 return query_values

def scatter_to_grid(inds, vals, size):
 """
 Scatter update values into empty tensor of size size.
 :param inds: (#values, dims)
 :param vals: (#values)
 :param size: tuple for size. len(size)=dims
 """
 dims = inds.shape[1]
 assert(inds.shape[0] == vals.shape[0])
 assert(len(size) == dims)
 dev = vals.device
 # result = torch.zeros(*size).view(-1).to(dev).type(vals.dtype) # flatten
 # # flatten inds
 result = torch.zeros(*size, device=dev).view(-1).type(vals.dtype) # flatten
 # flatten inds
 fac = [np.prod(size[i+1:]) for i in range(len(size)-1)] + [1]
 fac = torch.tensor(fac, device=dev).type(inds.dtype)
 inds_fold = torch.sum(inds*fac, dim=-1) # [#values,]
 result.scatter_add_(0, inds_fold, vals)
 result = result.view(*size)
 return result

def point_rasterize(pts, vals, size):
 """
 :param pts: point coords, tensor of shape (batch, num_points, dim) within range (0, 1)
 :param vals: point values, tensor of shape (batch, num_points, features)
 :param size: len(size)=dim tuple for grid size
 :return rasterized values (batch, features, res0, res1, res2)
 """
 dim = pts.shape[-1]
 assert(pts.shape[:2] == vals.shape[:2])
 assert(pts.shape[2] == dim)
 size_list = list(size)
 size = torch.tensor(size).to(pts.device).float()
 cubesize = 1.0 / size
 bs = pts.shape[0]
 nf = vals.shape[-1]
 npts = pts.shape[1]
 dev = pts.device
 
 ind0 = torch.floor(pts / cubesize).long() # (batch, num_points, dim)
 ind1 = torch.fmod(torch.ceil(pts / cubesize), size).long() # periodic wrap-around
 ind01 = torch.stack((ind0, ind1), dim=0) # (2, batch, num_points, dim)
 tmp = torch.tensor([0,1],dtype=torch.long)
 com_ = torch.stack(torch.meshgrid(tuple([tmp] * dim)), dim=-1).view(-1, dim)
 dim_ = torch.arange(dim).repeat(com_.shape[0], 1) # (2**dim, dim)
 ind_ = ind01[com_, ..., dim_] # (2**dim, dim, batch, num_points)
 ind_n = ind_.permute(2, 3, 0, 1) # (batch, num_points, 2**dim, dim)
 # ind_b = torch.arange(bs).expand(ind_n.shape[1], ind_n.shape[2], bs).permute(2, 0, 1) # (batch, num_points, 2**dim)
 ind_b = torch.arange(bs, device=dev).expand(ind_n.shape[1], ind_n.shape[2], bs).permute(2, 0, 1) # (batch, num_points, 2**dim)
 
 # weights of neighboring nodes
 xyz0 = ind0.type(cubesize.dtype) * cubesize # (batch, num_points, dim)
 xyz1 = (ind0.type(cubesize.dtype) + 1) * cubesize # (batch, num_points, dim)
 xyz01 = torch.stack((xyz0, xyz1), dim=0) # (2, batch, num_points, dim)
 pos = xyz01[com_, ..., dim_].permute(2,3,0,1) # (batch, num_points, 2**dim, dim)
 pos_ = xyz01[1-com_, ..., dim_].permute(2,3,0,1) # (batch, num_points, 2**dim, dim)
 pos_ = pos_.type(pts.dtype)
 dxyz_ = torch.abs(pts.unsqueeze(-2) - pos_) / cubesize # (batch, num_points, 2**dim, dim)
 weights = torch.prod(dxyz_, dim=-1, keepdim=False) # (batch, num_points, 2**dim)
 
 ind_b = ind_b.unsqueeze(-1).unsqueeze(-1) # (batch, num_points, 2**dim, 1, 1)
 ind_n = ind_n.unsqueeze(-2) # (batch, num_points, 2**dim, 1, dim)
 ind_f = torch.arange(nf, device=dev).view(1, 1, 1, nf, 1) # (1, 1, 1, nf, 1)
 # ind_f = torch.arange(nf).view(1, 1, 1, nf, 1) # (1, 1, 1, nf, 1)
 
 ind_b = ind_b.expand(bs, npts, 2**dim, nf, 1)
 ind_n = ind_n.expand(bs, npts, 2**dim, nf, dim).to(dev)
 ind_f = ind_f.expand(bs, npts, 2**dim, nf, 1)
 inds = torch.cat([ind_b, ind_f, ind_n], dim=-1) # (batch, num_points, 2**dim, nf, 1+1+dim)
 
 # weighted values
 vals = weights.unsqueeze(-1) * vals.unsqueeze(-2) # (batch, num_points, 2**dim, nf)
 
 inds = inds.view(-1, dim+2).permute(1, 0).long() # (1+dim+1, bs*npts*2**dim*nf)
 vals = vals.reshape(-1) # (bs*npts*2**dim*nf)
 tensor_size = [bs, nf] + size_list
 raster = scatter_to_grid(inds.permute(1, 0), vals, [bs, nf] + size_list)
 
 return raster # [batch, nf, res, res, res]



##################################################
# Below are the utilization functions in general

class AverageMeter(object):
 """Computes and stores the average and current value"""
 def __init__(self):
 self.reset()

 def reset(self):
 self.val = 0
 self.n = 0
 self.avg = 0
 self.sum = 0
 self.count = 0

 def update(self, val, n=1):
 self.val = val
 self.n = n
 self.sum += val * n
 self.count += n
 self.avg = self.sum / self.count

 @property
 def valcavg(self):
 return self.val.sum().item() / (self.n != 0).sum().item()

 @property
 def avgcavg(self):
 return self.avg.sum().item() / (self.count != 0).sum().item()

def load_model_manual(state_dict, model):
 new_state_dict = OrderedDict()
 is_model_parallel = isinstance(model, torch.nn.DataParallel)
 for k, v in state_dict.items():
 if k.startswith('module.') != is_model_parallel:
 if k.startswith('module.'):
 # remove module
 k = k[7:]
 else:
 # add module
 k = 'module.' + k

 new_state_dict[k]=v

 model.load_state_dict(new_state_dict)

def mc_from_psr(psr_grid, pytorchify=False, real_scale=False, zero_level=0):
 '''
 Run marching cubes from PSR grid
 '''
 batch_size = psr_grid.shape[0]
 s = psr_grid.shape[-1] # size of psr_grid
 psr_grid_numpy = psr_grid.squeeze().detach().cpu().numpy()
 
 if batch_size>1:
 verts, faces, normals = [], [], []
 for i in range(batch_size):
 verts_cur, faces_cur, normals_cur, values = measure.marching_cubes(psr_grid_numpy[i], level=0)
 verts.append(verts_cur)
 faces.append(faces_cur)
 normals.append(normals_cur)
 verts = np.stack(verts, axis = 0)
 faces = np.stack(faces, axis = 0)
 normals = np.stack(normals, axis = 0)
 else:
 try:
 verts, faces, normals, values = measure.marching_cubes(psr_grid_numpy, level=zero_level)
 except:
 verts, faces, normals, values = measure.marching_cubes(psr_grid_numpy)
 if real_scale:
 verts = verts / (s-1) # scale to range [0, 1]
 else:
 verts = verts / s # scale to range [0, 1)

 if pytorchify:
 device = psr_grid.device
 verts = torch.Tensor(np.ascontiguousarray(verts)).to(device)
 faces = torch.Tensor(np.ascontiguousarray(faces)).to(device)
 normals = torch.Tensor(np.ascontiguousarray(-normals)).to(device)

 return verts, faces, normals
 
def calc_inters_points(verts, faces, pose, img_size, mask_gt=None):
 verts = verts.squeeze()
 faces = faces.squeeze()
 pix_to_face, w, mask = mesh_rasterization(verts, faces, pose, img_size)
 if mask_gt is not None:
 #! only evaluate within the intersection
 mask = mask & mask_gt
 # find 3D points intesected on the mesh
 if True:
 w_masked = w[mask]
 f_p = faces[pix_to_face[mask]].long() # cooresponding faces for each pixel
 # corresponding vertices for p_closest
 v_a, v_b, v_c = verts[f_p[..., 0]], verts[f_p[..., 1]], verts[f_p[..., 2]]
 
 # calculate the intersection point of each pixel and the mesh
 p_inters = w_masked[..., 0, None] * v_a + \
 w_masked[..., 1, None] * v_b + \
 w_masked[..., 2, None] * v_c
 else:
 # backproject ndc to world coordinates using z-buffer
 W, H = img_size[1], img_size[0]
 xy = uv.to(mask.device)[mask]
 x_ndc = 1 - (2*xy[:, 0]) / (W - 1)
 y_ndc = 1 - (2*xy[:, 1]) / (H - 1)
 z = zbuf.squeeze().reshape(H * W)[mask]
 xy_depth = torch.stack((x_ndc, y_ndc, z), dim=1)
 
 p_inters = pose.unproject_points(xy_depth, world_coordinates=True)
 
 # if there are outlier points, we should remove it
 if (p_inters.max()>1) | (p_inters.min()<-1):
 mask_bound = (p_inters>=-1) & (p_inters<=1)
 mask_bound = (mask_bound.sum(dim=-1)==3)
 mask[mask==True] = mask_bound
 p_inters = p_inters[mask_bound]
 print('!!!!!find outlier!')

 return p_inters, mask, f_p, w_masked

def mesh_rasterization(verts, faces, pose, img_size):
 '''
 Use PyTorch3D to rasterize the mesh given a camera 
 '''
 transformed_v = pose.transform_points(verts.detach()) # world -> ndc coordinate system
 if isinstance(pose, PerspectiveCameras):
 transformed_v[..., 2] = 1/transformed_v[..., 2]
 # find p_closest on mesh of each pixel via rasterization
 transformed_mesh = Meshes(verts=[transformed_v], faces=[faces])
 pix_to_face, zbuf, bary_coords, dists = rasterize_meshes(
 transformed_mesh,
 image_size=img_size,
 blur_radius=0,
 faces_per_pixel=1,
 perspective_correct=False
 )
 pix_to_face = pix_to_face.reshape(1, -1) # B x reso x reso -> B x (reso x reso)
 mask = pix_to_face.clone() != -1
 mask = mask.squeeze()
 pix_to_face = pix_to_face.squeeze()
 w = bary_coords.reshape(-1, 3)

 return pix_to_face, w, mask

def verts_on_largest_mesh(verts, faces):
 '''
 verts: Numpy array or Torch.Tensor (N, 3)
 faces: Numpy array (N, 3)
 '''
 if torch.is_tensor(faces):
 verts = verts.squeeze().detach().cpu().numpy()
 faces = faces.squeeze().int().detach().cpu().numpy()

 A = adjacency_matrix(faces)
 num, conn_idx, conn_size = connected_components(A)
 if num == 0:
 v_large, f_large = verts, faces
 else:
 max_idx = conn_size.argmax() # find the index of the largest component
 v_large = verts[conn_idx==max_idx] # keep points on the largest component

 if True:
 mesh_largest = trimesh.Trimesh(verts, faces)
 connected_comp = mesh_largest.split(only_watertight=False)
 mesh_largest = connected_comp[max_idx]
 v_large, f_large = mesh_largest.vertices, mesh_largest.faces
 v_large = v_large.astype(np.float32)
 return v_large, f_large

def update_recursive(dict1, dict2):
 ''' Update two config dictionaries recursively.

 Args:
 dict1 (dict): first dictionary to be updated
 dict2 (dict): second dictionary which entries should be used

 '''
 for k, v in dict2.items():
 if k not in dict1:
 dict1[k] = dict()
 if isinstance(v, dict):
 update_recursive(dict1[k], v)
 else:
 dict1[k] = v

def scale2onet(p, scale=1.2):
 '''
 Scale the point cloud from SAP to ONet range
 '''
 return (p - 0.5) * scale
 
def update_optimizer(inputs, cfg, epoch, model=None, schedule=None):
 if model is not None:
 if schedule is not None:
 optimizer = torch.optim.Adam([
 {"params": model.parameters(),
 "lr": schedule[0].get_learning_rate(epoch)},
 {"params": inputs,
 "lr": schedule[1].get_learning_rate(epoch)}])
 elif 'lr' in cfg['train']:
 optimizer = torch.optim.Adam([
 {"params": model.parameters(),
 "lr": float(cfg['train']['lr'])},
 {"params": inputs,
 "lr": float(cfg['train']['lr_pcl'])}])
 else:
 raise Exception('no known learning rate')
 else:
 if schedule is not None:
 optimizer = torch.optim.Adam([inputs], lr=schedule[0].get_learning_rate(epoch))
 else:
 optimizer = torch.optim.Adam([inputs], lr=float(cfg['train']['lr_pcl']))

 return optimizer


def is_url(url):
 scheme = urllib.parse.urlparse(url).scheme
 return scheme in ('http', 'https')

def load_url(url):
 '''Load a module dictionary from url.
 
 Args:
 url (str): url to saved model
 '''
 print(url)
 print('=> Loading checkpoint from url...')
 state_dict = model_zoo.load_url(url, progress=True)
 
 return state_dict


class GaussianSmoothing(nn.Module):
 """
 Apply gaussian smoothing on a
 1d, 2d or 3d tensor. Filtering is performed seperately for each channel
 in the input using a depthwise convolution.
 Arguments:
 channels (int, sequence): Number of channels of the input tensors. Output will have this number of channels as well.
 kernel_size (int, sequence): Size of the gaussian kernel.
 sigma (float, sequence): Standard deviation of the gaussian kernel.
 dim (int, optional): The number of dimensions of the data.
 Default value is 2 (spatial).
 """
 def __init__(self, channels, kernel_size, sigma, dim=3):
 super(GaussianSmoothing, self).__init__()
 if isinstance(kernel_size, numbers.Number):
 kernel_size = [kernel_size] * dim
 if isinstance(sigma, numbers.Number):
 sigma = [sigma] * dim

 # The gaussian kernel is the product of the
 # gaussian function of each dimension.
 kernel = 1
 meshgrids = torch.meshgrid(
 [
 torch.arange(size, dtype=torch.float32)
 for size in kernel_size
 ]
 )
 for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
 mean = (size - 1) / 2
 kernel *= 1 / (std * math.sqrt(2 * math.pi)) * \
 torch.exp(-((mgrid - mean) / std) ** 2 / 2)

 # Make sure sum of values in gaussian kernel equals 1.
 kernel = kernel / torch.sum(kernel)

 # Reshape to depthwise convolutional weight
 kernel = kernel.view(1, 1, *kernel.size())
 kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))

 self.register_buffer('weight', kernel)
 self.groups = channels

 if dim == 1:
 self.conv = F.conv1d
 elif dim == 2:
 self.conv = F.conv2d
 elif dim == 3:
 self.conv = F.conv3d
 else:
 raise RuntimeError(
 'Only 1, 2 and 3 dimensions are supported. Received {}.'.format(dim)
 )

 def forward(self, input):
 """
 Apply gaussian filter to input.
 Arguments:
 input (torch.Tensor): Input to apply gaussian filter on.
 Returns:
 filtered (torch.Tensor): Filtered output.
 """
 return self.conv(input, weight=self.weight, groups=self.groups)
 
# Originally from https://github.com/amosgropp/IGR/blob/0db06b1273/code/utils/general.py
def get_learning_rate_schedules(schedule_specs):

 schedules = []

 for key in schedule_specs.keys():
 schedules.append(StepLearningRateSchedule(
 schedule_specs[key]['initial'],
 schedule_specs[key]["interval"],
 schedule_specs[key]["factor"], 
 schedule_specs[key]["final"]))
 return schedules

class LearningRateSchedule:
 def get_learning_rate(self, epoch):
 pass
class StepLearningRateSchedule(LearningRateSchedule):
 def __init__(self, initial, interval, factor, final=1e-6):
 self.initial = float(initial)
 self.interval = interval
 self.factor = factor
 self.final = float(final)

 def get_learning_rate(self, epoch):
 lr = np.maximum(self.initial * (self.factor ** (epoch // self.interval)), 5.0e-6)
 if lr > self.final:
 return lr
 else:
 return self.final

def adjust_learning_rate(lr_schedules, optimizer, epoch):
 for i, param_group in enumerate(optimizer.param_groups):
 param_group["lr"] = lr_schedules[i].get_learning_rate(epoch)