PyTorch 3D 模型渲染與姿勢估計最佳化實戰

– 9 分鐘閱讀

隨著深度學習技術的發展，使用 PyTorch 進行 3D 模型渲染和姿勢估計已成為一個熱門研究方向。本文將深入探討如何利用 PyTorch 的可微分渲染特性，實作 3D 模型的精確渲染和姿勢最佳化。我們將逐步介紹模型的構建、渲染器的設定、損失函式的設計以及最佳化流程，並最終展示如何視覺化最佳化結果。此方法的核心概念是利用可微分渲染器，計算渲染影像相對於模型引數的梯度，進而使用梯度下降等最佳化演算法調整模型引數，最終使渲染影像與目標影像儘可能接近。

模型架構

在這個例子中，我們定義了一個 PyTorch 的模型類別 Model，用於 3D 重建任務。這個模型類別繼承自 PyTorch 的 nn.Module 類別。

初始化模型

在 __init__ 方法中，我們初始化了模型的屬性，包括：

  • meshes：3D 網格資料
  • renderer：渲染器物件
  • image_ref：參考影像
  • camera_position：相機位置引數
class Model(nn.Module):
    def __init__(self, meshes, renderer, image_ref):
        super().__init__()
        self.meshes = meshes
        self.renderer = renderer
        image_ref = torch.from_numpy(
            (image_ref[..., :3].max(-1) != 1).astype(np.float32)
        )
        self.register_buffer('image_ref', image_ref)
        self.camera_position = nn.Parameter(
            torch.from_numpy(np.array([3.0, 6.9, +2.5]))
        )

前向傳播

在 forward 方法中，我們定義了模型的前向傳播過程。這個過程包括：

  • 根據相機位置計算旋轉矩陣 R
  • 根據旋轉矩陣和相機位置計算平移矩陣 T
  • 渲染 3D 網格資料以獲得預測影像
  • 計算預測影像和參考影像之間的損失函式
def forward(self):
    R = look_at_rotation(self.camera_position[None, :])
    T = -torch.bmm(
        R.transpose(1, 2),
        self.camera_position[None, :, None]
    )[:, :, 0]  # (1, 3)
    image = self.renderer(meshes_world=self.meshes.clone())
    loss = torch.mean((image - self.image_ref) ** 2)
    return loss

訓練模型

要訓練這個模型，我們需要定義一個損失函式和一個最佳化器。然後，我們可以使用 PyTorch 的 backward 方法計算梯度，並使用最佳化器更新模型引數。

model = Model(meshes, renderer, image_ref)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(100):
    optimizer.zero_grad()
    loss = model.forward()
    loss.backward()
    optimizer.step()
    print(f'Epoch {epoch+1}, Loss: {loss.item()}')

這個模型可以用於 3D 重建任務，例如從 2D 影像中重建 3D 物體。

使用 PyTorch 進行可微分渲染的物體姿勢估計

在這個例子中，我們將使用 PyTorch 進行可微分渲染的物體姿勢估計。首先，我們需要定義一個模型類別 Model，它包含了渲染器和引數更新的邏輯。

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

class Model(nn.Module):
    def __init__(self, meshes, renderer, image_ref):
        super(Model, self).__init__()
        self.meshes = meshes
        self.renderer = renderer
        self.image_ref = image_ref

    def forward(self):
        # 進行渲染
        image = self.renderer(self.meshes)
        # 計算損失
        loss = torch.sum((image[..., 3] - self.image_ref) ** 2)
        return loss, image

接下來，我們可以建立一個模型例項和定義最佳化器。最佳化器將用於更新模型引數以最小化損失。

# 建立模型例項
model = Model(meshes=teapot_mesh, renderer=silhouette_renderer, image_ref=image_ref).to(device)

# 定義最佳化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

現在，我們可以進行最佳化迭代。在每次迭代中，我們將儲存渲染的影像。

# 進行最佳化迭代
for i in range(200):
    if i % 10 == 0:
        # 渲染影像
        _, image_init = model()
        
        # 顯示和儲存影像
        plt.figure(figsize=(10, 10))
        plt.imshow(image_init.detach().squeeze().cpu().numpy()[..., 3])
        plt.grid(False)
        plt.title("渲染影像")
        plt.savefig(os.path.join(output_dir, f'iteration_{i}.png'))
        plt.close()
        
    # 更新模型引數
    optimizer.zero_grad()
    loss, _ = model()
    loss.backward()
    optimizer.step()

在這個例子中，我們使用了 PyTorch 的 Adam 最佳化器來更新模型引數。最佳化器的學習率設定為 0.05。在每次迭代中，我們計算損失、更新模型引數，並儲存渲染的影像。

內容解密：

  • 我們定義了一個 Model 類別，它包含了渲染器和引數更新的邏輯。
  • 我們建立了一個模型例項和定義最佳化器。
  • 我們進行最佳化迭代，在每次迭代中，我們渲染影像、更新模型引數，並儲存渲染的影像。

圖表翻譯：

  graph LR
    A[模型定義] --> B[模型例項化]
    B --> C[最佳化器定義]
    C --> D[最佳化迭代]
    D --> E[渲染影像]
    E --> F[更新模型引數]
    F --> G[儲存渲染影像]

這個流程圖顯示了我們的模型定義、模型例項化、最佳化器定義、最佳化迭代、渲染影像、更新模型引數和儲存渲染影像的過程。

使用 PyTorch 進行可微分渲染的物體姿勢估計

在這個例子中，我們將使用 PyTorch 進行可微分渲染，估計物體的姿勢。可微分渲染是一種技術，允許我們計算渲染影像的梯度，從而可以使用最佳化演算法來估計物體的姿勢。

安裝必要的套件

pip install torch torchvision

載入必要的套件

import torch
import torch.nn as nn
import torchvision
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

定義模型

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.meshes = ...  # 載入 3D 物體模型
        self.camera_position = ...  # 相機位置

    def forward(self):
        # 使用可微分渲染進行渲染
        image = phong_renderer(meshes_world=self.meshes.clone(), 
                               R=look_at_rotation(self.camera_position[None, :]), 
                               T=-torch.bmm(look_at_rotation(self.camera_position[None, :]).transpose(1, 2), 
                                              self.camera_position[None, :, None])[:, :, 0])
        return image

定義最佳化器

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

定義損失函式

def loss_function(image, target_image):
    return torch.mean((image - target_image) ** 2)

進行最佳化

for i in range(100):
    optimizer.zero_grad()
    image = model()
    loss = loss_function(image, target_image)
    loss.backward()
    optimizer.step()
    if loss.item() < 500:
        break
    image = image[0, ..., :3].detach().squeeze().cpu().numpy()
    image = img_as_ubyte(image)
    plt.figure()
    plt.imshow(image[..., :3])
    plt.title("iter: %d, loss: %0.2f" % (i, loss.data))
    plt.axis("off")
    plt.savefig(os.path.join(output_dir, 'fitting_' + str(i) + '.png'))
    plt.close()

視覺化結果

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot3D(model.meshes.clone().vertices[:, 0], model.meshes.clone().vertices[:, 1], model.meshes.clone().vertices[:, 2], 'b-')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.show()

結果

最終的結果將是一個估計的 3D 物體姿勢，與原始影像進行比較，可以看到估計的姿勢與原始姿勢非常接近。這個例子展示了使用可微分渲染進行物體姿勢估計的強大能力。

建立3D模型渲染器

首先，我們需要匯入必要的套件，包括 torch、numpy、matplotlib 和 skimage。同時，我們也需要從 pytorch3d 中匯入相關的類別和函式，例如 FoVPerspectiveCameras、look_at_view_transform、MeshRenderer 等。

import os
import torch
import numpy as np
import torch.nn as nn
import matplotlib.pyplot as plt
from skimage import img_as_ubyte
from pytorch3d.renderer import (
    FoVPerspectiveCameras, look_at_view_transform,
    look_at_rotation, RasterizationSettings,
    MeshRenderer, MeshRasterizer, BlendParams,
    SoftSilhouetteShader, HardPhongShader,
    PointLights, SoftPhongShader
)

接下來，我們需要設定 PyTorch 的裝置（device），如果有可用的 GPU，我們會使用 GPU，否則使用 CPU。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

設定輸出目錄（output directory）為 ./result_cow，這將是我們儲存渲染結果的位置。

output_dir = './result_cow'

接著，我們需要載入 3D 模型的網格資料。這裡我們使用 load_objs_as_meshes 函式從 cow.obj 檔案中載入網格資料，並指定裝置為我們之前設定的 device。

obj_filename = "./data/cow_mesh/cow.obj"
cow_mesh = load_objs_as_meshes([obj_filename], device=device)

最後，我們需要定義攝影機（camera）和光源（light source）。這裡我們使用 FoVPerspectiveCameras 類別建立攝影機，並使用 look_at_view_transform 函式設定攝影機的視角。同時，我們也需要定義光源的位置和強度。

# 定義攝影機
cameras = FoVPerspectiveCameras(device=device)

# 定義攝影機的視角
view_transform = look_at_view_transform(0.0, 0.0, 2.0, 0.0, 0.0, 0.0)

# 定義光源
lights = PointLights(device=device, locations=[[0.0, 0.0, 2.0]])

# 定義渲染設定
raster_settings = RasterizationSettings(
    image_size=256,
    pixel_depth=16,
    cull_backfaces=True,
    cull_frontfaces=False,
)

# 定義渲染器
renderer = MeshRenderer(
    cameras=cameras,
    raster_settings=raster_settings,
    device=device,
)

內容解密：

以上程式碼的作用是建立一個 3D 模型渲染器，包括攝影機、光源和渲染設定的定義。這些設定將用於渲染 3D 模型的網格資料。

圖表翻譯：

  flowchart TD
    A[載入網格資料] --> B[定義攝影機]
    B --> C[定義攝影機的視角]
    C --> D[定義光源]
    D --> E[定義渲染設定]
    E --> F[建立渲染器]

這個圖表描述了建立 3D 模型渲染器的流程，包括載入網格資料、定義攝影機、攝影機的視角、光源、渲染設定和建立渲染器的步驟。

建立渲染器

在這個步驟中，我們需要建立兩個渲染器：renderer_silhouette 和 renderer_textured。這兩個渲染器都使用了 MeshRenderer 類別，但它們的設定和用途不同。

建立輪廓渲染器（renderer_silhouette）

這個渲染器用於生成物體的輪廓影像。為了建立這個渲染器，我們需要設定 RasterizationSettings 和 SoftSilhouetteShader。

# 設定混合引數
blend_params = BlendParams(sigma=1e-4, gamma=1e-4)

# 設定光柵化設定
raster_settings = RasterizationSettings(
    image_size=256,
    blur_radius=np.log(1. / 1e-4 - 1.) * blend_params.sigma,
    faces_per_pixel=100,
)

# 建立輪廓渲染器
renderer_silhouette = MeshRenderer(
    rasterizer=MeshRasterizer(
        cameras=cameras,
        raster_settings=raster_settings
    ),
    shader=SoftSilhouetteShader(blend_params=blend_params)
)

建立紋理渲染器（renderer_textured）

這個渲染器用於生成物體的紋理影像。為了建立這個渲染器，我們需要設定 RasterizationSettings 和使用一個適合的紋理著色器。

# 設定sigma值
sigma = 1e-4

# 設定光柵化設定
raster_settings_soft = RasterizationSettings(
    image_size=256,
    blur_radius=np.log(1. / 1e-4 - 1.) * sigma,
    faces_per_pixel=50,
)

# 建立紋理渲染器
renderer_textured = MeshRenderer(
    rasterizer=MeshRasterizer(
        cameras=cameras,
        raster_settings=raster_settings_soft
    ),
    # 使用適合的紋理著色器
    shader=TexturedShader()
)

圖表翻譯：

  flowchart TD
    A[建立渲染器] --> B[設定混合引數]
    B --> C[設定光柵化設定]
    C --> D[建立輪廓渲染器]
    D --> E[設定sigma值]
    E --> F[設定光柵化設定]
    F --> G[建立紋理渲染器]

內容解密：

這兩個渲染器的建立過程中，我們需要設定不同的引數和著色器，以滿足不同的需求。輪廓渲染器用於生成物體的輪廓影像，而紋理渲染器用於生成物體的紋理影像。透過這兩個渲染器的建立，我們可以得到物體的輪廓和紋理影像，從而實作物體的渲染和視覺化。

3D繪圖設定與渲染

在進行3D繪圖時，設定適當的渲染器和燈光效果是非常重要的。以下是設定Phong渲染器和Mesh渲染器的步驟：

Phong渲染器設定

Phong渲染器是一種常用的渲染器，主要用於生成光滑的3D物體表面。以下是設定Phong渲染器的步驟：

raster_settings = RasterizationSettings(
    image_size=256,
    blur_radius=0.0,
    faces_per_pixel=1,
)

phong_renderer = MeshRenderer(
    rasterizer=MeshRasterizer(
        cameras=cameras,
        raster_settings=raster_settings
    ),
    shader=HardPhongShader(
        device=device,
        cameras=cameras,
        lights=lights
    )
)

Mesh渲染器設定

Mesh渲染器主要用於生成清晰的3D物體表面。以下是設定Mesh渲染器的步驟：

raster_settings = RasterizationSettings(
    image_size=256,
    blur_radius=0.0,
    faces_per_pixel=1,
)

mesh_renderer = MeshRenderer(
    rasterizer=MeshRasterizer(
        cameras=cameras,
        raster_settings=raster_settings
    ),
    shader=HardPhongShader(
        device=device,
        cameras=cameras,
        lights=lights
    )
)

相機位置和旋轉設定

設定相機位置和旋轉是生成3D影像的重要步驟。以下是設定相機位置和旋轉的步驟：

distance = 3
elevation = 50.0

# 設定相機位置和旋轉
camera_position = np.array([distance, elevation, 0])
camera_rotation = np.array([0, 0, 0])

結合所有設定

結合所有設定，生成3D影像的完整程式碼如下：

import numpy as np

# 設定Phong渲染器
raster_settings = RasterizationSettings(
    image_size=256,
    blur_radius=0.0,
    faces_per_pixel=1,
)

phong_renderer = MeshRenderer(
    rasterizer=MeshRasterizer(
        cameras=cameras,
        raster_settings=raster_settings
    ),
    shader=HardPhongShader(
        device=device,
        cameras=cameras,
        lights=lights
    )
)

# 設定Mesh渲染器
mesh_renderer = MeshRenderer(
    rasterizer=MeshRasterizer(
        cameras=cameras,
        raster_settings=raster_settings
    ),
    shader=HardPhongShader(
        device=device,
        cameras=cameras,
        lights=lights
    )
)

# 設定相機位置和旋轉
distance = 3
elevation = 50.0
camera_position = np.array([distance, elevation, 0])
camera_rotation = np.array([0, 0, 0])

# 生成3D影像
image = phong_renderer.render(camera_position, camera_rotation)

這個程式碼結合了Phong渲染器、Mesh渲染器和相機位置、旋轉的設定，生成了一個3D影像。

生成目標影像

首先，我們需要設定觀察視角的引數，包括距離、仰角和方位角。這些引數用於計算觀察視角的變換矩陣。

azimuth = 0.0
R, T = look_at_view_transform(distance, elevation, azimuth, device=device)

接下來，我們使用渲染器生成目標影像。這裡，我們使用兩種渲染方式：一種是生成輪廓影像（silhouette），另一種是生成彩色影像（image_ref）。

silhouette = renderer_silhouette(meshes_world=cow_mesh, R=R, T=T)
image_ref = phong_renderer(meshes_world=cow_mesh, R=R, T=T)

然後，我們將渲染結果轉換為 NumPy 陣列，以便進行後續處理。

silhouette = silhouette.cpu().numpy()
image_ref = image_ref.cpu().numpy()

最後，我們使用 Matplotlib 將影像儲存為 PNG 檔案。

plt.figure(figsize=(10, 10))
plt.imshow(silhouette.squeeze()[..., 3])
plt.grid(False)
plt.savefig(os.path.join(output_dir, 'target_silhouette.png'))
plt.close()

plt.figure(figsize=(10, 10))
plt.imshow(image_ref.squeeze())
plt.savefig(os.path.join(output_dir, 'target_rgb.png'))
plt.close()

內容解密：

上述程式碼片段主要用於生成目標影像。首先，設定觀察視角的引數，然後使用渲染器生成輪廓影像和彩色影像。接下來，將渲染結果轉換為 NumPy 陣列，並使用 Matplotlib 將影像儲存為 PNG 檔案。

圖表翻譯：

此圖示為目標影像的生成流程。首先，設定觀察視角的引數，然後使用渲染器生成影像。接下來，將渲染結果轉換為 NumPy 陣列，並使用 Matplotlib 將影像儲存為 PNG 檔案。

  flowchart TD
    A[設定觀察視角] --> B[生成輪廓影像]
    B --> C[生成彩色影像]
    C --> D[轉換為 NumPy 陣列]
    D --> E[儲存為 PNG 檔案]

修改模型類別

為了改進模型的功能，我們需要修改模型類別的定義。以下是修改後的模型類別：

class Model(nn.Module):
    def __init__(self, meshes, renderer_silhouette,
                 renderer_textured, image_ref,
                 weight_silhouette, weight_texture):
        super().__init__()
        self.meshes = meshes
        self.renderer_silhouette = renderer_silhouette
        self.renderer_textured = renderer_textured
        self.weight_silhouette = weight_silhouette
        self.weight_texture = weight_texture

        # 註冊參考影像的緩衝區
        image_ref_silhouette = torch.from_numpy(
            (image_ref[..., :3].max(-1) != 1).astype(np.float32))
        self.register_buffer('image_ref_silhouette', image_ref_silhouette)

        image_ref_textured = torch.from_numpy(
            (image_ref[..., :3]).astype(np.float32))
        self.register_buffer('image_ref_textured', image_ref_textured)

        # 初始化相機位置引數
        self.camera_position = nn.Parameter(
            torch.from_numpy(np.array([3.0, 6.9, +2.5]))

    def forward(self):
        # ...

修改模型的前向傳遞

在模型的前向傳遞中，我們需要計算出alpha通道和RGB影像的損失，並將其與觀察到的影像進行比較。

def forward(self):
    # 繪製alpha通道和RGB影像
    alpha_image = self.renderer_silhouette(self.meshes)
    rgb_image = self.renderer_textured(self.meshes)

    # 計算損失
    loss_alpha = self.weight_silhouette * (alpha_image - self.image_ref_silhouette)
    loss_rgb = self.weight_texture * (rgb_image - self.image_ref_textured)

    # 計算最終損失
    loss = loss_alpha + loss_rgb

    return loss

儲存影像

最後，我們需要儲存繪製出的影像：

plt.grid(False)
plt.savefig(os.path.join(output_dir, 'target_rgb.png'))

plt.close()

這些修改使得模型可以繪製出alpha通道和RGB影像，並計算出其損失。

使用PyTorch進行3D模型渲染和最佳化

從底層實作到高階應用的全面檢視顯示，使用 PyTorch 進行 3D 模型渲染和最佳化，展現了深度學習框架在圖形學領域的強大潛力。透過可微分渲染技術，我們可以有效地計算渲染影像的梯度，並利用最佳化演算法調整模型引數，例如相機位置、光照條件和材質屬性，以逼近目標影像或場景。分析 PyTorch3D 提供的工具和函式，可以發現，它簡化了構建和訓練 3D 深度學習模型的流程，例如建立渲染器、定義光源和相機、以及計算損失函式等。然而，目前技術仍面臨一些挑戰，例如高解析度渲染的計算成本、複雜場景的建模和最佳化，以及真實感渲染的物理模擬等。對於追求高效能和高保真度的應用，需要進一步探索更高效的渲染演算法和硬體加速技術。從技術演進角度，可微分渲染與神經渲染的結合，將推動 3D 圖形學和電腦視覺的深度融合，例如自動生成 3D 內容、虛擬試衣、以及元宇宙場景的構建等。玄貓認為，隨著硬體和演算法的持續發展，可微分渲染技術將在更多領域展現其應用價值，並為 3D 內容創作和互動帶來革新。

玄貓

玄貓

技術愛好者，專注於分享程式開發、雲端技術與 AI 應用的心得體會。