mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2026-02-06 22:12:16 +08:00
Update latest version of site
This commit is contained in:
Nikhila Ravi
File diff suppressed because one or more lines are too long
@@ -1,219 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
|
||||
# In[1]:
|
||||
|
||||
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
|
||||
|
||||
# # Absolute camera orientation given set of relative camera pairs
|
||||
#
|
||||
# This tutorial showcases the `cameras`, `transforms` and `so3` API.
|
||||
#
|
||||
# The problem we deal with is defined as follows:
|
||||
#
|
||||
# Given an optical system of $N$ cameras with extrinsics $\{g_1, ..., g_N | g_i \in SE(3)\}$, and a set of relative camera positions $\{g_{ij} | g_{ij}\in SE(3)\}$ that map between coordinate frames of randomly selected pairs of cameras $(i, j)$, we search for the absolute extrinsic parameters $\{g_1, ..., g_N\}$ that are consistent with the relative camera motions.
|
||||
#
|
||||
# More formally:
|
||||
# $$
|
||||
# g_1, ..., g_N =
|
||||
# {\arg \min}_{g_1, ..., g_N} \sum_{g_{ij}} d(g_{ij}, g_i^{-1} g_j),
|
||||
# $$,
|
||||
# where $d(g_i, g_j)$ is a suitable metric that compares the extrinsics of cameras $g_i$ and $g_j$.
|
||||
#
|
||||
# Visually, the problem can be described as follows. The picture below depicts the situation at the beginning of our optimization. The ground truth cameras are plotted in green while the randomly initialized estimated cameras are plotted in blue:
|
||||
# 
|
||||
#
|
||||
# Our optimization seeks to align the estimated (blue) cameras with the ground truth (green) cameras, by minimizing the discrepancies between pairs of relative cameras. Thus, the solution to the problem should look as follows:
|
||||
# 
|
||||
#
|
||||
# In practice, the camera extrinsics $g_{ij}$ and $g_i$ are represented using objects from the `SfMPerspectiveCameras` class initialized with the corresponding rotation and translation matrices `R_absolute` and `T_absolute` that define the extrinsic parameters $g = (R, T); R \in SO(3); T \in \mathbb{R}^3$. In order to ensure that `R_absolute` is a valid rotation matrix, we represent it using an exponential map (implemented with `so3_exponential_map`) of the axis-angle representation of the rotation `log_R_absolute`.
|
||||
#
|
||||
# Note that the solution to this problem could only be recovered up to an unknown global rigid transformation $g_{glob} \in SE(3)$. Thus, for simplicity, we assume knowledge of the absolute extrinsics of the first camera $g_0$. We set $g_0$ as a trivial camera $g_0 = (I, \vec{0})$.
|
||||
#
|
||||
|
||||
# ## 0. Import Modules
|
||||
|
||||
# In[2]:
|
||||
|
||||
|
||||
# imports
|
||||
import torch
|
||||
from pytorch3d.transforms.so3 import (
|
||||
so3_exponential_map,
|
||||
so3_relative_angle,
|
||||
)
|
||||
from pytorch3d.renderer.cameras import (
|
||||
SfMPerspectiveCameras,
|
||||
)
|
||||
|
||||
# add path for demo utils
|
||||
import sys
|
||||
import os
|
||||
sys.path.append(os.path.abspath(''))
|
||||
from utils import plot_camera_scene
|
||||
|
||||
# set for reproducibility
|
||||
torch.manual_seed(42)
|
||||
|
||||
|
||||
# ## 1. Set up Cameras and load ground truth positions
|
||||
|
||||
# In[3]:
|
||||
|
||||
|
||||
# load the SE3 graph of relative/absolute camera positions
|
||||
camera_graph_file = './data/camera_graph.pth'
|
||||
(R_absolute_gt, T_absolute_gt), (R_relative, T_relative), relative_edges = torch.load(camera_graph_file)
|
||||
|
||||
# create the relative cameras
|
||||
cameras_relative = SfMPerspectiveCameras(
|
||||
R = R_relative.cuda(),
|
||||
T = T_relative.cuda(),
|
||||
device = "cuda",
|
||||
)
|
||||
|
||||
# create the absolute ground truth cameras
|
||||
cameras_absolute_gt = SfMPerspectiveCameras(
|
||||
R = R_absolute_gt.cuda(),
|
||||
T = T_absolute_gt.cuda(),
|
||||
device = "cuda",
|
||||
)
|
||||
|
||||
# the number of absolute camera positions
|
||||
N = R_absolute_gt.shape[0]
|
||||
|
||||
|
||||
# ## 2. Define optimization functions
|
||||
#
|
||||
# ### Relative cameras and camera distance
|
||||
# We now define two functions crucial for the optimization.
|
||||
#
|
||||
# **`calc_camera_distance`** compares a pair of cameras. This function is important as it defines the loss that we are minimizing. The method utilizes the `so3_relative_angle` function from the SO3 API.
|
||||
#
|
||||
# **`get_relative_camera`** computes the parameters of a relative camera that maps between a pair of absolute cameras. Here we utilize the `compose` and `inverse` class methods from the PyTorch3d Transforms API.
|
||||
|
||||
# In[ ]:
|
||||
|
||||
|
||||
def calc_camera_distance(cam_1, cam_2):
|
||||
"""
|
||||
Calculates the divergence of a batch of pairs of cameras cam_1, cam_2.
|
||||
The distance is composed of the cosine of the relative angle between
|
||||
the rotation components of the camera extrinsics and the l2 distance
|
||||
between the translation vectors.
|
||||
"""
|
||||
# rotation distance
|
||||
R_distance = (1.-so3_relative_angle(cam_1.R, cam_2.R, cos_angle=True)).mean()
|
||||
# translation distance
|
||||
T_distance = ((cam_1.T - cam_2.T)**2).sum(1).mean()
|
||||
# the final distance is the sum
|
||||
return R_distance + T_distance
|
||||
|
||||
def get_relative_camera(cams, edges):
|
||||
"""
|
||||
For each pair of indices (i,j) in "edges" generate a camera
|
||||
that maps from the coordinates of the camera cams[i] to
|
||||
the coordinates of the camera cams[j]
|
||||
"""
|
||||
|
||||
# first generate the world-to-view Transform3d objects of each
|
||||
# camera pair (i, j) according to the edges argument
|
||||
trans_i, trans_j = [
|
||||
SfMPerspectiveCameras(
|
||||
R = cams.R[edges[:, i]],
|
||||
T = cams.T[edges[:, i]],
|
||||
device = "cuda",
|
||||
).get_world_to_view_transform()
|
||||
for i in (0, 1)
|
||||
]
|
||||
|
||||
# compose the relative transformation as g_i^{-1} g_j
|
||||
trans_rel = trans_i.inverse().compose(trans_j)
|
||||
|
||||
# generate a camera from the relative transform
|
||||
matrix_rel = trans_rel.get_matrix()
|
||||
cams_relative = SfMPerspectiveCameras(
|
||||
R = matrix_rel[:, :3, :3],
|
||||
T = matrix_rel[:, 3, :3],
|
||||
device = "cuda",
|
||||
)
|
||||
return cams_relative
|
||||
|
||||
|
||||
# ## 3. Optimization
|
||||
# Finally, we start the optimization of the absolute cameras.
|
||||
#
|
||||
# We use SGD with momentum and optimize over `log_R_absolute` and `T_absolute`.
|
||||
#
|
||||
# As mentioned earlier, `log_R_absolute` is the axis angle representation of the rotation part of our absolute cameras. We can obtain the 3x3 rotation matrix `R_absolute` that corresponds to `log_R_absolute` with:
|
||||
#
|
||||
# `R_absolute = so3_exponential_map(log_R_absolute)`
|
||||
#
|
||||
|
||||
# In[ ]:
|
||||
|
||||
|
||||
# initialize the absolute log-rotations/translations with random entries
|
||||
log_R_absolute_init = torch.randn(N, 3).float().cuda()
|
||||
T_absolute_init = torch.randn(N, 3).float().cuda()
|
||||
|
||||
# futhermore, we know that the first camera is a trivial one
|
||||
# (see the description above)
|
||||
log_R_absolute_init[0, :] = 0.
|
||||
T_absolute_init[0, :] = 0.
|
||||
|
||||
# instantiate a copy of the initialization of log_R / T
|
||||
log_R_absolute = log_R_absolute_init.clone().detach()
|
||||
log_R_absolute.requires_grad = True
|
||||
T_absolute = T_absolute_init.clone().detach()
|
||||
T_absolute.requires_grad = True
|
||||
|
||||
# the mask the specifies which cameras are going to be optimized
|
||||
# (since we know the first camera is already correct,
|
||||
# we only optimize over the 2nd-to-last cameras)
|
||||
camera_mask = torch.ones(N, 1).float().cuda()
|
||||
camera_mask[0] = 0.
|
||||
|
||||
# init the optimizer
|
||||
optimizer = torch.optim.SGD([log_R_absolute, T_absolute], lr=.1, momentum=0.9)
|
||||
|
||||
# run the optimization
|
||||
n_iter = 2000 # fix the number of iterations
|
||||
for it in range(n_iter):
|
||||
# re-init the optimizer gradients
|
||||
optimizer.zero_grad()
|
||||
|
||||
# compute the absolute camera rotations as
|
||||
# an exponential map of the logarithms (=axis-angles)
|
||||
# of the absolute rotations
|
||||
R_absolute = so3_exponential_map(log_R_absolute * camera_mask)
|
||||
|
||||
# get the current absolute cameras
|
||||
cameras_absolute = SfMPerspectiveCameras(
|
||||
R = R_absolute,
|
||||
T = T_absolute * camera_mask,
|
||||
device = "cuda",
|
||||
)
|
||||
|
||||
# compute the relative cameras as a compositon of the absolute cameras
|
||||
cameras_relative_composed = get_relative_camera(cameras_absolute, relative_edges)
|
||||
|
||||
# compare the composed cameras with the ground truth relative cameras
|
||||
# camera_distance corresponds to $d$ from the description
|
||||
camera_distance = calc_camera_distance(cameras_relative_composed, cameras_relative)
|
||||
|
||||
# our loss function is the camera_distance
|
||||
camera_distance.backward()
|
||||
|
||||
# apply the gradients
|
||||
optimizer.step()
|
||||
|
||||
# plot and print status message
|
||||
if it % 200==0 or it==n_iter-1:
|
||||
status = 'iteration=%3d; camera_distance=%1.3e' % (it, camera_distance)
|
||||
plot_camera_scene(cameras_absolute, cameras_absolute_gt, status)
|
||||
|
||||
print('Optimization finished.')
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -1,280 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
|
||||
# In[5]:
|
||||
|
||||
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
|
||||
|
||||
# # Camera position optimization using differentiable rendering
|
||||
#
|
||||
# In this tutorial we will learn the [x, y, z] position of a camera given a reference image using differentiable rendering.
|
||||
#
|
||||
# We will first initialize a renderer with a starting position for the camera. We will then use this to generate an image, compute a loss with the reference image, and finally backpropagate through the entire pipeline to update the position of the camera.
|
||||
#
|
||||
# This tutorial shows how to:
|
||||
# - load a mesh from an `.obj` file
|
||||
# - initialize a `Camera`, `Shader` and `Renderer`,
|
||||
# - render a mesh
|
||||
# - set up an optimization loop with a loss function and optimizer
|
||||
#
|
||||
|
||||
# ## Set up and imports
|
||||
|
||||
# In[5]:
|
||||
|
||||
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
from tqdm import tqdm_notebook
|
||||
import imageio
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import matplotlib.pyplot as plt
|
||||
from skimage import img_as_ubyte
|
||||
|
||||
# io utils
|
||||
from pytorch3d.io import load_obj
|
||||
|
||||
# datastructures
|
||||
from pytorch3d.structures import Meshes, Textures
|
||||
|
||||
# 3D transformations functions
|
||||
from pytorch3d.transforms import Rotate, Translate
|
||||
|
||||
# rendering components
|
||||
from pytorch3d.renderer import (
|
||||
OpenGLPerspectiveCameras, look_at_view_transform, look_at_rotation,
|
||||
RasterizationSettings, MeshRenderer, MeshRasterizer, BlendParams,
|
||||
SilhouetteShader, PhongShader, PointLights
|
||||
)
|
||||
|
||||
|
||||
# ### Load the Obj
|
||||
#
|
||||
# We will load an obj file and create a **Meshes** object. **Meshes** is a unique datastructure provided in PyTorch3d for working with **batches of meshes of different sizes**. It has several useful class methods which are used in the rendering pipeline.
|
||||
|
||||
# In[6]:
|
||||
|
||||
|
||||
# Set the cuda device
|
||||
device = torch.device("cuda:0")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
# Load the obj and ignore the textures and materials.
|
||||
verts, faces_idx, _ = load_obj("./data/teapot.obj")
|
||||
faces = faces_idx.verts_idx
|
||||
|
||||
# Initialize each vertex to be white in color.
|
||||
verts_rgb = torch.ones_like(verts)[None] # (1, V, 3)
|
||||
textures = Textures(verts_rgb=verts_rgb.to(device))
|
||||
|
||||
# Create a Meshes object for the teapot. Here we have only one mesh in the batch.
|
||||
teapot_mesh = Meshes(
|
||||
verts=[verts.to(device)],
|
||||
faces=[faces.to(device)],
|
||||
textures=textures
|
||||
)
|
||||
|
||||
|
||||
#
|
||||
#
|
||||
# ## Optimization setup
|
||||
|
||||
# ### Create a renderer
|
||||
#
|
||||
# A **renderer** in PyTorch3d is composed of a **rasterizer** and a **shader** which each have a number of subcomponents such as a **camera** (orthgraphic/perspective). Here we initialize some of these components and use default values for the rest.
|
||||
#
|
||||
# For optimizing the camera position we will use a renderer which produces a **silhouette** of the object only and does not apply any **lighting** or **shading**. We will also initialize another renderer which applies full **phong shading** and use this for visualizing the outputs.
|
||||
|
||||
# In[7]:
|
||||
|
||||
|
||||
# Initialize an OpenGL perspective camera.
|
||||
cameras = OpenGLPerspectiveCameras(device=device)
|
||||
|
||||
# To blend the 100 faces we set a few parameters which control the opacity and the sharpness of
|
||||
# edges. Refer to blending.py for more details.
|
||||
blend_params = BlendParams(sigma=1e-4, gamma=1e-4)
|
||||
|
||||
# Define the settings for rasterization and shading. Here we set the output image to be of size
|
||||
# 256x256. To form the blended image we use 100 faces for each pixel. Refer to rasterize_meshes.py
|
||||
# for an explanation of this parameter.
|
||||
raster_settings = RasterizationSettings(
|
||||
image_size=256,
|
||||
blur_radius=np.log(1. / 1e-4 - 1.) * blend_params.sigma,
|
||||
faces_per_pixel=100,
|
||||
bin_size=0
|
||||
)
|
||||
|
||||
# Create a silhouette mesh renderer by composing a rasterizer and a shader.
|
||||
silhouette_renderer = MeshRenderer(
|
||||
rasterizer=MeshRasterizer(
|
||||
cameras=cameras,
|
||||
raster_settings=raster_settings
|
||||
),
|
||||
shader=SilhouetteShader(blend_params=blend_params)
|
||||
)
|
||||
|
||||
|
||||
# We will also create a phong renderer. This is simpler and only needs to render one face per pixel.
|
||||
raster_settings = RasterizationSettings(
|
||||
image_size=256,
|
||||
blur_radius=0.0,
|
||||
faces_per_pixel=1,
|
||||
bin_size=0
|
||||
)
|
||||
# We can add a point light in front of the object.
|
||||
lights = PointLights(device=device, location=((2.0, 2.0, -2.0),))
|
||||
phong_renderer = MeshRenderer(
|
||||
rasterizer=MeshRasterizer(
|
||||
cameras=cameras,
|
||||
raster_settings=raster_settings
|
||||
),
|
||||
shader=PhongShader(device=device, lights=lights)
|
||||
)
|
||||
|
||||
|
||||
# ### Create a reference image
|
||||
#
|
||||
# We will first position the teapot and generate an image. We use helper functions to rotate the teapot to a desired viewpoint. Then we can use the renderers to produce an image. Here we will use both renderers and visualize the silhouette and full shaded image.
|
||||
|
||||
# In[8]:
|
||||
|
||||
|
||||
# Select the viewpoint using spherical angles
|
||||
distance = 3 # distance from camera to the object
|
||||
elevation = 40.0 # angle of elevation in degrees
|
||||
azimuth = 10.0 # angle of azimuth rotation in degrees
|
||||
|
||||
# Get the position of the camera based on the spherical angles
|
||||
R, T = look_at_view_transform(distance, elevation, azimuth, device=device)
|
||||
|
||||
# Render the teapot providing the values of R and T.
|
||||
silhouete = silhouette_renderer(meshes_world=teapot_mesh, R=R, T=T)
|
||||
image_ref = phong_renderer(meshes_world=teapot_mesh, R=R, T=T)
|
||||
|
||||
silhouete = silhouete.cpu().numpy()
|
||||
image_ref = image_ref.cpu().numpy()
|
||||
|
||||
plt.figure(figsize=(10, 10))
|
||||
plt.subplot(1, 2, 1)
|
||||
plt.imshow(silhouete.squeeze()[..., 3]) # only plot the alpha channel of the RGBA image
|
||||
plt.grid("off")
|
||||
plt.subplot(1, 2, 2)
|
||||
plt.imshow(image_ref.squeeze())
|
||||
plt.grid("off")
|
||||
|
||||
|
||||
# ### Set up a basic model
|
||||
#
|
||||
# Here we create a simple model class and initialize a parameter for the camera position.
|
||||
|
||||
# In[10]:
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, meshes, renderer, image_ref):
|
||||
super().__init__()
|
||||
self.meshes = meshes
|
||||
self.device = meshes.device
|
||||
self.renderer = renderer
|
||||
|
||||
# Get the silhouette of the reference RGB image by finding all the non zero values.
|
||||
image_ref = torch.from_numpy((image_ref[..., :3].max(-1) != 0).astype(np.float32))
|
||||
self.register_buffer('image_ref', image_ref)
|
||||
|
||||
# Create an optimizable parameter for the x, y, z position of the camera.
|
||||
self.camera_position = nn.Parameter(
|
||||
torch.from_numpy(np.array([3.0, 6.9, -2.5], dtype=np.float32)).to(meshes.device))
|
||||
|
||||
def forward(self):
|
||||
|
||||
# Render the image using the updated camera position. Based on the new position of the
|
||||
# camer we calculate the rotation and translation matrices
|
||||
R = look_at_rotation(self.camera_position[None, :], device=self.device) # (1, 3, 3)
|
||||
T = -torch.bmm(R.transpose(1, 2), self.camera_position[None, :, None])[:, :, 0] # (1, 3)
|
||||
|
||||
image = self.renderer(meshes_world=self.meshes.clone(), R=R, T=T)
|
||||
|
||||
# Calculate the silhouette loss
|
||||
loss = torch.sum((image[..., 3] - self.image_ref) ** 2)
|
||||
return loss, image
|
||||
|
||||
|
||||
|
||||
# ## Initialize the model and optimizer
|
||||
#
|
||||
# Now we can create an instance of the **model** above and set up an **optimizer** for the camera position parameter.
|
||||
|
||||
# In[11]:
|
||||
|
||||
|
||||
# We will save images periodically and compose them into a GIF.
|
||||
filename_output = "./teapot_optimization_demo.gif"
|
||||
writer = imageio.get_writer(filename_output, mode='I', duration=0.3)
|
||||
|
||||
# Initialize a model using the renderer, mesh and reference image
|
||||
model = Model(meshes=teapot_mesh, renderer=silhouette_renderer, image_ref=image_ref).to(device)
|
||||
|
||||
# Create an optimizer. Here we are using Adam and we pass in the parameters of the model
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
|
||||
|
||||
|
||||
# ## Visualise the starting position and the reference position
|
||||
|
||||
# In[12]:
|
||||
|
||||
|
||||
plt.figure(figsize=(10, 10))
|
||||
|
||||
_, image_init = model()
|
||||
plt.subplot(1, 2, 1)
|
||||
plt.imshow(image_init.detach().squeeze().cpu().numpy()[..., 3])
|
||||
plt.grid("off")
|
||||
plt.title("Starting position")
|
||||
|
||||
plt.subplot(1, 2, 2)
|
||||
plt.imshow(model.image_ref.cpu().numpy().squeeze())
|
||||
plt.grid("off")
|
||||
plt.title("Reference silhouette")
|
||||
|
||||
|
||||
# ## Run the optimization
|
||||
#
|
||||
# We run several iterations of the forward and backward pass and save outputs every 10 iterations. When this has finished take a look at `./teapot_optimization_demo.gif` for a cool gif of the optimization process!
|
||||
|
||||
# In[13]:
|
||||
|
||||
|
||||
loop = tqdm_notebook(range(200))
|
||||
for i in loop:
|
||||
optimizer.zero_grad()
|
||||
loss, _ = model()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
loop.set_description('Optimizing (loss %.4f)' % loss.data)
|
||||
|
||||
if loss.item() < 200:
|
||||
break
|
||||
|
||||
# Save outputs to create a GIF.
|
||||
if i % 10 == 0:
|
||||
R = look_at_rotation(model.camera_position[None, :], device=model.device)
|
||||
T = -torch.bmm(R.transpose(1, 2), model.camera_position[None, :, None])[:, :, 0] # (1, 3)
|
||||
image = phong_renderer(meshes_world=model.meshes.clone(), R=R, T=T)
|
||||
image = image[0, ..., :3].detach().squeeze().cpu().numpy()
|
||||
image = img_as_ubyte(image)
|
||||
writer.append_data(image)
|
||||
|
||||
plt.figure()
|
||||
plt.imshow(image[..., :3])
|
||||
plt.title("iter: %d, loss: %0.2f" % (i, loss.data))
|
||||
plt.grid("off")
|
||||
plt.axis("off")
|
||||
|
||||
writer.close()
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -1,253 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
|
||||
# In[1]:
|
||||
|
||||
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
|
||||
|
||||
# # Deform a source mesh to form a target mesh using 3D loss functions
|
||||
|
||||
# In this tutorial, we learn to deform an initial generic shape (e.g. sphere) to fit a target shape.
|
||||
#
|
||||
# We will cover:
|
||||
#
|
||||
# - How to **load a mesh** from an `.obj` file
|
||||
# - How to use the PyTorch3d **Meshes** datastructure
|
||||
# - How to use 4 different PyTorch3d **mesh loss functions**
|
||||
# - How to set up an **optimization loop**
|
||||
#
|
||||
#
|
||||
# Starting from a sphere mesh, we learn the offset to each vertex in the mesh such that
|
||||
# the predicted mesh is closer to the target mesh at each optimization step. To achieve this we minimize:
|
||||
#
|
||||
# + `chamfer_distance`, the distance between the predicted (deformed) and target mesh, defined as the chamfer distance between the set of pointclouds resulting from **differentiably sampling points** from their surfaces.
|
||||
#
|
||||
# However, solely minimizing the chamfer distance between the predicted and the target mesh will lead to a non-smooth shape (verify this by setting `w_chamfer=1.0` and all other weights to `0.0`).
|
||||
#
|
||||
# We enforce smoothness by adding **shape regularizers** to the objective. Namely, we add:
|
||||
#
|
||||
# + `mesh_edge_length`, which minimizes the length of the edges in the predicted mesh.
|
||||
# + `mesh_normal_consistency`, which enforces consistency across the normals of neighboring faces.
|
||||
# + `mesh_laplacian_smoothing`, which is the laplacian regularizer.
|
||||
|
||||
# ## 0. Import modules
|
||||
|
||||
# In[2]:
|
||||
|
||||
|
||||
import os
|
||||
import torch
|
||||
from pytorch3d.io import load_obj, save_obj
|
||||
from pytorch3d.structures import Meshes
|
||||
from pytorch3d.utils import ico_sphere
|
||||
from pytorch3d.ops import sample_points_from_meshes
|
||||
from pytorch3d.loss import (
|
||||
chamfer_distance,
|
||||
mesh_edge_loss,
|
||||
mesh_laplacian_smoothing,
|
||||
mesh_normal_consistency,
|
||||
)
|
||||
import numpy as np
|
||||
from tqdm import tqdm_notebook
|
||||
get_ipython().run_line_magic('matplotlib', 'notebook')
|
||||
from mpl_toolkits.mplot3d import Axes3D
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib as mpl
|
||||
mpl.rcParams['savefig.dpi'] = 80
|
||||
mpl.rcParams['figure.dpi'] = 80
|
||||
|
||||
# Set the device
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
|
||||
# ## 1. Load an obj file and create a Meshes object
|
||||
|
||||
# In[4]:
|
||||
|
||||
|
||||
# The path to the target 3D model we wish to fit
|
||||
# e.g. download https://free3d.com/3d-model/-dolphin-v1--12175.html and save in ./data/dolphin
|
||||
trg_obj = os.path.join('./data/doplhin', '10014_dolphin_v2_max2011_it2.obj')
|
||||
|
||||
|
||||
# In[6]:
|
||||
|
||||
|
||||
# We read the target 3D model using load_obj
|
||||
verts, faces, aux = load_obj(trg_obj)
|
||||
|
||||
# verts is a FloatTensor of shape (V, 3) where V is the number of vertices in the mesh
|
||||
# faces is an object which contains the following LongTensors: verts_idx, normals_idx and textures_idx
|
||||
# For this tutorial, normals and textures are ignored.
|
||||
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])
|
||||
|
||||
|
||||
# In[11]:
|
||||
|
||||
|
||||
# We initialize the source shape to be a sphere of radius 1
|
||||
src_mesh = ico_sphere(4, device)
|
||||
|
||||
|
||||
# ### Visualize the source and target meshes
|
||||
|
||||
# In[12]:
|
||||
|
||||
|
||||
def plot_pointcloud(mesh, title=""):
|
||||
verts = mesh.verts_packed()
|
||||
faces = mesh.faces_packed()
|
||||
x, y, z = verts.clone().detach().cpu().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)
|
||||
plt.show()
|
||||
|
||||
|
||||
# In[13]:
|
||||
|
||||
|
||||
get_ipython().run_line_magic('matplotlib', 'notebook')
|
||||
plot_pointcloud(trg_mesh, "Target mesh")
|
||||
plot_pointcloud(src_mesh, "Source mesh")
|
||||
|
||||
|
||||
# ## 3. Optimization loop
|
||||
|
||||
# In[14]:
|
||||
|
||||
|
||||
# We will learn to deform the source mesh by offsetting its vertices
|
||||
# The shape of the derform parameters is equal to the total number of vertices in src_mesh
|
||||
deform_verts = torch.full(src_mesh.verts_packed().shape, 0.0, device=device, requires_grad=True)
|
||||
|
||||
|
||||
# In[15]:
|
||||
|
||||
|
||||
# The optimizer
|
||||
optimizer = torch.optim.SGD([deform_verts], lr=1.0, momentum=0.9)
|
||||
|
||||
|
||||
# In[16]:
|
||||
|
||||
|
||||
# Number of optimization steps
|
||||
Niter = 2000
|
||||
# Weight for the chamfer loss
|
||||
w_chamfer = 1.0
|
||||
# Weight for mesh edge loss
|
||||
w_edge = 1.0
|
||||
# Weight for mesh normal consistency
|
||||
w_normal = 0.01
|
||||
# Weight for mesh laplacian smoothing
|
||||
w_laplacian = 0.1
|
||||
# Plot period for the losses
|
||||
plot_period = 250
|
||||
loop = tqdm_notebook(range(Niter))
|
||||
|
||||
chamfer_losses = []
|
||||
laplacian_losses = []
|
||||
edge_losses = []
|
||||
normal_losses = []
|
||||
|
||||
get_ipython().run_line_magic('matplotlib', 'inline')
|
||||
|
||||
for i in loop:
|
||||
# Initialize optimizer
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Deform the mesh
|
||||
new_src_mesh = src_mesh.offset_verts(deform_verts)
|
||||
|
||||
# We sample 5k points from the surface of each mesh
|
||||
sample_trg = sample_points_from_meshes(trg_mesh, 5000)
|
||||
sample_src = sample_points_from_meshes(new_src_mesh, 5000)
|
||||
|
||||
# We compare the two sets of pointclouds by computing (a) the chamfer loss
|
||||
loss_chamfer, _ = chamfer_distance(sample_trg, sample_src)
|
||||
|
||||
# and (b) the edge length of the predicted mesh
|
||||
loss_edge = mesh_edge_loss(new_src_mesh)
|
||||
|
||||
# mesh normal consistency
|
||||
loss_normal = mesh_normal_consistency(new_src_mesh)
|
||||
|
||||
# mesh laplacian smoothing
|
||||
loss_laplacian = mesh_laplacian_smoothing(new_src_mesh, method="uniform")
|
||||
|
||||
# Weighted sum of the losses
|
||||
loss = loss_chamfer * w_chamfer + loss_edge * w_edge + loss_normal * w_normal + loss_laplacian * w_laplacian
|
||||
|
||||
# Print the losses
|
||||
loop.set_description('total_loss = %.6f' % loss)
|
||||
|
||||
# Save the losses for plotting
|
||||
chamfer_losses.append(loss_chamfer)
|
||||
edge_losses.append(loss_edge)
|
||||
normal_losses.append(loss_normal)
|
||||
laplacian_losses.append(loss_laplacian)
|
||||
|
||||
# Plot mesh
|
||||
if i % plot_period == 0:
|
||||
plot_pointcloud(new_src_mesh, title="iter: %d" % i)
|
||||
|
||||
# Optimization step
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
|
||||
# ## 4. Visualize the loss
|
||||
|
||||
# In[17]:
|
||||
|
||||
|
||||
fig = plt.figure(figsize=(13, 5))
|
||||
ax = fig.gca()
|
||||
ax.plot(chamfer_losses, label="chamfer loss")
|
||||
ax.plot(edge_losses, label="edge loss")
|
||||
ax.plot(normal_losses, label="normal loss")
|
||||
ax.plot(laplacian_losses, label="laplacian loss")
|
||||
ax.legend(fontsize="16")
|
||||
ax.set_xlabel("Iteration", fontsize="16")
|
||||
ax.set_ylabel("Loss", fontsize="16")
|
||||
ax.set_title("Loss vs iterations", fontsize="16")
|
||||
|
||||
|
||||
# ## 5. Save the predicted mesh
|
||||
|
||||
# In[ ]:
|
||||
|
||||
|
||||
# Fetch the verts and faces of the final predicted mesh
|
||||
final_verts, final_faces = new_src_mesh.get_mesh_verts_faces(0)
|
||||
|
||||
# Scale normalize back to the original target size
|
||||
final_verts = final_verts * scale + center
|
||||
|
||||
# Store the predicted mesh using save_obj
|
||||
final_obj = os.path.join('./', 'final_model.obj')
|
||||
save_obj(final_obj, final_verts, final_faces)
|
||||
|
||||
|
||||
# ## 6. Conclusion
|
||||
#
|
||||
# In this tutorial we learnt how to load a mesh from an obj file, initialize a PyTorch3d datastructure called **Meshes**, set up an optimization loop and use four different PyTorch3d mesh loss functions.
|
||||
File diff suppressed because one or more lines are too long
@@ -1,268 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
|
||||
# In[ ]:
|
||||
|
||||
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
|
||||
|
||||
# # Render a textured mesh
|
||||
#
|
||||
# This tutorial shows how to:
|
||||
# - load a mesh and textures from an `.obj` file.
|
||||
# - set up a renderer
|
||||
# - render the mesh
|
||||
# - vary the rendering settings such as lighting and camera position
|
||||
# - use the batching features of the pytorch3d API to render the mesh from different viewpoints
|
||||
|
||||
# ## Import modules
|
||||
|
||||
# In[1]:
|
||||
|
||||
|
||||
import os
|
||||
import torch
|
||||
import matplotlib.pyplot as plt
|
||||
from skimage.io import imread
|
||||
|
||||
# Util function for loading meshes
|
||||
from pytorch3d.io import load_obj
|
||||
|
||||
# Data structures and functions for rendering
|
||||
from pytorch3d.structures import Meshes, Textures
|
||||
from pytorch3d.renderer import (
|
||||
look_at_view_transform,
|
||||
OpenGLPerspectiveCameras,
|
||||
PointLights,
|
||||
DirectionalLights,
|
||||
Materials,
|
||||
RasterizationSettings,
|
||||
MeshRenderer,
|
||||
MeshRasterizer,
|
||||
TexturedPhongShader
|
||||
)
|
||||
|
||||
# add path for demo utils
|
||||
import sys
|
||||
import os
|
||||
sys.path.append(os.path.abspath(''))
|
||||
from utils import image_grid
|
||||
|
||||
|
||||
# ### Load a mesh and texture file
|
||||
#
|
||||
# Load an `.obj` file and it's associated `.mtl` file and create a **Textures** and **Meshes** object.
|
||||
#
|
||||
# **Meshes** is a unique datastructure provided in PyTorch3d for working with batches of meshes of different sizes.
|
||||
#
|
||||
# **Textures** is an auxillary datastructure for storing texture information about meshes.
|
||||
#
|
||||
# **Meshes** has several class methods which are used throughout the rendering pipeline.
|
||||
|
||||
# In[2]:
|
||||
|
||||
|
||||
# Setup
|
||||
device = torch.device("cuda:0")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
# Set paths
|
||||
DATA_DIR = "./data"
|
||||
obj_filename = os.path.join(DATA_DIR, "cow_mesh/cow.obj")
|
||||
|
||||
# Load obj file
|
||||
verts, faces, aux = load_obj(obj_filename)
|
||||
faces_idx = faces.verts_idx.to(device)
|
||||
verts = verts.to(device)
|
||||
|
||||
# Get textures from the outputs of the load_obj function
|
||||
# the `aux` variable contains the texture maps and vertex uv coordinates.
|
||||
# Refer to the `obj_io.load_obj` function for full API reference.
|
||||
# Here we only have one texture map for the whole mesh.
|
||||
verts_uvs = aux.verts_uvs[None, ...].to(device) # (N, V, 2)
|
||||
faces_uvs = faces.textures_idx[None, ...].to(device) # (N, F, 3)
|
||||
tex_maps = aux.texture_images
|
||||
texture_image = list(tex_maps.values())[0]
|
||||
texture_image = texture_image[None, ...].to(device) # (N, H, W, 3)
|
||||
|
||||
# Create a textures object
|
||||
tex = Textures(verts_uvs=verts_uvs, faces_uvs=faces_uvs, maps=texture_image)
|
||||
|
||||
# Create a meshes object with textures
|
||||
mesh = Meshes(verts=[verts], faces=[faces_idx], textures=tex)
|
||||
|
||||
|
||||
# #### Let's visualize the texture map
|
||||
|
||||
# In[3]:
|
||||
|
||||
|
||||
plt.figure(figsize=(7,7))
|
||||
plt.imshow(texture_image.squeeze().cpu().numpy())
|
||||
plt.grid("off")
|
||||
plt.axis('off')
|
||||
|
||||
|
||||
# ## Create a renderer
|
||||
#
|
||||
# A renderer in PyTorch3d is composed of a **rasterizer** and a **shader** which each have a number of subcomponents such as a **camera** (orthgraphic/perspective). Here we initialize some of these components and use default values for the rest.
|
||||
#
|
||||
# In this example we will first create a **renderer** which uses a **perspective camera**, a **point light** and applies **phong shading**. Then we learn how to vary different components using the modular API.
|
||||
|
||||
# In[4]:
|
||||
|
||||
|
||||
# Initialize an OpenGL perspective camera.
|
||||
R, T = look_at_view_transform(2.7, 10, 20)
|
||||
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)
|
||||
|
||||
# Define the settings for rasterization and shading. Here we set the output image to be of size
|
||||
# 512x512. As we are rendering images for visualization purposes only we will set faces_per_pixel=1
|
||||
# and blur_radius=0.0. Refer to rasterize_meshes.py for explanations of these parameters.
|
||||
raster_settings = RasterizationSettings(
|
||||
image_size=512,
|
||||
blur_radius=0.0,
|
||||
faces_per_pixel=1,
|
||||
bin_size=0
|
||||
)
|
||||
|
||||
# Place a point light in front of the object
|
||||
lights = PointLights(device=device, location=[[1.0, 1.0, -2.0]])
|
||||
|
||||
# Create a phong renderer by composing a rasterizer and a shader. The textured phong shader will
|
||||
# interpolate the texture uv coordinates for each vertex, sample from a texture image and
|
||||
# apply the Phong lighting model
|
||||
renderer = MeshRenderer(
|
||||
rasterizer=MeshRasterizer(
|
||||
cameras=cameras,
|
||||
raster_settings=raster_settings
|
||||
),
|
||||
shader=TexturedPhongShader(
|
||||
device=device,
|
||||
cameras=cameras,
|
||||
lights=lights
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ## Render the mesh
|
||||
|
||||
# The light is in front of the object so it is bright and the image has specular highlights.
|
||||
|
||||
# In[5]:
|
||||
|
||||
|
||||
images = renderer(mesh)
|
||||
plt.figure(figsize=(10, 10))
|
||||
plt.imshow(images[0, ..., :3].cpu().numpy())
|
||||
plt.grid("off")
|
||||
plt.axis("off")
|
||||
|
||||
|
||||
# ## Move the light behind the object and re-render
|
||||
#
|
||||
# We can pass arbirary keyword arguments to the `rasterizer`/`shader` via the call to the `renderer` so the renderer does not need to be reinitialized if any of the settings change/
|
||||
#
|
||||
# In this case, we can simply update the location of the lights and pass them into the call to the renderer.
|
||||
#
|
||||
# The image is now dark as there is only ambient lighting, and there are no specular highlights.
|
||||
|
||||
# In[6]:
|
||||
|
||||
|
||||
lights.location = torch.tensor([0.0, 0.0, +1.0], device=device)[None]
|
||||
images = renderer(mesh, lights=lights)
|
||||
|
||||
|
||||
# In[7]:
|
||||
|
||||
|
||||
plt.figure(figsize=(10, 10))
|
||||
plt.imshow(images[0, ..., :3].cpu().numpy())
|
||||
plt.grid("off")
|
||||
plt.axis("off")
|
||||
|
||||
|
||||
# ## Rotate the object, modify the material properties or light properties
|
||||
#
|
||||
# We can also change many other settings in the rendering pipeline. Here we:
|
||||
#
|
||||
# - change the **viewing angle** of the camera
|
||||
# - change the **position** of the point light
|
||||
# - change the **material reflectance** properties of the mesh
|
||||
|
||||
# In[8]:
|
||||
|
||||
|
||||
# Rotate the object by increasing the azimuth angle
|
||||
R, T = look_at_view_transform(dist=2.7, elev=10, azim=50)
|
||||
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)
|
||||
|
||||
# Move the light location to be in front of the object again
|
||||
lights.location = torch.tensor([[5.0, 5.0, -2.0]], device=device)
|
||||
|
||||
# Change specular color to green and change material shininess
|
||||
materials = Materials(
|
||||
device=device,
|
||||
specular_color=[[0.0, 1.0, 0.0]],
|
||||
shininess=10.0
|
||||
)
|
||||
|
||||
# Re render the mesh, passing in keyword arguments for the modified components.
|
||||
images = renderer(mesh, lights=lights, materials=materials, cameras=cameras)
|
||||
|
||||
|
||||
# In[9]:
|
||||
|
||||
|
||||
plt.figure(figsize=(10, 10))
|
||||
plt.imshow(images[0, ..., :3].cpu().numpy())
|
||||
plt.grid("off")
|
||||
plt.axis("off")
|
||||
|
||||
|
||||
# ## Batched Rendering
|
||||
#
|
||||
# One of the core design choices of the PyTorch3d API is to suport **batched inputs for all components**.
|
||||
# The renderer and associated components can take batched inputs and **render a batch of output images in one forward pass**. We will now use this feature to render the mesh from many different viewpoints.
|
||||
#
|
||||
|
||||
# In[10]:
|
||||
|
||||
|
||||
# Set batch size - this is the number of different viewpoints from which we want to render the mesh.
|
||||
batch_size = 20
|
||||
|
||||
# Create a batch of meshes by repeating the cow mesh and associated textures.
|
||||
# Meshes has a useful `extend` method which allows us do this very easily.
|
||||
# This also extends the textures.
|
||||
meshes = mesh.extend(batch_size)
|
||||
|
||||
# Get a batch of viewing angles.
|
||||
elev = torch.linspace(0, 360, batch_size)
|
||||
azim = torch.linspace(0, 360, batch_size)
|
||||
|
||||
# All the cameras helper methods support mixed type inputs and broadcasting. So we can
|
||||
# view the camera from the same distance and specify dist=2.7 as a float,
|
||||
# and then specify elevation and azimuth angles for each viewpoint as tensors.
|
||||
R, T = look_at_view_transform(dist=2.7, elev=elev, azim=azim)
|
||||
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)
|
||||
|
||||
# Move the light back in front of the object
|
||||
lights.location = torch.tensor([[1.0, 1.0, -5.0]], device=device)
|
||||
|
||||
|
||||
# In[11]:
|
||||
|
||||
|
||||
# We can pass arbirary keyword arguments to the rasterizer/shader via the renderer
|
||||
# so the renderer does not need to be reinitialized if any of the settings change.
|
||||
images = renderer(meshes, cameras=cameras, lights=lights)
|
||||
|
||||
|
||||
# In[12]:
|
||||
|
||||
|
||||
image_grid(images.cpu().numpy(), rows=4, cols=5, rgb=True)
|
||||
|
||||
Reference in New Issue
Block a user