# -*- 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
from lib.dataset.mesh_util import projection
from lib.common.render import Render
import numpy as np
import torch
from torchvision.utils import make_grid
from pytorch3d import _C
from torch.autograd import Function
from torch.autograd.function import once_differentiable
from pytorch3d.structures import Pointclouds
from PIL import Image
from typing import Tuple
from pytorch3d.ops.mesh_face_areas_normals import mesh_face_areas_normals
from pytorch3d.ops.packed_to_padded import packed_to_padded
_DEFAULT_MIN_TRIANGLE_AREA: float = 5e-3
# PointFaceDistance
class _PointFaceDistance(Function):
"""
Torch autograd Function wrapper PointFaceDistance Cuda implementation
"""
@staticmethod
def forward(
ctx,
points,
points_first_idx,
tris,
tris_first_idx,
max_points,
min_triangle_area=_DEFAULT_MIN_TRIANGLE_AREA,
):
"""
Args:
ctx: Context object used to calculate gradients.
points: FloatTensor of shape `(P, 3)`
points_first_idx: LongTensor of shape `(N,)` indicating the first point
index in each example in the batch
tris: FloatTensor of shape `(T, 3, 3)` of triangular faces. The `t`-th
triangular face is spanned by `(tris[t, 0], tris[t, 1], tris[t, 2])`
tris_first_idx: LongTensor of shape `(N,)` indicating the first face
index in each example in the batch
max_points: Scalar equal to maximum number of points in the batch
min_triangle_area: (float, defaulted) Triangles of area less than this
will be treated as points/lines.
Returns:
dists: FloatTensor of shape `(P,)`, where `dists[p]` is the squared
euclidean distance of `p`-th point to the closest triangular face
in the corresponding example in the batch
idxs: LongTensor of shape `(P,)` indicating the closest triangular face
in the corresponding example in the batch.
`dists[p]` is
`d(points[p], tris[idxs[p], 0], tris[idxs[p], 1], tris[idxs[p], 2])`
where `d(u, v0, v1, v2)` is the distance of point `u` from the triangular
face `(v0, v1, v2)`
"""
dists, idxs = _C.point_face_dist_forward(
points,
points_first_idx,
tris,
tris_first_idx,
max_points,
min_triangle_area,
)
ctx.save_for_backward(points, tris, idxs)
ctx.min_triangle_area = min_triangle_area
return dists, idxs
@staticmethod
@once_differentiable
def backward(ctx, grad_dists):
grad_dists = grad_dists.contiguous()
points, tris, idxs = ctx.saved_tensors
min_triangle_area = ctx.min_triangle_area
grad_points, grad_tris = _C.point_face_dist_backward(
points, tris, idxs, grad_dists, min_triangle_area
)
return grad_points, None, grad_tris, None, None, None
def _rand_barycentric_coords(
size1, size2, dtype: torch.dtype, device: torch.device
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Helper function to generate random barycentric coordinates which are uniformly
distributed over a triangle.
Args:
size1, size2: The number of coordinates generated will be size1*size2.
Output tensors will each be of shape (size1, size2).
dtype: Datatype to generate.
device: A torch.device object on which the outputs will be allocated.
Returns:
w0, w1, w2: Tensors of shape (size1, size2) giving random barycentric
coordinates
"""
uv = torch.rand(2, size1, size2, dtype=dtype, device=device)
u, v = uv[0], uv[1]
u_sqrt = u.sqrt()
w0 = 1.0 - u_sqrt
w1 = u_sqrt * (1.0 - v)
w2 = u_sqrt * v
w = torch.cat([w0[..., None], w1[..., None], w2[..., None]], dim=2)
return w
def sample_points_from_meshes(meshes, num_samples: int = 10000):
"""
Convert a batch of meshes to a batch of pointclouds by uniformly sampling
points on the surface of the mesh with probability proportional to the
face area.
Args:
meshes: A Meshes object with a batch of N meshes.
num_samples: Integer giving the number of point samples per mesh.
return_normals: If True, return normals for the sampled points.
return_textures: If True, return textures for the sampled points.
Returns:
3-element tuple containing
- **samples**: FloatTensor of shape (N, num_samples, 3) giving the
coordinates of sampled points for each mesh in the batch. For empty
meshes the corresponding row in the samples array will be filled with 0.
- **normals**: FloatTensor of shape (N, num_samples, 3) giving a normal vector
to each sampled point. Only returned if return_normals is True.
For empty meshes the corresponding row in the normals array will
be filled with 0.
- **textures**: FloatTensor of shape (N, num_samples, C) giving a C-dimensional
texture vector to each sampled point. Only returned if return_textures is True.
For empty meshes the corresponding row in the textures array will
be filled with 0.
Note that in a future releases, we will replace the 3-element tuple output
with a `Pointclouds` datastructure, as follows
.. code-block:: python
Pointclouds(samples, normals=normals, features=textures)
"""
if meshes.isempty():
raise ValueError("Meshes are empty.")
verts = meshes.verts_packed()
if not torch.isfinite(verts).all():
raise ValueError("Meshes contain nan or inf.")
faces = meshes.faces_packed()
mesh_to_face = meshes.mesh_to_faces_packed_first_idx()
num_meshes = len(meshes)
num_valid_meshes = torch.sum(meshes.valid) # Non empty meshes.
# Initialize samples tensor with fill value 0 for empty meshes.
samples = torch.zeros((num_meshes, num_samples, 3), device=meshes.device)
# Only compute samples for non empty meshes
with torch.no_grad():
areas, _ = mesh_face_areas_normals(verts, faces) # Face areas can be zero.
max_faces = meshes.num_faces_per_mesh().max().item()
areas_padded = packed_to_padded(areas, mesh_to_face[meshes.valid], max_faces) # (N, F)
# TODO (gkioxari) Confirm multinomial bug is not present with real data.
samples_face_idxs = areas_padded.multinomial(
num_samples, replacement=True
) # (N, num_samples)
samples_face_idxs += mesh_to_face[meshes.valid].view(num_valid_meshes, 1)
# Randomly generate barycentric coords.
# w (N, num_samples, 3)
# sample_face_idxs (N, num_samples)
# samples_verts (N, num_samples, 3, 3)
samples_bw = _rand_barycentric_coords(num_valid_meshes, num_samples, verts.dtype, verts.device)
sample_verts = verts[faces][samples_face_idxs]
samples[meshes.valid] = (sample_verts * samples_bw[..., None]).sum(dim=-2)
return samples, samples_face_idxs, samples_bw
def econ_point_mesh_distance(meshes, pcls, weighted=True):
if len(meshes) != len(pcls):
raise ValueError("meshes and pointclouds must be equal sized batches")
# packed representation for pointclouds
points = pcls.points_packed() # (P, 3)
points_first_idx = pcls.cloud_to_packed_first_idx()
max_points = pcls.num_points_per_cloud().max().item()
# packed representation for faces
verts_packed = meshes.verts_packed()
faces_packed = meshes.faces_packed()
tris = verts_packed[faces_packed] # (T, 3, 3)
tris_first_idx = meshes.mesh_to_faces_packed_first_idx()
# point to face distance: shape (P,)
point_to_face, idxs = _PointFaceDistance.apply(
points, points_first_idx, tris, tris_first_idx, max_points, 5e-3
)
if weighted:
# weight each example by the inverse of number of points in the example
point_to_cloud_idx = pcls.packed_to_cloud_idx() # (sum(P_i),)
num_points_per_cloud = pcls.num_points_per_cloud() # (N,)
weights_p = num_points_per_cloud.gather(0, point_to_cloud_idx)
weights_p = 1.0 / weights_p.float()
point_to_face = torch.sqrt(point_to_face) * weights_p
return point_to_face, idxs
class Evaluator:
def __init__(self, device):
self.render = Render(size=512, device=device)
self.device = device
def set_mesh(self, result_dict, scale=True):
for k, v in result_dict.items():
setattr(self, k, v)
if scale:
self.verts_pr -= self.recon_size / 2.0
self.verts_pr /= self.recon_size / 2.0
self.verts_gt = projection(self.verts_gt, self.calib)
self.verts_gt[:, 1] *= -1
self.render.load_meshes(self.verts_pr, self.faces_pr)
self.src_mesh = self.render.meshes
self.render.load_meshes(self.verts_gt, self.faces_gt)
self.tgt_mesh = self.render.meshes
def calculate_normal_consist(self, normal_path):
self.render.meshes = self.src_mesh
src_normal_imgs = self.render.get_image(cam_type="all", bg="black")
self.render.meshes = self.tgt_mesh
tgt_normal_imgs = self.render.get_image(cam_type="all", bg="black")
error_list = []
if len(src_normal_imgs)>4:
# for i in range(len(src_normal_imgs)):
src_normal_arr = make_grid(torch.cat(src_normal_imgs, dim=0), nrow=6,padding=1) # [0,1]
tgt_normal_arr = make_grid(torch.cat(tgt_normal_imgs, dim=0), nrow=6,padding=1) # [0,1]
# src_normal_arr = make_grid(torch.cat(src_normal_imgs, dim=0), nrow=4,padding=0) # [0,1]
# tgt_normal_arr = make_grid(torch.cat(tgt_normal_imgs, dim=0), nrow=4,padding=0) # [0,1]
src_norm = torch.norm(src_normal_arr, dim=0, keepdim=True)
tgt_norm = torch.norm(tgt_normal_arr, dim=0, keepdim=True)
src_norm[src_norm == 0.0] = 1.0
tgt_norm[tgt_norm == 0.0] = 1.0
src_normal_arr /= src_norm
tgt_normal_arr /= tgt_norm
# sim_mask = self.get_laplacian_2d(tgt_normal_arr).to(self.device)
src_normal_arr = (src_normal_arr + 1.0) * 0.5
tgt_normal_arr = (tgt_normal_arr + 1.0) * 0.5
error = ((
(src_normal_arr - tgt_normal_arr)**2).sum(dim=0).mean()) * 4
#error_list.append(error)
normal_img = Image.fromarray(
(torch.cat([src_normal_arr, tgt_normal_arr], dim=1).permute(
1, 2, 0).detach().cpu().numpy() * 255.0).astype(np.uint8))
normal_img.save(normal_path)
return error
else:
src_normal_arr = make_grid(torch.cat(src_normal_imgs, dim=0), nrow=4,padding=0) # [0,1]
tgt_normal_arr = make_grid(torch.cat(tgt_normal_imgs, dim=0), nrow=4,padding=0) # [0,1]
src_norm = torch.norm(src_normal_arr, dim=0, keepdim=True)
tgt_norm = torch.norm(tgt_normal_arr, dim=0, keepdim=True)
src_norm[src_norm == 0.0] = 1.0
tgt_norm[tgt_norm == 0.0] = 1.0
src_normal_arr /= src_norm
tgt_normal_arr /= tgt_norm
# sim_mask = self.get_laplacian_2d(tgt_normal_arr).to(self.device)
src_normal_arr = (src_normal_arr + 1.0) * 0.5
tgt_normal_arr = (tgt_normal_arr + 1.0) * 0.5
error = ((
(src_normal_arr - tgt_normal_arr)**2).sum(dim=0).mean()) * 4
return error
def calculate_chamfer_p2s(self, num_samples=1000):
samples_tgt, _, _ = sample_points_from_meshes(self.tgt_mesh, num_samples)
samples_src, _, _ = sample_points_from_meshes(self.src_mesh, num_samples)
tgt_points = Pointclouds(samples_tgt)
src_points = Pointclouds(samples_src)
p2s_dist = point_mesh_distance(self.src_mesh, tgt_points)[0].sum() * 100.0
chamfer_dist = (
point_mesh_distance(self.tgt_mesh, src_points)[0].sum() * 100.0 + p2s_dist
) * 0.5
return chamfer_dist, p2s_dist
def calc_acc(self, output, target, thres=0.5, use_sdf=False):
# # remove the surface points with thres
# non_surf_ids = (target != thres)
# output = output[non_surf_ids]
# target = target[non_surf_ids]
with torch.no_grad():
output = output.masked_fill(output < thres, 0.0)
output = output.masked_fill(output > thres, 1.0)
if use_sdf:
target = target.masked_fill(target < thres, 0.0)
target = target.masked_fill(target > thres, 1.0)
acc = output.eq(target).float().mean()
# iou, precison, recall
output = output > thres
target = target > thres
union = output | target
inter = output & target
_max = torch.tensor(1.0).to(output.device)
union = max(union.sum().float(), _max)
true_pos = max(inter.sum().float(), _max)
vol_pred = max(output.sum().float(), _max)
vol_gt = max(target.sum().float(), _max)
return acc, true_pos / union, true_pos / vol_pred, true_pos / vol_gt