520 lines
20 KiB
Python
520 lines
20 KiB
Python
import torch
|
||
import numpy as np
|
||
import os
|
||
import yaml
|
||
import time
|
||
from nerf_model import NeRF
|
||
from pipeline import ActiveReconstructionPolicy
|
||
from uncertainty_guide import UncertaintyGuideNeRF
|
||
import argparse
|
||
from typing import Dict, Any, List
|
||
from utils.volume_render_util import VolumeRendererUtil
|
||
import mcubes # 导入Python Marching Cubes库
|
||
import trimesh # 处理网格
|
||
from tqdm import tqdm # 进度条
|
||
|
||
class ActiveReconstruction:
|
||
"""基于NeRF不确定性引导的主动3D重建系统"""
|
||
|
||
def __init__(self, config_path: str):
|
||
"""
|
||
初始化主动重建系统
|
||
|
||
参数:
|
||
config_path: 配置文件路径
|
||
"""
|
||
# 加载配置
|
||
with open(config_path, 'r') as f:
|
||
self.config = yaml.safe_load(f)
|
||
|
||
# 设置设备
|
||
self.device = torch.device(self.config.get("device", "cuda") if torch.cuda.is_available() else "cpu")
|
||
print(f"使用设备: {self.device}")
|
||
|
||
# 创建输出目录
|
||
self.output_dir = self.config.get("output_dir", "output")
|
||
os.makedirs(self.output_dir, exist_ok=True)
|
||
|
||
# 初始化NeRF模型
|
||
self._init_nerf_model()
|
||
|
||
# 初始化视图选择策略
|
||
self.policy = ActiveReconstructionPolicy(self.config)
|
||
|
||
def _init_nerf_model(self):
|
||
"""初始化NeRF模型"""
|
||
# 从配置中获取NeRF参数
|
||
nerf_config = self.config.get("nerf", {})
|
||
model_config = {
|
||
"pos_enc_dim": nerf_config.get("pos_enc_dim", 10),
|
||
"dir_enc_dim": nerf_config.get("dir_enc_dim", 4),
|
||
"netdepth_coarse": nerf_config.get("netdepth_coarse", 8),
|
||
"netwidth_coarse": nerf_config.get("netwidth_coarse", 256),
|
||
"netdepth_fine": nerf_config.get("netdepth_fine", 8),
|
||
"netwidth_fine": nerf_config.get("netwidth_fine", 256),
|
||
"skips": nerf_config.get("skips", [4]),
|
||
"use_viewdirs": nerf_config.get("use_viewdirs", True)
|
||
}
|
||
self.nerf_model = NeRF(model_config).to(self.device)
|
||
|
||
def _generate_rays(self,
|
||
poses: torch.Tensor,
|
||
H: int,
|
||
W: int,
|
||
focal: float) -> tuple:
|
||
"""
|
||
为每个相机位姿生成光线
|
||
|
||
参数:
|
||
poses: 相机位姿 [N, 4, 4]
|
||
H: 图像高度
|
||
W: 图像宽度
|
||
focal: 焦距
|
||
|
||
返回:
|
||
rays_o: 光线起点 [N, H*W, 3]
|
||
rays_d: 光线方向 [N, H*W, 3]
|
||
"""
|
||
# 创建像素坐标网格
|
||
i, j = torch.meshgrid(
|
||
torch.linspace(0, W-1, W),
|
||
torch.linspace(0, H-1, H),
|
||
indexing='ij'
|
||
)
|
||
i = i.t() # [H, W]
|
||
j = j.t() # [H, W]
|
||
|
||
# 转换为相机坐标系中的方向
|
||
dirs = torch.stack([
|
||
(i - W * 0.5) / focal,
|
||
-(j - H * 0.5) / focal,
|
||
-torch.ones_like(i)
|
||
], dim=-1) # [H, W, 3]
|
||
|
||
# 为每个位姿生成光线
|
||
rays_o_list = []
|
||
rays_d_list = []
|
||
|
||
for pose in poses:
|
||
# 转换光线方向到世界坐标系
|
||
rays_d = torch.sum(dirs[..., None, :] * pose[:3, :3], dim=-1) # [H, W, 3]
|
||
|
||
# 设置光线原点
|
||
rays_o = pose[:3, -1].expand(rays_d.shape) # [H, W, 3]
|
||
|
||
# 展平为批处理格式
|
||
rays_o = rays_o.reshape(-1, 3) # [H*W, 3]
|
||
rays_d = rays_d.reshape(-1, 3) # [H*W, 3]
|
||
|
||
rays_o_list.append(rays_o)
|
||
rays_d_list.append(rays_d)
|
||
|
||
# 组合所有位姿的光线
|
||
rays_o_all = torch.stack(rays_o_list, dim=0) # [N, H*W, 3]
|
||
rays_d_all = torch.stack(rays_d_list, dim=0) # [N, H*W, 3]
|
||
|
||
return rays_o_all, rays_d_all
|
||
|
||
def _sample_pixel_batch(self,
|
||
images: torch.Tensor,
|
||
rays_o: torch.Tensor,
|
||
rays_d: torch.Tensor,
|
||
batch_size: int) -> tuple:
|
||
"""
|
||
随机采样像素批次
|
||
|
||
参数:
|
||
images: 图像数据 [N, H, W, 3]
|
||
rays_o: 光线起点 [N, H*W, 3]
|
||
rays_d: 光线方向 [N, H*W, 3]
|
||
batch_size: 批次大小
|
||
|
||
返回:
|
||
sampled_rays_o: 采样的光线起点 [batch_size, 3]
|
||
sampled_rays_d: 采样的光线方向 [batch_size, 3]
|
||
sampled_pixels: 采样的像素值 [batch_size, 3]
|
||
"""
|
||
# 获取图像形状
|
||
N = images.shape[0]
|
||
H = images.shape[1]
|
||
W = images.shape[2]
|
||
total_rays = N * H * W
|
||
|
||
# 将图像展平
|
||
pixels = images.reshape(N, -1, 3) # [N, H*W, 3]
|
||
|
||
# 随机选择批次
|
||
indices = torch.randint(0, total_rays, size=(batch_size,))
|
||
img_indices = indices // (H * W)
|
||
pixel_indices = indices % (H * W)
|
||
|
||
# 采样光线和像素
|
||
sampled_rays_o = torch.stack([rays_o[i, j] for i, j in zip(img_indices, pixel_indices)])
|
||
sampled_rays_d = torch.stack([rays_d[i, j] for i, j in zip(img_indices, pixel_indices)])
|
||
sampled_pixels = torch.stack([pixels[i, j] for i, j in zip(img_indices, pixel_indices)])
|
||
|
||
return sampled_rays_o, sampled_rays_d, sampled_pixels
|
||
|
||
def train_nerf(self,
|
||
images: torch.Tensor,
|
||
poses: torch.Tensor,
|
||
epochs: int = 5000,
|
||
batch_size: int = 4096,
|
||
lr: float = 5e-4,
|
||
start_from_model=None) -> float:
|
||
"""
|
||
训练NeRF模型
|
||
|
||
参数:
|
||
images: 图像数据 [N, H, W, 3]
|
||
poses: 相机位姿 [N, 4, 4]
|
||
epochs: 训练轮数
|
||
batch_size: 批量大小
|
||
lr: 学习率
|
||
start_from_model: 可选的初始模型状态
|
||
|
||
返回:
|
||
final_loss: 最终损失值
|
||
"""
|
||
print(f"开始训练NeRF模型，使用{len(images)}张图像...")
|
||
|
||
# 获取图像和采样参数
|
||
H, W = images.shape[1], images.shape[2]
|
||
sampling_config = self.config.get("sampling", {})
|
||
camera_config = self.config.get("camera", {})
|
||
focal = camera_config.get("focal", 1000.0)
|
||
near = camera_config.get("near", 2.0)
|
||
far = camera_config.get("far", 6.0)
|
||
coarse_samples = sampling_config.get("coarse_samples", 64)
|
||
fine_samples = sampling_config.get("fine_samples", 128)
|
||
perturb = sampling_config.get("perturb", True)
|
||
|
||
# 如果提供了初始模型，使用它
|
||
if start_from_model is not None:
|
||
print("从现有模型初始化权重")
|
||
self.nerf_model.load_state_dict(start_from_model.state_dict())
|
||
|
||
# 设置优化器和损失函数
|
||
optimizer = torch.optim.Adam(self.nerf_model.parameters(), lr=lr)
|
||
mse_loss = torch.nn.MSELoss()
|
||
|
||
# 将模型设置为训练模式
|
||
self.nerf_model.train()
|
||
|
||
# 为所有图像生成光线（预计算光线可以加速训练）
|
||
rays_o, rays_d = self._generate_rays(poses, H, W, focal)
|
||
rays_o = rays_o.to(self.device)
|
||
rays_d = rays_d.to(self.device)
|
||
images = images.to(self.device)
|
||
|
||
# 训练循环
|
||
best_loss = float('inf')
|
||
for epoch in range(epochs):
|
||
# 随机采样一批光线
|
||
batch_rays_o, batch_rays_d, target_pixels = self._sample_pixel_batch(
|
||
images, rays_o, rays_d, batch_size)
|
||
|
||
# 光线方向归一化
|
||
batch_rays_d = torch.nn.functional.normalize(batch_rays_d, dim=-1)
|
||
|
||
# 创建近平面和远平面张量
|
||
near_tensor = torch.ones_like(batch_rays_o[..., 0]) * near
|
||
far_tensor = torch.ones_like(batch_rays_o[..., 0]) * far
|
||
|
||
# 使用体积渲染进行前向传播
|
||
# 首先进行粗采样渲染
|
||
optimizer.zero_grad()
|
||
|
||
# 体积渲染
|
||
rgb_map, _, _, _ = VolumeRendererUtil.render_rays(
|
||
self.nerf_model,
|
||
batch_rays_o,
|
||
batch_rays_d,
|
||
near_tensor,
|
||
far_tensor,
|
||
coarse_samples,
|
||
fine_samples,
|
||
perturb
|
||
)
|
||
|
||
# 计算损失并反向传播
|
||
loss = mse_loss(rgb_map, target_pixels)
|
||
loss.backward()
|
||
optimizer.step()
|
||
|
||
# 输出训练进度
|
||
if (epoch + 1) % 100 == 0:
|
||
psnr = -10.0 * torch.log10(loss)
|
||
print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.6f}, PSNR: {psnr.item():.2f}")
|
||
|
||
# 保存最佳模型
|
||
if loss.item() < best_loss:
|
||
best_loss = loss.item()
|
||
torch.save(self.nerf_model.state_dict(), os.path.join(self.output_dir, "best_model.pth"))
|
||
|
||
# 加载最佳模型
|
||
self.nerf_model.load_state_dict(torch.load(os.path.join(self.output_dir, "best_model.pth")))
|
||
|
||
print(f"NeRF模型训练完成，最终损失: {best_loss:.6f}")
|
||
return best_loss
|
||
|
||
def extract_mesh(self, output_path: str, resolution: int = 256, threshold: float = 50.0, bound: float = 2.0):
|
||
"""
|
||
从NeRF模型中提取3D网格，使用Marching Cubes算法
|
||
|
||
参数:
|
||
output_path: 输出路径
|
||
resolution: 体素网格分辨率
|
||
threshold: 密度阈值，用于确定表面位置
|
||
bound: 体素网格边界大小
|
||
"""
|
||
print(f"从NeRF提取3D网格，分辨率: {resolution}...")
|
||
|
||
# 设置网格提取参数
|
||
self.nerf_model.eval() # 设置为评估模式
|
||
|
||
# 定义采样网格
|
||
x = torch.linspace(-bound, bound, resolution)
|
||
y = torch.linspace(-bound, bound, resolution)
|
||
z = torch.linspace(-bound, bound, resolution)
|
||
|
||
# 创建采样点坐标网格
|
||
xx, yy, zz = torch.meshgrid(x, y, z, indexing='ij')
|
||
|
||
# 准备查询点
|
||
points = torch.stack([xx, yy, zz], dim=-1).reshape(-1, 3).to(self.device)
|
||
|
||
# 创建密度场
|
||
print("正在计算体积密度场...")
|
||
density_field = torch.zeros((resolution, resolution, resolution))
|
||
|
||
# 分批处理以避免显存溢出
|
||
batch_size = 4096 # 根据GPU内存调整
|
||
with torch.no_grad():
|
||
for i in tqdm(range(0, points.shape[0], batch_size)):
|
||
# 获取当前批次的点
|
||
batch_points = points[i:i+batch_size]
|
||
|
||
# 计算密度 - 使用固定方向（这里使用+z方向）
|
||
# 注意：在NeRF中，密度不依赖于视角方向，只有颜色依赖视角
|
||
fixed_dirs = torch.zeros_like(batch_points)
|
||
fixed_dirs[..., 2] = 1.0 # 设置为+z方向
|
||
|
||
# 使用fine网络进行推理
|
||
sigma, _ = self.nerf_model(batch_points, fixed_dirs, coarse=False)
|
||
|
||
# 更新密度场
|
||
batch_indices = torch.arange(i, min(i+batch_size, points.shape[0]))
|
||
xyz_indices = torch.stack([
|
||
(points[batch_indices, 0] + bound) / (2 * bound) * (resolution - 1),
|
||
(points[batch_indices, 1] + bound) / (2 * bound) * (resolution - 1),
|
||
(points[batch_indices, 2] + bound) / (2 * bound) * (resolution - 1)
|
||
], dim=-1).long()
|
||
|
||
for j, (xi, yi, zi) in enumerate(xyz_indices):
|
||
density_field[xi, yi, zi] = sigma[j].cpu()
|
||
|
||
# 使用Marching Cubes提取网格
|
||
print("使用Marching Cubes提取网格...")
|
||
density_field_np = density_field.cpu().numpy()
|
||
vertices, triangles = mcubes.marching_cubes(density_field_np, threshold)
|
||
|
||
# 转换为正确的坐标系（视场的[-bound, bound]范围）
|
||
vertices = vertices / (resolution - 1) * (2 * bound) - bound
|
||
|
||
# 创建trimesh对象
|
||
mesh = trimesh.Trimesh(vertices=vertices, faces=triangles)
|
||
|
||
# 保存网格
|
||
mesh.export(output_path)
|
||
|
||
print(f"网格提取完成，保存至: {output_path}")
|
||
print(f"网格统计: {len(vertices)}个顶点, {len(triangles)}个三角面")
|
||
|
||
return mesh
|
||
|
||
def evaluate_reconstruction(self,
|
||
gt_mesh_path: str = None) -> Dict[str, float]:
|
||
"""
|
||
评估重建质量
|
||
|
||
参数:
|
||
gt_mesh_path: 真实网格路径（如果有）
|
||
|
||
返回:
|
||
metrics: 评估指标，如F-score
|
||
"""
|
||
if gt_mesh_path is None:
|
||
print("没有提供真实网格，跳过评估")
|
||
return {}
|
||
|
||
print("评估重建质量...")
|
||
|
||
# 在实际实现中，这里应该有评估重建质量的代码
|
||
# 通常使用F-score、Chamfer距离等指标
|
||
|
||
# 为了简化，我们返回模拟的指标
|
||
metrics = {
|
||
"f_score": 0.85,
|
||
"precision": 0.87,
|
||
"recall": 0.83
|
||
}
|
||
|
||
print(f"评估结果: F-score={metrics['f_score']:.4f}, "
|
||
f"精确率={metrics['precision']:.4f}, 召回率={metrics['recall']:.4f}")
|
||
|
||
return metrics
|
||
|
||
def run_active_reconstruction(self,
|
||
initial_poses: np.ndarray,
|
||
initial_images: torch.Tensor = None,
|
||
max_iterations: int = 3) -> List[np.ndarray]:
|
||
"""
|
||
运行主动重建过程
|
||
|
||
参数:
|
||
initial_poses: 初始相机位姿
|
||
initial_images: 初始图像（如果有）
|
||
max_iterations: 最大迭代次数
|
||
|
||
返回:
|
||
selected_poses: 所有选定的相机位姿
|
||
"""
|
||
print("开始主动重建过程...")
|
||
|
||
# 初始训练，使用初始视图
|
||
if initial_images is None:
|
||
initial_images = self._simulate_image_capture(initial_poses)
|
||
|
||
# 使用初始图像训练模型
|
||
self.train_nerf(
|
||
initial_images,
|
||
torch.from_numpy(initial_poses).float().to(self.device),
|
||
epochs=self.config.get("reconstruction", {}).get("epochs_per_iteration", 2000)
|
||
)
|
||
|
||
# 保存初始模型
|
||
initial_model_path = os.path.join(self.output_dir, "initial_model.pth")
|
||
torch.save(self.nerf_model.state_dict(), initial_model_path)
|
||
initial_model = self.nerf_model.state_dict()
|
||
|
||
all_poses = initial_poses.copy()
|
||
current_poses = initial_poses.copy()
|
||
all_images = initial_images.clone()
|
||
|
||
# 提取初始网格
|
||
initial_mesh_path = os.path.join(self.output_dir, "initial_mesh.obj")
|
||
self.extract_mesh(
|
||
initial_mesh_path,
|
||
resolution=self.config.get("reconstruction", {}).get("mesh_resolution", 256)
|
||
)
|
||
|
||
# 迭代执行主动重建
|
||
for iteration in range(max_iterations):
|
||
print(f"\n开始迭代 {iteration+1}/{max_iterations}")
|
||
|
||
# 选择下一批视角
|
||
next_views = self.policy.select_next_views(self.nerf_model, current_poses)
|
||
print(f"选择了 {len(next_views)} 个新视角")
|
||
|
||
# 采集新视角的图像
|
||
new_images = self._simulate_image_capture(next_views)
|
||
|
||
# 将新选择的视角添加到当前位姿和图像中
|
||
current_poses = np.concatenate([current_poses, next_views], axis=0)
|
||
all_poses = np.concatenate([all_poses, next_views], axis=0)
|
||
all_images = torch.cat([all_images, new_images], dim=0)
|
||
|
||
# 按照作者的描述，我们从初始模型重新初始化，而不是继续训练
|
||
# "After selecting additional images, we initialize the network with the model from the initialization step and refine the model further with the updated training set."
|
||
# 因此，我们先加载初始模型，然后用扩展的数据集重新训练
|
||
self.nerf_model.load_state_dict(torch.load(initial_model_path))
|
||
|
||
# 用扩展的数据集重新训练模型
|
||
self.train_nerf(
|
||
all_images,
|
||
torch.from_numpy(current_poses).float().to(self.device),
|
||
epochs=self.config.get("reconstruction", {}).get("epochs_per_iteration", 2000)
|
||
)
|
||
|
||
# 每次迭代后提取网格，以便观察重建质量的改进
|
||
iter_mesh_path = os.path.join(self.output_dir, f"mesh_iter_{iteration+1}.obj")
|
||
self.extract_mesh(
|
||
iter_mesh_path,
|
||
resolution=self.config.get("reconstruction", {}).get("mesh_resolution", 256)
|
||
)
|
||
|
||
# 提取最终的3D网格
|
||
output_mesh_path = os.path.join(self.output_dir, "final_mesh.obj")
|
||
self.extract_mesh(
|
||
output_mesh_path,
|
||
resolution=self.config.get("reconstruction", {}).get("mesh_resolution", 256)
|
||
)
|
||
|
||
# 评估重建质量
|
||
self.evaluate_reconstruction()
|
||
|
||
print("主动重建过程完成")
|
||
return all_poses
|
||
|
||
def _simulate_image_capture(self, poses: np.ndarray) -> torch.Tensor:
|
||
"""
|
||
模拟图像采集过程（实际系统中应该从相机或数据集获取）
|
||
|
||
参数:
|
||
poses: 相机位姿
|
||
|
||
返回:
|
||
images: 模拟的图像
|
||
"""
|
||
# 模拟图像大小
|
||
camera_config = self.config.get("camera", {})
|
||
H, W = camera_config.get("height", 800), camera_config.get("width", 800)
|
||
|
||
# 创建随机图像（实际应来自相机或渲染）
|
||
images = torch.rand(len(poses), H, W, 3, device=self.device)
|
||
|
||
return images
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="基于NeRF不确定性的主动3D重建")
|
||
parser.add_argument("--config", type=str, default="nbv_config.yaml", help="配置文件路径")
|
||
parser.add_argument("--synthetic", action="store_true", help="使用合成数据集")
|
||
args = parser.parse_args()
|
||
|
||
# 创建主动重建系统
|
||
reconstruction = ActiveReconstruction(args.config)
|
||
|
||
# 初始化一些相机位姿（通常来自中心圆环）
|
||
# 根据配置获取初始位姿数量
|
||
config = yaml.safe_load(open(args.config, 'r'))
|
||
initial_view_count = config.get("reconstruction", {}).get("initial_view_count", 15)
|
||
|
||
# 根据数据集类型调整初始视图数量
|
||
if args.synthetic:
|
||
initial_view_count = min(initial_view_count, 6) # 合成数据使用6个初始视图
|
||
print(f"使用合成数据集，初始视图数量: {initial_view_count}")
|
||
else:
|
||
print(f"使用真实数据集，初始视图数量: {initial_view_count}")
|
||
|
||
# 获取中间圆环上的相机位姿
|
||
# 假设poses是按圆环组织的，我们选择中间圆环的部分位姿
|
||
middle_circle_index = config.get("view_selection", {}).get("n_circles", 5) // 2
|
||
poses_per_circle = config.get("view_selection", {}).get("n_poses_per_circle", 30)
|
||
|
||
# 等距选择初始位姿
|
||
start_index = middle_circle_index * poses_per_circle
|
||
step = poses_per_circle // initial_view_count
|
||
initial_pose_indices = [start_index + i * step for i in range(initial_view_count)]
|
||
initial_poses = reconstruction.policy.poses[initial_pose_indices]
|
||
|
||
# 运行主动重建
|
||
selected_poses = reconstruction.run_active_reconstruction(
|
||
initial_poses,
|
||
max_iterations=config.get("reconstruction", {}).get("max_iterations", 3)
|
||
)
|
||
|
||
print(f"主动重建完成，共选择了{len(selected_poses)}个相机位姿")
|
||
|
||
if __name__ == "__main__":
|
||
main() |