Add utility function to tesselate a torus

Summary: Add utility function to tesselate a torus, to be used in more complex mesh I/O benchmarks

Reviewed By: bottler

Differential Revision: D20390724

fbshipit-source-id: 882bbbe9cac81cf340a34495b9aa66e3c1ddeebc
This commit is contained in:
Patrick Labatut 2020-03-13 04:28:30 -07:00 committed by Facebook GitHub Bot
parent 098554d323
commit 327868b86e
2 changed files with 72 additions and 0 deletions

View File

@ -1,5 +1,6 @@
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
from .ico_sphere import ico_sphere
from .torus import torus
__all__ = [k for k in globals().keys() if not k.startswith("_")]

71
pytorch3d/utils/torus.py Normal file
View File

@ -0,0 +1,71 @@
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
from itertools import tee
from math import cos, pi, sin
from typing import Iterator, Optional, Tuple
import torch
from pytorch3d.structures.meshes import Meshes
# Make an iterator over the adjacent pairs: (-1, 0), (0, 1), ..., (N - 2, N - 1)
def _make_pair_range(N: int) -> Iterator[Tuple[int, int]]:
i, j = tee(range(-1, N))
next(j, None)
return zip(i, j)
def torus(
r: float,
R: float,
sides: int,
rings: int,
device: Optional[torch.device] = None,
) -> Meshes:
"""
Create vertices and faces for a torus.
Args:
r: Inner radius of the torus.
R: Outer radius of the torus.
sides: Number of inner divisions.
rings: Number of outer divisions.
device: Device on which the outputs will be allocated.
Returns:
Meshes object with the generated vertices and faces.
"""
if not (sides > 0):
raise ValueError("sides must be > 0.")
if not (rings > 0):
raise ValueError("rings must be > 0.")
device = device if device else torch.device("cpu")
verts = []
for i in range(rings):
# phi ranges from 0 to 2 pi (rings - 1) / rings
phi = 2 * pi * i / rings
for j in range(sides):
# theta ranges from 0 to 2 pi (sides - 1) / sides
theta = 2 * pi * j / sides
x = (R + r * cos(theta)) * cos(phi)
y = (R + r * cos(theta)) * sin(phi)
z = r * sin(theta)
# This vertex has index i * sides + j
verts.append([x, y, z])
faces = []
for i0, i1 in _make_pair_range(rings):
index0 = (i0 % rings) * sides
index1 = (i1 % rings) * sides
for j0, j1 in _make_pair_range(sides):
index00 = index0 + (j0 % sides)
index01 = index0 + (j1 % sides)
index10 = index1 + (j0 % sides)
index11 = index1 + (j1 % sides)
faces.append([index00, index10, index11])
faces.append([index11, index01, index00])
verts_list = [torch.tensor(verts, dtype=torch.float32, device=device)]
faces_list = [torch.tensor(faces, dtype=torch.int64, device=device)]
return Meshes(verts_list, faces_list)