Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is | |
| # holder of all proprietary rights on this computer program. | |
| # You can only use this computer program if you have closed | |
| # a license agreement with MPG or you get the right to use the computer | |
| # program from someone who is authorized to grant you that right. | |
| # Any use of the computer program without a valid license is prohibited and | |
| # liable to prosecution. | |
| # | |
| # Copyright©2019 Max-Planck-Gesellschaft zur Förderung | |
| # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute | |
| # for Intelligent Systems. All rights reserved. | |
| # | |
| # Contact: ps-license@tuebingen.mpg.de | |
| import torch | |
| from torch import nn | |
| import trimesh | |
| import math | |
| from typing import NewType | |
| from pytorch3d.structures import Meshes | |
| from pytorch3d.renderer.mesh import rasterize_meshes | |
| Tensor = NewType('Tensor', torch.Tensor) | |
| def solid_angles(points: Tensor, | |
| triangles: Tensor, | |
| thresh: float = 1e-8) -> Tensor: | |
| ''' Compute solid angle between the input points and triangles | |
| Follows the method described in: | |
| The Solid Angle of a Plane Triangle | |
| A. VAN OOSTEROM AND J. STRACKEE | |
| IEEE TRANSACTIONS ON BIOMEDICAL ENGINEERING, | |
| VOL. BME-30, NO. 2, FEBRUARY 1983 | |
| Parameters | |
| ----------- | |
| points: BxQx3 | |
| Tensor of input query points | |
| triangles: BxFx3x3 | |
| Target triangles | |
| thresh: float | |
| float threshold | |
| Returns | |
| ------- | |
| solid_angles: BxQxF | |
| A tensor containing the solid angle between all query points | |
| and input triangles | |
| ''' | |
| # Center the triangles on the query points. Size should be BxQxFx3x3 | |
| centered_tris = triangles[:, None] - points[:, :, None, None] | |
| # BxQxFx3 | |
| norms = torch.norm(centered_tris, dim=-1) | |
| # Should be BxQxFx3 | |
| cross_prod = torch.cross(centered_tris[:, :, :, 1], | |
| centered_tris[:, :, :, 2], | |
| dim=-1) | |
| # Should be BxQxF | |
| numerator = (centered_tris[:, :, :, 0] * cross_prod).sum(dim=-1) | |
| del cross_prod | |
| dot01 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 1]).sum(dim=-1) | |
| dot12 = (centered_tris[:, :, :, 1] * centered_tris[:, :, :, 2]).sum(dim=-1) | |
| dot02 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 2]).sum(dim=-1) | |
| del centered_tris | |
| denominator = (norms.prod(dim=-1) + dot01 * norms[:, :, :, 2] + | |
| dot02 * norms[:, :, :, 1] + dot12 * norms[:, :, :, 0]) | |
| del dot01, dot12, dot02, norms | |
| # Should be BxQ | |
| solid_angle = torch.atan2(numerator, denominator) | |
| del numerator, denominator | |
| torch.cuda.empty_cache() | |
| return 2 * solid_angle | |
| def winding_numbers(points: Tensor, | |
| triangles: Tensor, | |
| thresh: float = 1e-8) -> Tensor: | |
| ''' Uses winding_numbers to compute inside/outside | |
| Robust inside-outside segmentation using generalized winding numbers | |
| Alec Jacobson, | |
| Ladislav Kavan, | |
| Olga Sorkine-Hornung | |
| Fast Winding Numbers for Soups and Clouds SIGGRAPH 2018 | |
| Gavin Barill | |
| NEIL G. Dickson | |
| Ryan Schmidt | |
| David I.W. Levin | |
| and Alec Jacobson | |
| Parameters | |
| ----------- | |
| points: BxQx3 | |
| Tensor of input query points | |
| triangles: BxFx3x3 | |
| Target triangles | |
| thresh: float | |
| float threshold | |
| Returns | |
| ------- | |
| winding_numbers: BxQ | |
| A tensor containing the Generalized winding numbers | |
| ''' | |
| # The generalized winding number is the sum of solid angles of the point | |
| # with respect to all triangles. | |
| return 1 / (4 * math.pi) * solid_angles(points, triangles, | |
| thresh=thresh).sum(dim=-1) | |
| def batch_contains(verts, faces, points): | |
| B = verts.shape[0] | |
| N = points.shape[1] | |
| verts = verts.detach().cpu() | |
| faces = faces.detach().cpu() | |
| points = points.detach().cpu() | |
| contains = torch.zeros(B, N) | |
| for i in range(B): | |
| contains[i] = torch.as_tensor( | |
| trimesh.Trimesh(verts[i], faces[i]).contains(points[i])) | |
| return 2.0 * (contains - 0.5) | |
| def dict2obj(d): | |
| # if isinstance(d, list): | |
| # d = [dict2obj(x) for x in d] | |
| if not isinstance(d, dict): | |
| return d | |
| class C(object): | |
| pass | |
| o = C() | |
| for k in d: | |
| o.__dict__[k] = dict2obj(d[k]) | |
| return o | |
| def face_vertices(vertices, faces): | |
| """ | |
| :param vertices: [batch size, number of vertices, 3] | |
| :param faces: [batch size, number of faces, 3] | |
| :return: [batch size, number of faces, 3, 3] | |
| """ | |
| bs, nv = vertices.shape[:2] | |
| bs, nf = faces.shape[:2] | |
| device = vertices.device | |
| faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) * | |
| nv)[:, None, None] | |
| vertices = vertices.reshape((bs * nv, vertices.shape[-1])) | |
| return vertices[faces.long()] | |
| class Pytorch3dRasterizer(nn.Module): | |
| """ Borrowed from https://github.com/facebookresearch/pytorch3d | |
| Notice: | |
| x,y,z are in image space, normalized | |
| can only render squared image now | |
| """ | |
| def __init__(self, image_size=224): | |
| """ | |
| use fixed raster_settings for rendering faces | |
| """ | |
| super().__init__() | |
| raster_settings = { | |
| 'image_size': image_size, | |
| 'blur_radius': 0.0, | |
| 'faces_per_pixel': 1, | |
| 'bin_size': None, | |
| 'max_faces_per_bin': None, | |
| 'perspective_correct': True, | |
| 'cull_backfaces': True, | |
| } | |
| raster_settings = dict2obj(raster_settings) | |
| self.raster_settings = raster_settings | |
| def forward(self, vertices, faces, attributes=None): | |
| fixed_vertices = vertices.clone() | |
| fixed_vertices[..., :2] = -fixed_vertices[..., :2] | |
| meshes_screen = Meshes(verts=fixed_vertices.float(), | |
| faces=faces.long()) | |
| raster_settings = self.raster_settings | |
| pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( | |
| meshes_screen, | |
| image_size=raster_settings.image_size, | |
| blur_radius=raster_settings.blur_radius, | |
| faces_per_pixel=raster_settings.faces_per_pixel, | |
| bin_size=raster_settings.bin_size, | |
| max_faces_per_bin=raster_settings.max_faces_per_bin, | |
| perspective_correct=raster_settings.perspective_correct, | |
| ) | |
| vismask = (pix_to_face > -1).float() | |
| D = attributes.shape[-1] | |
| attributes = attributes.clone() | |
| attributes = attributes.view(attributes.shape[0] * attributes.shape[1], | |
| 3, attributes.shape[-1]) | |
| N, H, W, K, _ = bary_coords.shape | |
| mask = pix_to_face == -1 | |
| pix_to_face = pix_to_face.clone() | |
| pix_to_face[mask] = 0 | |
| idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D) | |
| pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D) | |
| pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2) | |
| pixel_vals[mask] = 0 # Replace masked values in output. | |
| pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2) | |
| pixel_vals = torch.cat( | |
| [pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1) | |
| return pixel_vals | |