def test_cpu_vs_cuda_naive(self):
"""
Compare naive versions of cuda and cpp
"""
torch.manual_seed(231)
image_size = 64
radius = 0.1**2
faces_per_pixel = 3
device = torch.device("cpu")
meshes_cpu = ico_sphere(0, device)
verts1, faces1 = meshes_cpu.get_mesh_verts_faces(0)
verts1.requires_grad = True
meshes_cpu = Meshes(verts=[verts1], faces=[faces1])
device = torch.device("cuda:0")
meshes_cuda = ico_sphere(0, device)
verts2, faces2 = meshes_cuda.get_mesh_verts_faces(0)
verts2.requires_grad = True
meshes_cuda = Meshes(verts=[verts2], faces=[faces2])
args_cpu = (meshes_cpu, image_size, radius, faces_per_pixel)
args_cuda = (meshes_cuda, image_size, radius, faces_per_pixel, 0, 0)
self._compare_impls(
rasterize_meshes,
rasterize_meshes,
args_cpu,
args_cuda,
verts1,
verts2,
compare_grads=True,
)
def test_compare_coarse_cpu_vs_cuda(self):
torch.manual_seed(231)
N = 1
image_size = 512
blur_radius = 0.0
bin_size = 32
max_faces_per_bin = 20
device = torch.device("cpu")
meshes = ico_sphere(2, device)
faces = meshes.faces_packed()
verts = meshes.verts_packed()
faces_verts = verts[faces]
num_faces_per_mesh = meshes.num_faces_per_mesh()
mesh_to_face_first_idx = meshes.mesh_to_faces_packed_first_idx()
args = (
faces_verts,
mesh_to_face_first_idx,
num_faces_per_mesh,
image_size,
blur_radius,
bin_size,
max_faces_per_bin,
)
bin_faces_cpu = _C._rasterize_meshes_coarse(*args)
device = torch.device("cuda:0")
meshes = ico_sphere(2, device)
faces = meshes.faces_packed()
verts = meshes.verts_packed()
faces_verts = verts[faces]
num_faces_per_mesh = meshes.num_faces_per_mesh()
mesh_to_face_first_idx = meshes.mesh_to_faces_packed_first_idx()
args = (
faces_verts,
mesh_to_face_first_idx,
num_faces_per_mesh,
image_size,
blur_radius,
bin_size,
max_faces_per_bin,
)
bin_faces_cuda = _C._rasterize_meshes_coarse(*args)
# Bin faces might not be the same: CUDA version might write them in
# any order. But if we sort the non-(-1) elements of the CUDA output
# then they should be the same.
for n in range(N):
for by in range(bin_faces_cpu.shape[1]):
for bx in range(bin_faces_cpu.shape[2]):
K = (bin_faces_cuda[n, by, bx] != -1).sum().item()
idxs_cpu = bin_faces_cpu[n, by, bx].tolist()
idxs_cuda = bin_faces_cuda[n, by, bx].tolist()
idxs_cuda[:K] = sorted(idxs_cuda[:K])
self.assertEqual(idxs_cpu, idxs_cuda)
def rasterize_meshes_cuda_with_init(
num_meshes: int,
ico_level: int,
image_size: int,
blur_radius: float,
bin_size: int,
max_faces_per_bin: int,
):
meshes = ico_sphere(ico_level, torch.device("cuda:0"))
meshes_batch = meshes.extend(num_meshes)
torch.cuda.synchronize()
def rasterize():
rasterize_meshes(
meshes_batch,
image_size,
blur_radius,
8,
bin_size,
max_faces_per_bin,
)
torch.cuda.synchronize()
return rasterize
# Helpful comments below.# Helpful comments below.# Helpful comments below.# Helpful comments below.# Helpful comments below.# Helpful comments below.
def test_gather_scatter(self):
"""
Check gather_scatter cuda and python versions give the same results.
Check that gather_scatter cuda version throws an error if cpu tensors
are given as input.
"""
device = get_random_cuda_device()
mesh = ico_sphere()
verts = mesh.verts_packed()
edges = mesh.edges_packed()
w0 = nn.Linear(3, 1)
input = w0(verts)
# undirected
output_python = gather_scatter_python(input, edges, False)
output_cuda = _C.gather_scatter(input.to(device=device),
edges.to(device=device), False, False)
self.assertClose(output_cuda.cpu(), output_python)
output_cpu = _C.gather_scatter(input.cpu(), edges.cpu(), False, False)
self.assertClose(output_cpu, output_python)
# directed
output_python = gather_scatter_python(input, edges, True)
output_cuda = _C.gather_scatter(input.to(device=device),
edges.to(device=device), True, False)
self.assertClose(output_cuda.cpu(), output_python)
output_cpu = _C.gather_scatter(input.cpu(), edges.cpu(), True, False)
self.assertClose(output_cpu, output_python)
def test_gather_scatter(self):
"""
Check gather_scatter cuda and python versions give the same results.
Check that gather_scatter cuda version throws an error if cpu tensors
are given as input.
"""
device = torch.device("cuda:0")
mesh = ico_sphere()
verts = mesh.verts_packed()
edges = mesh.edges_packed()
w0 = nn.Linear(3, 1)
input = w0(verts)
# output
output_cpu = gather_scatter_python(input, edges, False)
output_cuda = _C.gather_scatter(input.to(device=device),
edges.to(device=device), False, False)
self.assertClose(output_cuda.cpu(), output_cpu)
with self.assertRaises(Exception) as err:
_C.gather_scatter(input.cpu(), edges.cpu(), False, False)
self.assertTrue("Not implemented on the CPU" in str(err.exception))
# directed
output_cpu = gather_scatter_python(input, edges, True)
output_cuda = _C.gather_scatter(input.to(device=device),
edges.to(device=device), True, False)
self.assertClose(output_cuda.cpu(), output_cpu)
def test_ball_query_output_simple(self):
device = get_random_cuda_device()
N, P1, P2, K = 5, 8, 16, 4
sphere = ico_sphere(level=2, device=device).extend(N)
points_1 = sample_points_from_meshes(sphere, P1)
points_2 = sample_points_from_meshes(sphere, P2) * 5.0
radius = 6.0
naive_out = self._ball_query_naive(
points_1, points_2, lengths1=None, lengths2=None, K=K, radius=radius
)
cuda_out = ball_query(points_1, points_2, K=K, radius=radius)
# All points should have N sample neighbors as radius is large
# Zero is a valid index but can only be present once (i.e. no zero padding)
naive_out_zeros = (naive_out.idx == 0).sum(dim=-1).max()
cuda_out_zeros = (cuda_out.idx == 0).sum(dim=-1).max()
self.assertTrue(naive_out_zeros == 0 or naive_out_zeros == 1)
self.assertTrue(cuda_out_zeros == 0 or cuda_out_zeros == 1)
# All points should now have zero sample neighbors as radius is small
radius = 0.5
naive_out = self._ball_query_naive(
points_1, points_2, lengths1=None, lengths2=None, K=K, radius=radius
)
cuda_out = ball_query(points_1, points_2, K=K, radius=radius)
naive_out_allzeros = (naive_out.idx == -1).all()
cuda_out_allzeros = (cuda_out.idx == -1).sum()
self.assertTrue(naive_out_allzeros)
self.assertTrue(cuda_out_allzeros)
def test_cuda_naive_vs_binned_perspective_correct(self):
meshes = ico_sphere(2, device=torch.device("cuda"))
verts1, faces1 = meshes.get_mesh_verts_faces(0)
verts1.requires_grad = True
meshes1 = Meshes(verts=[verts1], faces=[faces1])
verts2 = verts1.detach().clone().requires_grad_(True)
faces2 = faces1.detach().clone()
meshes2 = Meshes(verts=[verts2], faces=[faces2])
kwargs = {"image_size": 64, "perspective_correct": True}
fn1 = functools.partial(rasterize_meshes,
meshes1,
bin_size=0,
**kwargs)
fn2 = functools.partial(rasterize_meshes,
meshes2,
bin_size=8,
**kwargs)
args = ()
self._compare_impls(fn1,
fn2,
args,
args,
verts1,
verts2,
compare_grads=True)
def transform_meshes_to_camera_coord_system(meshes, boxes, zranges, Ks, imsize):
device = meshes.device
new_verts, new_faces = [], []
h, w = imsize
im_size = torch.tensor([w, h], device=device).view(1, 2)
assert len(meshes) == len(zranges)
for i in range(len(meshes)):
verts, faces = meshes.get_mesh_verts_faces(i)
if verts.numel() == 0:
verts, faces = ico_sphere(level=3, device=device).get_mesh_verts_faces(0)
assert not torch.isnan(verts).any()
assert not torch.isnan(faces).any()
roi = boxes[i].view(1, 4)
zrange = zranges[i].view(1, 2)
K = Ks[i].view(1, 3)
cub3D = shape_utils.box2D_to_cuboid3D(zrange, K, roi, im_size)
txz, tyz = shape_utils.cuboid3D_to_unitbox3D(cub3D)
# image to camera coords
verts[:, 0] = -verts[:, 0]
verts[:, 1] = -verts[:, 1]
# transform to destination size
xz = verts[:, [0, 2]]
yz = verts[:, [1, 2]]
pxz = txz.inverse(xz.view(1, -1, 2)).squeeze(0)
pyz = tyz.inverse(yz.view(1, -1, 2)).squeeze(0)
verts = torch.stack([pxz[:, 0], pyz[:, 0], pxz[:, 1]], dim=1).to(
device, dtype=torch.float32
)
new_verts.append(verts)
new_faces.append(faces)
return Meshes(verts=new_verts, faces=new_faces)
def rasterize_meshes_cpu_with_init(num_meshes: int, ico_level: int,
image_size: int, blur_radius: float):
meshes = ico_sphere(ico_level, torch.device("cpu"))
meshes_batch = meshes.extend(num_meshes)
def rasterize():
rasterize_meshes(meshes_batch, image_size, blur_radius, bin_size=0)
return rasterize
def deform_sphere():
meshes = ico_sphere(3)
meshes = ARAP_from_meshes(meshes) # convert to ARAP obejct
N = meshes.num_verts_per_mesh()[0]
handle_verts = [26]
handle_pos = meshes.verts_padded()[0][handle_verts]
handle_pos_shifted = handle_pos.clone()
# static as furthest vert
static_verts = [
max(range(N),
key=lambda x: torch.norm(meshes.verts_padded()[0][x] - handle_pos[
0]))
]
static_verts = add_n_ring_neighbours(meshes, static_verts, n=5)
trisurfs = plot_meshes(ax,
meshes,
handle_verts=handle_verts,
static_verts=static_verts,
prop=False,
change_lims=True,
color="gray")
disp_vec = meshes.C[0] - handle_pos[0] # displace towards centre of mass
n_frames = 100
disp_frac = 1.2 # fraction of full disp_vec to move in animation
step = disp_frac * 4 / n_frames # moves
def anim(i):
[x.remove() for x in trisurfs] # remove previous frame's mesh
if i < n_frames / 4 or i > 3 * n_frames / 4:
direction = 1
else:
direction = -1
handle_pos_shifted[0] += direction * step * disp_vec
## deform, replot
meshes.solve(static_verts=static_verts,
handle_verts=handle_verts,
handle_verts_pos=handle_pos_shifted,
n_its=1) ## run ARAP
trisurfs[:] = plot_meshes(ax,
meshes,
handle_verts=handle_verts,
static_verts=static_verts,
prop=False,
color="gray")
ax.axis("off")
save_animation(fig, anim, n_frames=n_frames, title="sphere", fps=30)
def rasterize_meshes_python_with_init(num_meshes: int, ico_level: int,
image_size: int, blur_radius: float):
device = torch.device("cpu")
meshes = ico_sphere(ico_level, device)
meshes_batch = meshes.extend(num_meshes)
def rasterize():
rasterize_meshes_python(meshes_batch, image_size, blur_radius)
return rasterize
def forward(self, imgs):
N = imgs.shape[0]
device = imgs.device
img_feats = self.backbone(imgs)
P = self._get_projection_matrix(N, device)
init_meshes = ico_sphere(self.ico_sphere_level, device).extend(N)
refined_meshes = self.mesh_head(img_feats, init_meshes, P, subdivide=True)
return None, refined_meshes
def create_model(cfg, device, mode="train", camera_model=None, **kwargs):
''' Returns model
Args:
cfg (edict): imported yaml config
device (device): pytorch device
'''
if cfg.model.type == 'point':
decoder = None
texture = None
use_lighting = (cfg.renderer is not None
and not cfg.renderer.get('is_neural_texture', True))
if use_lighting:
texture = LightingTexture()
else:
if 'rgb' not in cfg.model.decoder_kwargs.out_dims:
Texture = get_class_from_string(cfg.model.texture_type)
cfg.model.texture_kwargs[
'c_dim'] = cfg.model.decoder_kwargs.out_dims.get('latent', 0)
texture_decoder = Texture(**cfg.model.texture_kwargs)
else:
texture_decoder = decoder
logger_py.info("Decoder used as NeuralTexture")
texture = NeuralTexture(
view_dependent=cfg.model.texture_kwargs.view_dependent,
decoder=texture_decoder).to(device=device)
logger_py.info("Created NeuralTexture {}".format(texture.__class__))
logger_py.info(texture)
Model = get_class_from_string("DSS.models.{}_modeling.Model".format(
cfg.model.type))
# if not using decoder, then use non-parameterized point renderer
# create icosphere as initial point cloud
sphere_mesh = ico_sphere(level=4)
sphere_mesh.scale_verts_(0.5)
points, normals = sample_points_from_meshes(
sphere_mesh,
num_samples=int(cfg['model']['model_kwargs']['n_points_per_cloud']),
return_normals=True)
colors = torch.ones_like(points)
renderer = create_renderer(cfg.renderer).to(device)
model = Model(
points,
normals,
colors,
renderer,
device=device,
texture=texture,
**cfg.model.model_kwargs,
).to(device=device)
return model
def __init__(
self,
batch_size=2,
radius=0.2,
camextr=None,
hand_links=None,
debug=True,
random_rot=True,
z_off=0.5,
y_off=0,
x_off=0,
mesh_type="box",
):
super().__init__()
self.batch_size = batch_size
if mesh_type == "box":
box = tricreation.box([1, 1, 1])
faces = torch.Tensor(np.array(box.faces))
verts_loc = torch.Tensor(np.array(box.vertices))
elif mesh_type == "sphere":
icomesh = py3dutils.ico_sphere(2)
verts_loc = icomesh.verts_padded()[0]
faces = icomesh.faces_padded()[0]
else:
raise ValueError(f"{mesh_type} not in [sphere|box]")
# Normalize and save verts
norm_verts = normalize.normalize_verts(verts_loc, 1)
self.register_buffer(
"verts", norm_verts.unsqueeze(0).repeat(batch_size, 1, 1)
)
self.register_buffer(
"faces", faces.unsqueeze(0).repeat(batch_size, 1, 1)
)
# Initialize translation and rotation
rot_vecs = torch.stack(
[rotutils.get_rotvec6d(random_rot) for _ in range(batch_size)]
)
# Initialize rotation parameters
self.rot_vecs = torch.nn.Parameter(rot_vecs, requires_grad=True)
# Initialize scale and translation
self.trans = torch.nn.Parameter(
norm_verts.new_zeros(batch_size, 3)
+ norm_verts.new([x_off, y_off, z_off]).view(1, 3),
requires_grad=True,
)
# Scale is shared for all object views
self.scale = torch.nn.Parameter(
torch.Tensor([radius]), requires_grad=True
)
def get_initial_sphere_meshes(level, device):
"""
wrapper around pytorch3d.utils.ico_sphere
returns mesh at proper position/scale
"""
init_meshes = ico_sphere(level, device)
# TODO: Don't use these magic numbers
init_meshes.scale_verts_(0.25)
offset = init_meshes.verts_packed().clone()
offset[:, :] = 0
# average depth of objects is 1.4m
# z is negative because of coordinate from convention
offset[:, 2] = -1.4
init_meshes.offset_verts_(offset)
return init_meshes
def test_backward(self):
device = torch.device("cuda:0")
mesh = ico_sphere()
verts = mesh.verts_packed()
edges = mesh.edges_packed()
verts_cuda = verts.clone().to(device)
edges_cuda = edges.clone().to(device)
verts.requires_grad = True
verts_cuda.requires_grad = True
neighbor_sums_cuda = gather_scatter(verts_cuda, edges_cuda, False)
neighbor_sums = gather_scatter_python(verts, edges, False)
neighbor_sums_cuda.sum().backward()
neighbor_sums.sum().backward()
self.assertClose(verts.grad.cpu(), verts_cuda.grad.cpu())
def forward(self, imgs):
N = imgs.shape[0]
device = imgs.device
img_feats = self.backbone(imgs)
# add view dimension (single view)
img_feats = [i.unsqueeze(1) for i in img_feats]
P = [self._get_projection_matrix(N, device)]
init_meshes = ico_sphere(self.ico_sphere_level, device).extend(N)
refined_meshes = self.mesh_head(img_feats, init_meshes, P)
return {
"voxel_scores": None,
"meshes_pred": refined_meshes,
}
def create_data(folder_path='meshes/',
nb_of_pointclouds=50,
nb_of_points=5000,
sphere_level=4,
normalize_data=True):
device = torch.device("cuda:0")
data_path = os.path.join(os.getcwd(), folder_path)
src_mesh = ico_sphere(sphere_level, device)
for filename in os.listdir(data_path):
print(f"{datetime.now()} Starting:{filename}")
file_path = os.path.join(data_path, filename)
cur_mesh = utils.load_mesh(file_path)
cur_deform_verts = deformation.get_deform_verts(
cur_mesh, nb_of_points, sphere_level)
data_verts = np.expand_dims(cur_deform_verts.detach().cpu().numpy(),
axis=0)
data_input = None
data_output = None
for _ in range(nb_of_pointclouds):
data_a = sample_points_from_meshes(
cur_mesh, nb_of_points).squeeze().cpu().numpy()
if normalize_data:
data_a = data_a - np.mean(data_a, axis=0)
data_a = data_a / np.max(data_a, axis=0)
data_a_sort_indices = np.argsort(np.linalg.norm(data_a,
axis=1))
data_a = data_a[data_a_sort_indices]
data_a = np.expand_dims(data_a, axis=0)
data_input = data_a if data_input is None else np.concatenate(
(data_input, data_a))
data_output = data_verts if data_output is None else np.concatenate(
(data_output, data_verts))
np.save(f'data/{os.path.splitext(filename)[0]}_input.npy', data_input)
np.save(f'data/{os.path.splitext(filename)[0]}_output.npy',
data_output)
deformed_mesh = src_mesh.offset_verts(cur_deform_verts)
final_verts, final_faces = deformed_mesh.get_mesh_verts_faces(0)
final_obj = os.path.join(
'deformed_meshes/',
f'{os.path.splitext(filename)[0]}_deformed.obj')
save_obj(final_obj, final_verts, final_faces)
print(
f"{datetime.now()} Finished:{filename}, Point Cloud Shape:{data_input.shape} Deform Verts Shape:{data_output.shape}"
)
def forward(self,
imgs,
z=None): # z is the latent vector sampled from P(z|x)
N = imgs.shape[0]
device = imgs.device
img_feats = self.backbone(imgs)
# concat_feats = torch.cat((img_feats,z),dim=1)
P = self._get_projection_matrix(N, device)
# print(P)
init_meshes = ico_sphere(self.ico_sphere_level, device).extend(N)
refined_meshes = self.mesh_head(img_feats,
z,
init_meshes,
P,
subdivide=True)
return None, refined_meshes
def get_deform_verts(target_mesh, points_to_sample=5000, sphere_level=4):
device = torch.device("cuda:0")
src_mesh = ico_sphere(sphere_level, device)
deform_verts = torch.full(src_mesh.verts_packed().shape,
0.0,
device=device,
requires_grad=True)
learning_rate = 0.01
num_iter = 500
w_chamfer = 1.0
w_edge = 0.05
w_normal = 0.0005
w_laplacian = 0.005
optimizer = torch.optim.Adam([deform_verts],
lr=learning_rate,
betas=(0.5, 0.999))
for _ in range(num_iter):
optimizer.zero_grad()
new_src_mesh = src_mesh.offset_verts(deform_verts)
sample_trg = sample_points_from_meshes(target_mesh, points_to_sample)
sample_src = sample_points_from_meshes(new_src_mesh, points_to_sample)
loss_chamfer, _ = chamfer_distance(sample_trg, sample_src)
loss_edge = mesh_edge_loss(new_src_mesh)
loss_normal = mesh_normal_consistency(new_src_mesh)
loss_laplacian = mesh_laplacian_smoothing(new_src_mesh,
method="uniform")
loss = loss_chamfer * w_chamfer + loss_edge * w_edge + loss_normal * w_normal + loss_laplacian * w_laplacian
loss.backward()
optimizer.step()
print(
f"{datetime.now()} Loss Chamfer:{loss_chamfer * w_chamfer}, Loss Edge:{loss_edge * w_edge}, Loss Normal:{loss_normal * w_normal}, Loss Laplacian:{loss_laplacian * w_laplacian}"
)
return deform_verts
def test_taubin(self):
N = 3
device = get_random_cuda_device()
mesh = ico_sphere(4, device).extend(N)
ico_verts = mesh.verts_padded()
ico_faces = mesh.faces_padded()
rand_noise = torch.rand_like(ico_verts) * 0.2 - 0.1
z_mask = (ico_verts[:, :, -1] > 0).view(N, -1, 1)
rand_noise = rand_noise * z_mask
verts = ico_verts + rand_noise
mesh = Meshes(verts=verts, faces=ico_faces)
smooth_mesh = taubin_smoothing(mesh, num_iter=50)
smooth_verts = smooth_mesh.verts_padded()
smooth_dist = (smooth_verts - ico_verts).norm(dim=-1).mean()
dist = (verts - ico_verts).norm(dim=-1).mean()
self.assertTrue(smooth_dist < dist)
def test_backward(self):
device = get_random_cuda_device()
mesh = ico_sphere()
verts = mesh.verts_packed()
edges = mesh.edges_packed()
verts_cpu = verts.clone()
edges_cpu = edges.clone()
verts_cuda = verts.clone().to(device)
edges_cuda = edges.clone().to(device)
verts.requires_grad = True
verts_cpu.requires_grad = True
verts_cuda.requires_grad = True
neighbor_sums_cuda = gather_scatter(verts_cuda, edges_cuda, False)
neighbor_sums_cpu = gather_scatter(verts_cpu, edges_cpu, False)
neighbor_sums = gather_scatter_python(verts, edges, False)
randoms = torch.rand_like(neighbor_sums)
(neighbor_sums_cuda * randoms.to(device)).sum().backward()
(neighbor_sums_cpu * randoms).sum().backward()
(neighbor_sums * randoms).sum().backward()
self.assertClose(verts.grad, verts_cuda.grad.cpu())
self.assertClose(verts.grad, verts_cpu.grad)
def __init__(self,
batch_size=100,
latent_size=128,
img_size=128,
seed_sphere_divisions=3):
super(Generator, self).__init__()
src_mesh = ico_sphere(seed_sphere_divisions)
output_shape = src_mesh.verts_packed().shape
self.src_meshes = src_mesh.extend(batch_size).cuda()
self.layers = LinearNet(latent_size, output_shape).cuda()
""" Setup rendering. """
num_views = batch_size
self.lights = PointLights(location=[[0.0, 0.0, -3.0]], device=device)
sigma = 1e-4
raster_settings_silhouette = RasterizationSettings(
image_size=128,
blur_radius=np.log(1. / 1e-4 - 1.) * sigma,
faces_per_pixel=50,
)
self.renderer_silhouette = MeshRenderer(
rasterizer=MeshRasterizer(
cameras=None, raster_settings=raster_settings_silhouette),
shader=SoftSilhouetteShader()).cuda()
faces_idx = faces.verts_idx.to(device)
verts = verts.to(device)
# We scale normalize and center the target mesh to fit in a sphere of radius 1 centered at (0,0,0).
# (scale, center) will be used to bring the predicted mesh to its original center and scale
# Note that normalizing the target mesh, speeds up the optimization but is not necessary!
center = verts.mean(0)
verts = verts - center
scale = max(verts.abs().max(0)[0])
verts = verts / scale
# We construct a Meshes structure for the target mesh
trg_mesh = Meshes(verts=[verts], faces=[faces_idx])
# We initialize the source shape to be a sphere of radius 1
block1 = ico_sphere(2, device)
uppooling1 = SubdivideMeshes(block1)
block2 = uppooling1(block1)
uppooling2 = SubdivideMeshes(block2)
block3 = uppooling2(block2)
#TODO batch size, features
data1 = Data(x=block1.verts_packed().reshape(-1, 3),
edge_index=block1.edges_packed().reshape(2, -1),
face=block1.faces_packed().reshape(3, -1))
data1.pos = data1.x
data2 = Data(x=block2.verts_packed().reshape(-1, 3),
edge_index=block2.edges_packed().reshape(2, -1),
face=block2.faces_packed().reshape(3, -1))
faces_idx = faces.verts_idx.to(device)
verts = verts.to(device)
# We scale normalize and center the target mesh to fit in a sphere of radius 1 centered at (0,0,0).
# (scale, center) will be used to bring the predicted mesh to its original center and scale
# Note that normalizing the target mesh, speeds up the optimization but is not necessary!
center = verts.mean(0)
verts = verts - center
scale = max(verts.abs().max(0)[0])
verts = verts / scale
# We construct a Meshes structure for the target mesh
trg_mesh = Meshes(verts=[verts], faces=[faces_idx])
# We initialize the source shape to be a sphere of radius 1
src_mesh = ico_sphere(4, device)
### Visualize the source and target meshes
def plot_pointcloud(mesh, title=""):
# Sample points uniformly from the surface of the mesh.
points = sample_points_from_meshes(mesh, 5000)
x, y, z = points.clone().detach().cpu().squeeze().unbind(1)
fig = plt.figure(figsize=(5, 5))
ax = Axes3D(fig)
ax.scatter3D(x, z, -y)
ax.set_xlabel('x')
ax.set_ylabel('z')
ax.set_zlabel('y')
ax.set_title(title)
def test_python_vs_cpu_vs_cuda(self):
torch.manual_seed(231)
device = torch.device("cpu")
image_size = 32
blur_radius = 0.1**2
faces_per_pixel = 3
for d in ["cpu", "cuda"]:
device = torch.device(d)
compare_grads = True
# Mesh with a single face.
verts1 = torch.tensor(
[[0.0, 0.6, 0.1], [-0.7, -0.4, 0.5], [0.7, -0.4, 0.7]],
dtype=torch.float32,
requires_grad=True,
device=device,
)
faces1 = torch.tensor([[0, 1, 2]],
dtype=torch.int64,
device=device)
meshes1 = Meshes(verts=[verts1], faces=[faces1])
args1 = (meshes1, image_size, blur_radius, faces_per_pixel)
verts2 = verts1.detach().clone()
verts2.requires_grad = True
meshes2 = Meshes(verts=[verts2], faces=[faces1])
args2 = (meshes2, image_size, blur_radius, faces_per_pixel)
self._compare_impls(
rasterize_meshes_python,
rasterize_meshes,
args1,
args2,
verts1,
verts2,
compare_grads=compare_grads,
)
# Mesh with multiple faces.
# fmt: off
verts1 = torch.tensor(
[
[-0.5, 0.0, 0.1], # noqa: E241, E201
[0.0, 0.6, 0.5], # noqa: E241, E201
[0.5, 0.0, 0.7], # noqa: E241, E201
[-0.25, 0.0, 0.9], # noqa: E241, E201
[0.26, 0.5, 0.8], # noqa: E241, E201
[0.76, 0.0, 0.8], # noqa: E241, E201
[-0.41, 0.0, 0.5], # noqa: E241, E201
[0.61, 0.6, 0.6], # noqa: E241, E201
[0.41, 0.0, 0.5], # noqa: E241, E201
[-0.2, 0.0, -0.5], # noqa: E241, E201
[0.3, 0.6, -0.5], # noqa: E241, E201
[0.4, 0.0, -0.5], # noqa: E241, E201
],
dtype=torch.float32,
device=device,
requires_grad=True)
faces1 = torch.tensor(
[
[1, 0, 2], # noqa: E241, E201
[4, 3, 5], # noqa: E241, E201
[7, 6, 8], # noqa: E241, E201
[10, 9, 11] # noqa: E241, E201
],
dtype=torch.int64,
device=device,
)
# fmt: on
meshes = Meshes(verts=[verts1], faces=[faces1])
args1 = (meshes, image_size, blur_radius, faces_per_pixel)
verts2 = verts1.clone().detach()
verts2.requires_grad = True
meshes2 = Meshes(verts=[verts2], faces=[faces1])
args2 = (meshes2, image_size, blur_radius, faces_per_pixel)
self._compare_impls(
rasterize_meshes_python,
rasterize_meshes,
args1,
args2,
verts1,
verts2,
compare_grads=compare_grads,
)
# Icosphere
meshes = ico_sphere(device=device)
verts1, faces1 = meshes.get_mesh_verts_faces(0)
verts1.requires_grad = True
meshes = Meshes(verts=[verts1], faces=[faces1])
args1 = (meshes, image_size, blur_radius, faces_per_pixel)
verts2 = verts1.detach().clone()
verts2.requires_grad = True
meshes2 = Meshes(verts=[verts2], faces=[faces1])
args2 = (meshes2, image_size, blur_radius, faces_per_pixel)
self._compare_impls(
rasterize_meshes_python,
rasterize_meshes,
args1,
args2,
verts1,
verts2,
compare_grads=compare_grads,
)
def test_cpp_vs_cuda_naive_vs_cuda_binned(self):
# Make sure that the backward pass runs for all pathways
image_size = 64 # test is too slow for very large images.
N = 1
radius = 0.1**2
faces_per_pixel = 3
grad_zbuf = torch.randn(N, image_size, image_size, faces_per_pixel)
grad_dist = torch.randn(N, image_size, image_size, faces_per_pixel)
grad_bary = torch.randn(N, image_size, image_size, faces_per_pixel, 3)
device = torch.device("cpu")
meshes = ico_sphere(0, device)
verts, faces = meshes.get_mesh_verts_faces(0)
verts.requires_grad = True
meshes = Meshes(verts=[verts], faces=[faces])
# Option I: CPU, naive
args = (meshes, image_size, radius, faces_per_pixel)
idx1, zbuf1, bary1, dist1 = rasterize_meshes(*args)
loss = ((zbuf1 * grad_zbuf).sum() + (dist1 * grad_dist).sum() +
(bary1 * grad_bary).sum())
loss.backward()
idx1 = idx1.data.cpu().clone()
zbuf1 = zbuf1.data.cpu().clone()
dist1 = dist1.data.cpu().clone()
grad1 = verts.grad.data.cpu().clone()
# Option II: CUDA, naive
device = torch.device("cuda:0")
meshes = ico_sphere(0, device)
verts, faces = meshes.get_mesh_verts_faces(0)
verts.requires_grad = True
meshes = Meshes(verts=[verts], faces=[faces])
args = (meshes, image_size, radius, faces_per_pixel, 0, 0)
idx2, zbuf2, bary2, dist2 = rasterize_meshes(*args)
grad_zbuf = grad_zbuf.cuda()
grad_dist = grad_dist.cuda()
grad_bary = grad_bary.cuda()
loss = ((zbuf2 * grad_zbuf).sum() + (dist2 * grad_dist).sum() +
(bary2 * grad_bary).sum())
loss.backward()
idx2 = idx2.data.cpu().clone()
zbuf2 = zbuf2.data.cpu().clone()
dist2 = dist2.data.cpu().clone()
grad2 = verts.grad.data.cpu().clone()
# Option III: CUDA, binned
device = torch.device("cuda:0")
meshes = ico_sphere(0, device)
verts, faces = meshes.get_mesh_verts_faces(0)
verts.requires_grad = True
meshes = Meshes(verts=[verts], faces=[faces])
args = (meshes, image_size, radius, faces_per_pixel, 32, 500)
idx3, zbuf3, bary3, dist3 = rasterize_meshes(*args)
loss = ((zbuf3 * grad_zbuf).sum() + (dist3 * grad_dist).sum() +
(bary3 * grad_bary).sum())
loss.backward()
idx3 = idx3.data.cpu().clone()
zbuf3 = zbuf3.data.cpu().clone()
dist3 = dist3.data.cpu().clone()
grad3 = verts.grad.data.cpu().clone()
# Make sure everything was the same
self.assertTrue((idx1 == idx2).all().item())
self.assertTrue((idx1 == idx3).all().item())
self.assertTrue(torch.allclose(zbuf1, zbuf2, atol=1e-6))
self.assertTrue(torch.allclose(zbuf1, zbuf3, atol=1e-6))
self.assertTrue(torch.allclose(dist1, dist2, atol=1e-6))
self.assertTrue(torch.allclose(dist1, dist3, atol=1e-6))
self.assertTrue(torch.allclose(grad1, grad2, rtol=5e-3)) # flaky test
self.assertTrue(torch.allclose(grad1, grad3, rtol=5e-3))
self.assertTrue(torch.allclose(grad2, grad3, rtol=5e-3))
def _compare_with_meshes_renderer(self,
image_size,
batch_size=11,
sphere_diameter=0.6):
"""
Generate a spherical RGB volumetric function and its corresponding mesh
and check whether MeshesRenderer returns the same images as the
corresponding ImplicitRenderer.
"""
# generate NDC camera extrinsics and intrinsics
cameras = init_cameras(batch_size, image_size=image_size, ndc=True)
# get rand offset of the volume
sphere_centroid = torch.randn(batch_size, 3,
device=cameras.device) * 0.1
sphere_centroid.requires_grad = True
# init the grid raysampler with the ndc grid
raysampler = NDCMultinomialRaysampler(
image_width=image_size[1],
image_height=image_size[0],
n_pts_per_ray=256,
min_depth=0.1,
max_depth=2.0,
)
# get the EA raymarcher
raymarcher = EmissionAbsorptionRaymarcher()
# jitter the camera intrinsics a bit for each render
cameras_randomized = cameras.clone()
cameras_randomized.principal_point = (
torch.randn_like(cameras.principal_point) * 0.3)
cameras_randomized.focal_length = (
cameras.focal_length +
torch.randn_like(cameras.focal_length) * 0.2)
# the list of differentiable camera vars
cam_vars = ("R", "T", "focal_length", "principal_point")
# enable the gradient caching for the camera variables
for cam_var in cam_vars:
getattr(cameras_randomized, cam_var).requires_grad = True
# get the implicit renderer
images_opacities = ImplicitRenderer(
raysampler=raysampler, raymarcher=raymarcher)(
cameras=cameras_randomized,
volumetric_function=spherical_volumetric_function,
sphere_centroid=sphere_centroid,
sphere_diameter=sphere_diameter,
)[0]
# check that the renderer does not erase gradients
loss = images_opacities.sum()
loss.backward()
for check_var in (
*[
getattr(cameras_randomized, cam_var)
for cam_var in cam_vars
],
sphere_centroid,
):
self.assertIsNotNone(check_var.grad)
# instantiate the corresponding spherical mesh
ico = ico_sphere(level=4, device=cameras.device).extend(batch_size)
verts = (torch.nn.functional.normalize(ico.verts_padded(), dim=-1) *
sphere_diameter + sphere_centroid[:, None])
meshes = Meshes(
verts=verts,
faces=ico.faces_padded(),
textures=TexturesVertex(verts_features=(
torch.nn.functional.normalize(verts, dim=-1) * 0.5 + 0.5)),
)
# instantiate the corresponding mesh renderer
lights = PointLights(device=cameras.device, location=[[0.0, 0.0, 0.0]])
renderer_textured = MeshRenderer(
rasterizer=MeshRasterizer(
cameras=cameras_randomized,
raster_settings=RasterizationSettings(
image_size=image_size,
blur_radius=1e-3,
faces_per_pixel=10,
z_clip_value=None,
perspective_correct=False,
),
),
shader=SoftPhongShader(
device=cameras.device,
cameras=cameras_randomized,
lights=lights,
materials=Materials(
ambient_color=((2.0, 2.0, 2.0), ),
diffuse_color=((0.0, 0.0, 0.0), ),
specular_color=((0.0, 0.0, 0.0), ),
shininess=64,
device=cameras.device,
),
blend_params=BlendParams(sigma=1e-3,
gamma=1e-4,
background_color=(0.0, 0.0, 0.0)),
),
)
# get the mesh render
images_opacities_meshes = renderer_textured(meshes,
cameras=cameras_randomized,
lights=lights)
if DEBUG:
outdir = tempfile.gettempdir() + "/test_implicit_vs_mesh_renderer"
os.makedirs(outdir, exist_ok=True)
frames = []
for (image_opacity,
image_opacity_mesh) in zip(images_opacities,
images_opacities_meshes):
image, opacity = image_opacity.split([3, 1], dim=-1)
image_mesh, opacity_mesh = image_opacity_mesh.split([3, 1],
dim=-1)
diff_image = (((image - image_mesh) * 0.5 + 0.5).mean(
dim=2, keepdim=True).repeat(1, 1, 3))
image_pil = Image.fromarray((torch.cat(
(
image,
image_mesh,
diff_image,
opacity.repeat(1, 1, 3),
opacity_mesh.repeat(1, 1, 3),
),
dim=1,
).detach().cpu().numpy() * 255.0).astype(np.uint8))
frames.append(image_pil)
# export gif
outfile = os.path.join(outdir, "implicit_vs_mesh_render.gif")
frames[0].save(
outfile,
save_all=True,
append_images=frames[1:],
duration=batch_size // 15,
loop=0,
)
print(f"exported {outfile}")
# export concatenated frames
outfile_cat = os.path.join(outdir, "implicit_vs_mesh_render.png")
Image.fromarray(
np.concatenate([np.array(f) for f in frames],
axis=0)).save(outfile_cat)
print(f"exported {outfile_cat}")
# compare the renders
diff = (images_opacities - images_opacities_meshes).abs().mean(dim=-1)
mu_diff = diff.mean(dim=(1, 2))
std_diff = diff.std(dim=(1, 2))
self.assertClose(mu_diff, torch.zeros_like(mu_diff), atol=5e-2)
self.assertClose(std_diff, torch.zeros_like(std_diff), atol=6e-2)
def _forward_shape(self, features, instances):
"""
Forward logic for the voxel and mesh refinement branch.
Args:
features (list[Tensor]): #level input features for voxel prediction
instances (list[Instances]): the per-image instances to train/predict meshes.
In training, they can be the proposals.
In inference, they can be the predicted boxes.
Returns:
In training, a dict of losses.
In inference, update `instances` with new fields "pred_voxels" & "pred_meshes" and return it.
"""
if not self.voxel_on and not self.mesh_on:
return {} if self.training else instances
if self.training:
# The loss is only defined on positive proposals.
proposals, _ = select_foreground_proposals(instances,
self.num_classes)
proposal_boxes = [x.proposal_boxes for x in proposals]
losses = {}
if self.voxel_on:
voxel_features = self.voxel_pooler(features, proposal_boxes)
voxel_logits = self.voxel_head(voxel_features)
loss_voxel, target_voxels = voxel_rcnn_loss(
voxel_logits,
proposals,
loss_weight=self.voxel_loss_weight)
losses.update({"loss_voxel": loss_voxel})
if self._vis:
self._misc["target_voxels"] = target_voxels
if self.cls_agnostic_voxel:
with torch.no_grad():
vox_in = voxel_logits.sigmoid().squeeze(
1) # (N, V, V, V)
init_mesh = cubify(vox_in, self.cubify_thresh) # 1
else:
raise ValueError(
"No support for class specific predictions")
if self.mesh_on:
mesh_features = self.mesh_pooler(features, proposal_boxes)
if not self.voxel_on:
if mesh_features.shape[0] > 0:
init_mesh = ico_sphere(self.ico_sphere_level,
mesh_features.device)
init_mesh = init_mesh.extend(mesh_features.shape[0])
else:
init_mesh = Meshes(verts=[], faces=[])
pred_meshes = self.mesh_head(mesh_features, init_mesh)
# loss weights
loss_weights = {
"chamfer": self.chamfer_loss_weight,
"normals": self.normals_loss_weight,
"edge": self.edge_loss_weight,
}
if not pred_meshes[0].isempty():
loss_chamfer, loss_normals, loss_edge, target_meshes = mesh_rcnn_loss(
pred_meshes,
proposals,
loss_weights=loss_weights,
gt_num_samples=self.gt_num_samples,
pred_num_samples=self.pred_num_samples,
gt_coord_thresh=self.gt_coord_thresh,
)
if self._vis:
self._misc["init_meshes"] = init_mesh
self._misc["target_meshes"] = target_meshes
else:
loss_chamfer = sum(
k.sum() for k in self.mesh_head.parameters()) * 0.0
loss_normals = sum(
k.sum() for k in self.mesh_head.parameters()) * 0.0
loss_edge = sum(k.sum()
for k in self.mesh_head.parameters()) * 0.0
losses.update({
"loss_chamfer": loss_chamfer,
"loss_normals": loss_normals,
"loss_edge": loss_edge,
})
return losses
else:
pred_boxes = [x.pred_boxes for x in instances]
if self.voxel_on:
voxel_features = self.voxel_pooler(features, pred_boxes)
voxel_logits = self.voxel_head(voxel_features)
voxel_rcnn_inference(voxel_logits, instances)
if self.cls_agnostic_voxel:
with torch.no_grad():
vox_in = voxel_logits.sigmoid().squeeze(
1) # (N, V, V, V)
init_mesh = cubify(vox_in, self.cubify_thresh) # 1
else:
raise ValueError(
"No support for class specific predictions")
if self.mesh_on:
mesh_features = self.mesh_pooler(features, pred_boxes)
if not self.voxel_on:
if mesh_features.shape[0] > 0:
init_mesh = ico_sphere(self.ico_sphere_level,
mesh_features.device)
init_mesh = init_mesh.extend(mesh_features.shape[0])
else:
init_mesh = Meshes(verts=[], faces=[])
pred_meshes = self.mesh_head(mesh_features, init_mesh)
mesh_rcnn_inference(pred_meshes[-1], instances)
else:
assert self.voxel_on
mesh_rcnn_inference(init_mesh, instances)
return instances
def test_save_load_icosphere(self):
# Test that saving a mesh as an off file and loading it results in the
# same data on the correct device, for all permitted types of textures.
# Standard test is for random colors, but also check totally white,
# because there's a different in OFF semantics between "1.0" color (=full)
# and "1" (= 1/255 color)
sphere = ico_sphere(0)
io = IO()
device = torch.device("cuda:0")
atlas_padded = torch.rand(1, sphere.faces_list()[0].shape[0], 1, 1, 3)
atlas = TexturesAtlas(atlas_padded)
atlas_padded_white = torch.ones(1,
sphere.faces_list()[0].shape[0], 1, 1,
3)
atlas_white = TexturesAtlas(atlas_padded_white)
verts_colors_padded = torch.rand(1, sphere.verts_list()[0].shape[0], 3)
vertex_texture = TexturesVertex(verts_colors_padded)
verts_colors_padded_white = torch.ones(1,
sphere.verts_list()[0].shape[0],
3)
vertex_texture_white = TexturesVertex(verts_colors_padded_white)
# No colors case
with NamedTemporaryFile(mode="w", suffix=".off") as f:
io.save_mesh(sphere, f.name)
f.flush()
mesh1 = io.load_mesh(f.name, device=device)
self.assertEqual(mesh1.device, device)
mesh1 = mesh1.cpu()
self.assertClose(mesh1.verts_padded(), sphere.verts_padded())
self.assertClose(mesh1.faces_padded(), sphere.faces_padded())
self.assertIsNone(mesh1.textures)
# Atlas case
sphere.textures = atlas
with NamedTemporaryFile(mode="w", suffix=".off") as f:
io.save_mesh(sphere, f.name)
f.flush()
mesh2 = io.load_mesh(f.name, device=device)
self.assertEqual(mesh2.device, device)
mesh2 = mesh2.cpu()
self.assertClose(mesh2.verts_padded(), sphere.verts_padded())
self.assertClose(mesh2.faces_padded(), sphere.faces_padded())
self.assertClose(mesh2.textures.atlas_padded(),
atlas_padded,
atol=1e-4)
# White atlas case
sphere.textures = atlas_white
with NamedTemporaryFile(mode="w", suffix=".off") as f:
io.save_mesh(sphere, f.name)
f.flush()
mesh3 = io.load_mesh(f.name)
self.assertClose(mesh3.textures.atlas_padded(),
atlas_padded_white,
atol=1e-4)
# TexturesVertex case
sphere.textures = vertex_texture
with NamedTemporaryFile(mode="w", suffix=".off") as f:
io.save_mesh(sphere, f.name)
f.flush()
mesh4 = io.load_mesh(f.name, device=device)
self.assertEqual(mesh4.device, device)
mesh4 = mesh4.cpu()
self.assertClose(mesh4.verts_padded(), sphere.verts_padded())
self.assertClose(mesh4.faces_padded(), sphere.faces_padded())
self.assertClose(mesh4.textures.verts_features_padded(),
verts_colors_padded,
atol=1e-4)
# white TexturesVertex case
sphere.textures = vertex_texture_white
with NamedTemporaryFile(mode="w", suffix=".off") as f:
io.save_mesh(sphere, f.name)
f.flush()
mesh5 = io.load_mesh(f.name)
self.assertClose(mesh5.textures.verts_features_padded(),
verts_colors_padded_white,
atol=1e-4)
