【3D再構成】【ディープラーニング】NeuSコードのPytorch実装 – テストフェーズのコード解析（前編）
この論文では、2D 画像入力からオブジェクトとシーンを高忠実度で再構成するための、NeuS と呼ばれる新しい神経表面再構成方法を提案しています。NeuS では、表面を符号付き距離関数 (SDF) のゼロ次セットとして表現し、ニューラル SDF 表現をトレーニングするための新しいボリューム レンダリング方法を開発して、マスク監視なしでもより正確な表面再構成を可能にすることが提案されています。NeuS は、特に複雑な構造と自己オクルージョンを持つオブジェクトやシーンに対して、高品質のサーフェス再構成において既存の技術を上回ります。このブログ投稿では、コード実行プロセスに基づいて、テスト段階で特定の機能モジュール コードを分析します。
記事ディレクトリ
序文
NeuS ネットワークを詳細に分析する前に、最初のタスクはNeuSに必要な動作環境[ win10 リファレンス チュートリアル] を構築し、モデルのトレーニングとテストを完了することです。
このブログ投稿では、NeuS のテスト段階に関係する機能コード モジュールを分析します。
各機能モジュールのコードはブロガーが別のブログ記事で詳細に分析しており、[Win10 でのリファレンス チュートリアル] をクリックすると、そのブログ記事のディレクトリ リンクが序文に掲載されています。
ここでのコード スニペットは、exp_runner.py ファイルの train 関数部分であり、広義にはトレーニング フェーズの一部ですが、NeuS ネットワークの更新には関与せず、NeuS の段階的な検証のみを実行するため、ネットワーク、ブロガーがここに投稿しました。詳細な説明はブログ投稿にあります。
if self.iter_step % self.save_freq == 0:
self.save_checkpoint()
if self.iter_step % self.val_freq == 0:
self.validate_image()
if self.iter_step % self.val_mesh_freq == 0:
self.validate_mesh()
self.update_learning_rate()
if self.iter_step % len(image_perm) == 0:
image_perm = self.get_image_perm()
チェックポイントの保存
exp_runner.py ファイルに属する Runner クラスのメンバー メソッド。目的は、完了したフェーズ トレーニングの NeuS 重みを保存することです。
def save_checkpoint(self):
checkpoint = {
'nerf': self.nerf_outside.state_dict(), # 各深度学习网络参数权重
'sdf_network_fine': self.sdf_network.state_dict(),
'variance_network_fine': self.deviation_network.state_dict(),
'color_network_fine': self.color_network.state_dict(),
'optimizer': self.optimizer.state_dict(), # 优化器
'iter_step': self.iter_step, # 训练的次数
}
# 创建放置权重模型的文件夹
os.makedirs(os.path.join(self.base_exp_dir, 'checkpoints'), exist_ok=True)
# 保存
torch.save(checkpoint, os.path.join(self.base_exp_dir, 'checkpoints', 'ckpt_{:0>6d}.pth'.format(self.iter_step)))
検証画像
NeuS モデルのトレーニングを段階的に完了した後、画像をレンダリングして実際のトレーニング画像と比較して、モデル トレーニングの効果を検証する必要があります。
まず、gen_rays_at関数は、画像全体のレイを生成するために必要です (ダウンサンプリング後)。次に、レイ上のサンプリング ポイント (前景) の最遠点と最近点を取得し、最後にレンダラー関数を通じて必要な結果を取得します。 。
def validate_image(self, idx=-1, resolution_level=-1):
# 假设验证图像的序号小于0,随机获取一个图片序号
if idx < 0:
idx = np.random.randint(self.dataset.n_images)
print('Validate: iter: {}, camera: {}'.format(self.iter_step, idx))
if resolution_level < 0:
# 下采样倍数
resolution_level = self.validate_resolution_level
# [W, H, 3]
rays_o, rays_d = self.dataset.gen_rays_at(idx, resolution_level=resolution_level)
H, W, _ = rays_o.shape
# 按照batch_size切分,[W*H,3]=>tuple形式:W*H/batch_size个[batch_size, 3]
rays_o = rays_o.reshape(-1, 3).split(self.batch_size)
rays_d = rays_d.reshape(-1, 3).split(self.batch_size)
out_rgb_fine = []
out_normal_fine = []
for rays_o_batch, rays_d_batch in zip(rays_o, rays_d):
# 最近点和最远点
near, far = self.dataset.near_far_from_sphere(rays_o_batch, rays_d_batch)
# 背景颜色
background_rgb = torch.ones([1, 3]) if self.use_white_bkgd else None
render_out = self.renderer.render(rays_o_batch,
rays_d_batch,
near,
far,
cos_anneal_ratio=self.get_cos_anneal_ratio(),
background_rgb=background_rgb)
def feasible(key): return (key in render_out) and (render_out[key] is not None)
# 前景颜色
if feasible('color_fine'):
out_rgb_fine.append(render_out['color_fine'].detach().cpu().numpy())
# 梯度信息和采样点权重
if feasible('gradients') and feasible('weights'):
n_samples = self.renderer.n_samples + self.renderer.n_importance
# 梯度信息权重加成
normals = render_out['gradients'] * render_out['weights'][:, :n_samples, None] # [batch_size,n_samples，3]
# 采样点是否在球体内
if feasible('inside_sphere'):
# 只保留采样点在球体内的部分
normals = normals * render_out['inside_sphere'][..., None] # [batch_size,n_samples，3]
# normals是带有权重的有效梯度信息
normals = normals.sum(dim=1).detach().cpu().numpy() # [batch_size，3]
out_normal_fine.append(normals)
del render_out
gen_rays_at
データセット データ マネージャーによって定義された関数は、models/dataset.py ファイルにあります。ブロガー [ NeuS 概要] は、ブログ投稿でこのプロセスを簡単に紹介しました。
def gen_rays_at(self, img_idx, resolution_level=1):
"""
Generate rays at world space from one camera.
一个摄影机在世界空间中生成光线
"""
# 下采样倍数
l = resolution_level
# 获取2D图像上所有的像素点(下采样后的)
tx = torch.linspace(0, self.W - 1, self.W // l)
ty = torch.linspace(0, self.H - 1, self.H // l)
# 生成网格用于生成坐标
pixels_x, pixels_y = torch.meshgrid(tx, ty) # [W, H]
# 相机坐标系下的方向向量:内参(逆)×像素坐标系
p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # [W, H, 3]
p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze() # [W, H, 3]
# 单位方向向量:对方向向量做归一化处理
rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True) # [W, H, 3]
# 世界坐标系下的方向向量：外参(逆)×相机坐标系
rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze() # [W, H, 3]
# 世界坐标系下的光心位置(外参的逆对应的平移矩阵t)
rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape) # [W, H, 3]
return rays_o.transpose(0, 1), rays_v.transpose(0, 1) # [H, W, 3]
コードの実行図は次の図に示されており、この関数は Rays_o (光学中心) と Rays_v (単位方向ベクトル) を返します。
レイを生成するトレーニング プロセスと検証プロセスの違いに注意してください。トレーニング プロセスでは、batch_size ピクセルがランダムに選択され、これらのピクセルを通過するレイが生成されますが、検証プロセスでは、画像全体のすべてのピクセルが選択されます。画像全体のピクセルの光線を通過する光線を生成します。
メッシュの検証
NeuS モデルの段階的なトレーニングが完了したら、モデル トレーニングの効果を検証するために物理モデルを 3 次元で再構成することも必要です。
まず、再構成の空間範囲を定義する必要があり、次に描画アルゴリズムを通じて頂点座標と面インデックスが取得され、最後に実際の 3 次元モデル ファイルが出力されます。
def validate_mesh(self, world_space=False, resolution=64, threshold=0.0):
# 获取提取域(方体)的对角线顶点
bound_min = torch.tensor(self.dataset.object_bbox_min, dtype=torch.float32)
bound_max = torch.tensor(self.dataset.object_bbox_max, dtype=torch.float32)
# 面绘制算法获取vertices顶点坐标和triangles面索引
vertices, triangles =\
self.renderer.extract_geometry(bound_min, bound_max, resolution=resolution, threshold=threshold)
os.makedirs(os.path.join(self.base_exp_dir, 'meshes'), exist_ok=True)
if world_space:
# 再次缩放位移
vertices = vertices * self.dataset.scale_mats_np[0][0, 0] + self.dataset.scale_mats_np[0][:3, 3][None]
# 表示和操作三角网格模型
mesh = trimesh.Trimesh(vertices, triangles)
# 保存mesh模型
mesh.export(os.path.join(self.base_exp_dir, 'meshes', '{:0>8d}.ply'.format(self.iter_step)))
logging.info('End')
以下の図は、bound_min とbound_max が 3 次元再構成範囲を示すことを示しています。
なお、3D再構築の範囲と2D画像へのレンダリングの範囲は異なり、それぞれ設定があるので混同しないように注意してください。
抽出ジオメトリ
これらはすべてmodels/renderer.pyファイルの下にあり、ソースコードの作者はマトリョーシカ人形を作成しています。前者のextract_geometryはNeuSRendererクラスに属するクラスメンバメソッドで、後者は独立した関数です。
def extract_geometry(self, bound_min, bound_max, resolution, threshold=0.0):
return extract_geometry(bound_min,
bound_max,
resolution=resolution,
threshold=threshold,
query_func=lambda pts: -self.sdf_network.sdf(pts))
Marching_cubes表面レンダリング アルゴリズム リファレンス、extract_fieldsは 3 次元再構成範囲内の各点の SDF 値を取得します。
def extract_geometry(bound_min, bound_max, resolution, threshold, query_func):
print('threshold: {}'.format(threshold))
# 获取提取域多的sdf
u = extract_fields(bound_min, bound_max, resolution, query_func)
# 面绘制算法
# vertices 顶点坐标[N,3] N是根据具有情况而通过算法得出,与其他无关
# triangles 面索引[M,3] 索引指向顶点坐标数组中的对应顶点，3个顶点一个面
vertices, triangles = mcubes.marching_cubes(u, threshold)
# 提取域的对角顶点
b_max_np = bound_max.detach().cpu().numpy() # [3]
b_min_np = bound_min.detach().cpu().numpy() # [3]
# 缩小位移
vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
return vertices, triangles
抽出フィールド
この関数の機能は、3 次元再構成範囲内で適切な抽出点 (ボクセル) を取得し、各抽出点 (ボクセル) に対応する sdf 値を計算することです。
def extract_fields(bound_min, bound_max, resolution, query_func):
N = 64
# 根据提取域(方体)的对角顶点,获取提取域在各xyz轴的范围(max-min)和单位刻度((max-min)/resolution)
X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
# 初始化对应方体的sdf值
u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
with torch.no_grad():
for xi, xs in enumerate(X):
for yi, ys in enumerate(Y):
for zi, zs in enumerate(Z):
# 网格化
xx, yy, zz = torch.meshgrid(xs, ys, zs) # [N,N,N]
# [N^3,3]
pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
# 找到对应点的sdf
val = query_func(pts).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy()
# 为方体正确的赋sdf值
u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
return u
コードの実行図は下図のとおりで、オレンジ色の四角が抽出点(ボクセル)であり、分割要件に応じて、小さな抽出点(ボクセル)をより詳細に小さな抽出点(ボクセル)に分割することができます。
要約する
NeuS テスト フェーズのコードの一部 (画像をレンダリングする validate_image とモデルを再構築する validate_mesh のプロセス) をできるだけ簡単かつ詳細に紹介します。テストフェーズの残りのコードについては後で説明します。