import torch import numpy as np import open3d as o3d from tqdm import tqdm from scipy.spatial import KDTree from skimage.metrics import dice_score class MultiViewConsistency: def init(self, rgb_paths, mask_paths, depth_paths): self.device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”) # 初始化模型参数 self.point_cloud = self.init_point_cloud() self.optimizer = torch.optim.Adam([self.point_cloud], lr=0.001) self.load_data(rgb_paths, mask_paths, depth_paths) def load_data(self, rgb_paths, mask_paths, depth_paths): # 数据加载与预处理 self.views = [] for i in range(len(rgb_paths)): view_data = { 'rgb': self.preprocess_image(rgb_paths[i]), 'mask': self.load_mask(mask_paths[i]), 'depth': self.load_depth(depth_paths[i]), 'pose': self.generate_camera_pose(i) # 假设的相机位姿生成 } self.views.append(view_data) def multi_view_dice_loss(self, pred_masks, gt_masks): """多视角Dice损失计算""" total_loss = 0.0 for pred, gt in zip(pred_masks, gt_masks): intersection = torch.sum(pred * gt) union = torch.sum(pred) + torch.sum(gt) total_loss += 1 - (2. * intersection) / (union + 1e-8) return total_loss / len(pred_masks) def project_pointcloud_to_views(self): """将点云投影到各视角生成预测mask""" projected_masks = [] for view in self.views: # 使用深度信息和相机位姿进行投影 mask = self.render_projection(view['pose'], view['depth']) projected_masks.append(mask) return projected_masks def optimize_consistency(self, epochs=100): """多视角一致性优化循环""" best_loss = float('inf') for epoch in tqdm(range(epochs), desc="Optimizing Consistency"): self.optimizer.zero_grad() # 生成预测mask pred_masks = self.project_pointcloud_to_views() gt_masks = [view['mask'] for view in self.views] # 计算组合损失 dice_loss = self.multi_view_dice_loss(pred_masks, gt_masks) geometric_loss = self.calculate_geometric_constraints() total_loss = dice_loss + 0.5 * geometric_loss # 权重可调 # 反向传播优化 total_loss.backward() self.optimizer.step() # 记录最佳结果 if total_loss.item() < best_loss: best_loss = total_loss.item() self.save_best_pointcloud() def calculate_geometric_constraints(self): """几何约束项（法线一致性、曲率平滑性）""" # 使用Open3D计算点云法线 pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(self.point_cloud.detach().cpu().numpy()) pcd.estimate_normals() # 构建KDTree进行邻近点查询 kdtree = KDTree(np.asarray(pcd.points)) distances, indices = kdtree.query(np.asarray(pcd.points), k=5) # 计算法线一致性损失 normals = torch.tensor(np.asarray(pcd.normals)).to(self.device) normal_loss = torch.mean(1 - torch.abs(torch.sum(normals * normals[indices], dim=1))) return normal_loss def save_best_pointcloud(self, output_path="best_pointcloud.ply"): # 保存优化后的点云 pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(self.point_cloud.detach().cpu().numpy()) o3d.io.write_point_cloud(output_path, pcd)以上代码你生成的吗？