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)