1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
|
import torch
from torch_geometric.data import Data
from torch_geometric.data.datapipes import functional_transform
from torch_geometric.transforms import BaseTransform
@functional_transform('sample_points')
class SamplePoints(BaseTransform):
r"""Uniformly samples a fixed number of points on the mesh faces according
to their face area (functional name: :obj:`sample_points`).
Args:
num (int): The number of points to sample.
remove_faces (bool, optional): If set to :obj:`False`, the face tensor
will not be removed. (default: :obj:`True`)
include_normals (bool, optional): If set to :obj:`True`, then compute
normals for each sampled point. (default: :obj:`False`)
"""
def __init__(
self,
num: int,
remove_faces: bool = True,
include_normals: bool = False,
):
self.num = num
self.remove_faces = remove_faces
self.include_normals = include_normals
def forward(self, data: Data) -> Data:
assert data.pos is not None
assert data.face is not None
pos, face = data.pos, data.face
assert pos.size(1) == 3 and face.size(0) == 3
pos_max = pos.abs().max()
pos = pos / pos_max
area = (pos[face[1]] - pos[face[0]]).cross(
pos[face[2]] - pos[face[0]],
dim=1,
)
area = area.norm(p=2, dim=1).abs() / 2
prob = area / area.sum()
sample = torch.multinomial(prob, self.num, replacement=True)
face = face[:, sample]
frac = torch.rand(self.num, 2, device=pos.device)
mask = frac.sum(dim=-1) > 1
frac[mask] = 1 - frac[mask]
vec1 = pos[face[1]] - pos[face[0]]
vec2 = pos[face[2]] - pos[face[0]]
if self.include_normals:
data.normal = torch.nn.functional.normalize(
vec1.cross(vec2, dim=1), p=2)
pos_sampled = pos[face[0]]
pos_sampled += frac[:, :1] * vec1
pos_sampled += frac[:, 1:] * vec2
pos_sampled = pos_sampled * pos_max
data.pos = pos_sampled
if self.remove_faces:
data.face = None
return data
def __repr__(self) -> str:
return f'{self.__class__.__name__}({self.num})'
|
