mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-04 04:42:49 +08:00
9eeb456e82
Summary: Update all FB license strings to the new format. Reviewed By: patricklabatut Differential Revision: D33403538 fbshipit-source-id: 97a4596c5c888f3c54f44456dc07e718a387a02c
68 lines
2.6 KiB
Python
68 lines
2.6 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import torch
|
|
from pytorch3d import _C
|
|
from torch.autograd import Function
|
|
from torch.autograd.function import once_differentiable
|
|
|
|
|
|
class _MeshFaceAreasNormals(Function):
|
|
"""
|
|
Torch autograd Function wrapper for face areas & normals C++/CUDA implementations.
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(ctx, verts, faces):
|
|
"""
|
|
Args:
|
|
ctx: Context object used to calculate gradients.
|
|
verts: FloatTensor of shape (V, 3), representing the packed
|
|
batch verts tensor.
|
|
faces: LongTensor of shape (F, 3), representing the packed
|
|
batch faces tensor
|
|
Returns:
|
|
areas: FloatTensor of shape (F,) with the areas of each face
|
|
normals: FloatTensor of shape (F,3) with the normals of each face
|
|
"""
|
|
if not (verts.dim() == 2):
|
|
raise ValueError("verts need to be of shape Vx3.")
|
|
if not (verts.shape[1] == 3):
|
|
raise ValueError("verts need to be of shape Vx3.")
|
|
if not (faces.dim() == 2):
|
|
raise ValueError("faces need to be of shape Fx3.")
|
|
if not (faces.shape[1] == 3):
|
|
raise ValueError("faces need to be of shape Fx3.")
|
|
if not (faces.dtype == torch.int64):
|
|
raise ValueError("faces need to be of type torch.int64.")
|
|
# TODO(gkioxari) Change cast to floats once we add support for doubles.
|
|
if not (verts.dtype == torch.float32):
|
|
verts = verts.float()
|
|
|
|
ctx.save_for_backward(verts, faces)
|
|
areas, normals = _C.face_areas_normals_forward(verts, faces)
|
|
return areas, normals
|
|
|
|
@staticmethod
|
|
@once_differentiable
|
|
def backward(ctx, grad_areas, grad_normals):
|
|
grad_areas = grad_areas.contiguous()
|
|
grad_normals = grad_normals.contiguous()
|
|
verts, faces = ctx.saved_tensors
|
|
# TODO(gkioxari) Change cast to floats once we add support for doubles.
|
|
if not (grad_areas.dtype == torch.float32):
|
|
grad_areas = grad_areas.float()
|
|
if not (grad_normals.dtype == torch.float32):
|
|
grad_normals = grad_normals.float()
|
|
grad_verts = _C.face_areas_normals_backward(
|
|
grad_areas, grad_normals, verts, faces
|
|
)
|
|
return grad_verts, None
|
|
|
|
|
|
# pyre-fixme[16]: `_MeshFaceAreasNormals` has no attribute `apply`.
|
|
mesh_face_areas_normals = _MeshFaceAreasNormals.apply
|