# -*- 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