diff --git a/MASK_TO_COORDINATES_README.md b/MASK_TO_COORDINATES_README.md deleted file mode 100644 index 0a9a514..0000000 --- a/MASK_TO_COORDINATES_README.md +++ /dev/null @@ -1,109 +0,0 @@ -# VVL Mask转坐标节点 - -## 功能描述 - -`MaskToCoordinates` 节点可以从输入的 mask 图像中自动提取坐标信息，生成符合 SAM2 分割节点要求的 `coordinates_positive` 和 `coordinates_negative` 参数。 - -- **白色区域** (值 >= threshold) → `coordinates_positive` (正样本点) -- **黑色区域** (值 < threshold) → `coordinates_negative` (负样本点) - -## 输入参数 - -### 必需参数 -- **mask** (MASK): 输入的 mask 图像 -- **sample_method** (选择): 采样方法 - - `random`: 随机采样 - - `grid`: 网格采样 - - `contour`: 轮廓采样 -- **positive_points** (INT): 正样本点数量 (1-1000, 默认: 10) -- **negative_points** (INT): 负样本点数量 (0-1000, 默认: 10) -- **threshold** (FLOAT): 二值化阈值 (0.0-1.0, 默认: 0.5) - -### 可选参数 -- **min_distance** (INT): 点之间最小距离 (1-100, 默认: 10) -- **edge_margin** (INT): 边缘留白像素 (0-50, 默认: 5) - -## 输出结果 - -- **coordinates_positive** (STRING): 正坐标 JSON 字符串，格式: `[{"x": x1, "y": y1}, {"x": x2, "y": y2}, ...]` -- **coordinates_negative** (STRING): 负坐标 JSON 字符串，格式同上 -- **visualization** (IMAGE): 可视化图像，显示采样点位置 - - 绿色圆点: 正样本点 - - 红色圆点: 负样本点 - -## 使用示例 - -### 与 SAM2 分割节点连接 - -``` -[Mask输入] → [MaskToCoordinates] → [Sam2Segmentation] - ↓ - coordinates_positive - coordinates_negative -``` - -### 典型工作流程 - -1. 准备或生成 mask 图像 -2. 连接到 `MaskToCoordinates` 节点 -3. 调整采样参数 (点数量、采样方法等) -4. 将生成的坐标输出连接到 `Sam2Segmentation` 节点的对应输入 -5. 查看可视化输出确认采样效果 - -## 采样方法说明 - -### Random (随机采样) -- 在有效区域内随机选择点 -- 支持最小距离约束，避免点过于密集 -- 适用于大多数情况 - -### Grid (网格采样) -- 按网格模式均匀分布采样点 -- 确保点分布均匀 -- 适用于需要规整分布的场景 - -### Contour (轮廓采样) -- 沿着区域轮廓边缘采样 -- 优先选择最大轮廓 -- 适用于边界敏感的分割任务 - -## 技术细节 - -### 坐标格式 -生成的坐标完全兼容 SAM2 分割节点，格式为: -```json -[ - {"x": 100, "y": 150}, - {"x": 200, "y": 250}, - ... -] -``` - -### 边缘处理 -- `edge_margin` 参数可设置边缘留白，避免在图像边缘采样 -- 防止分割时边界效应 - -### 性能优化 -- 支持批处理 (自动处理第一帧) -- 内存高效的采样算法 -- 快速的轮廓检测 - -## 注意事项 - -1. **Mask 格式**: 支持标准的 ComfyUI MASK 格式 -2. **坐标范围**: 生成的坐标在 mask 图像尺寸范围内 -3. **点数量**: 实际生成的点数可能少于设定值（当有效区域不足时） -4. **阈值设置**: 建议根据具体 mask 特征调整 threshold 值 - -## 故障排除 - -- **生成点数不足**: 检查 mask 是否有足够的有效区域 -- **点分布不理想**: 尝试不同的采样方法或调整参数 -- **边缘点过多**: 增大 `edge_margin` 值 -- **坐标格式错误**: 确保输出连接到正确的 SAM2 节点输入 - -## 版本信息 - -- 版本: 1.0.0 -- 兼容: ComfyUI SAM2 节点 -- 依赖: PyTorch, OpenCV, NumPy \ No newline at end of file diff --git a/__init__.py b/__init__.py index 53169bf..02a413c 100644 --- a/__init__.py +++ b/__init__.py @@ -4,35 +4,44 @@ ComfyUI VVL Video Camera Advanced """ from .comfyui_vggt_nodes import VGGTMultiInputNode +from .comfyui_vggt_native_nodes import VGGTNativeFullOutputNode +from .comfyui_vggt_mask_processor import VGGTMaskProcessorNode from .vggt_model_loader import VVLVGGTLoader -from .glb_point_cloud_processor import * -from .mask_to_coordinates import MaskToCoordinates -from .json_merger import JSONMerger, UESceneGenerator +# 以下节点已迁移到 ComfyUI_VVL_DataProcessing 插件 +# from .glb_point_cloud_processor import * +# from .mask_to_coordinates import MaskToCoordinates +# from .json_merger import JSONMerger, UESceneGenerator # 节点类映射 NODE_CLASS_MAPPINGS = { "VVLVGGTLoader": VVLVGGTLoader, - "VGGTMultiInputNode": VGGTMultiInputNode, - "GLBPointCloudProcessor": GLBPointCloudProcessor, - "GLBPointCloudBounds": GLBPointCloudBounds, - "GLBPointCloudOriginAdjuster": GLBPointCloudOriginAdjuster, - "GLBPointCloudDensityFilter": GLBPointCloudDensityFilter, - "MaskToCoordinates": MaskToCoordinates, - "JSONMerger": JSONMerger, - "UESceneGenerator": UESceneGenerator, + "VGGTMultiInputNode": VGGTMultiInputNode, + "VGGTNativeFullOutputNode": VGGTNativeFullOutputNode, + "VGGTMaskProcessorNode": VGGTMaskProcessorNode, + # 以下节点已迁移到 ComfyUI_VVL_DataProcessing 插件 + # "GLBPointCloudProcessor": GLBPointCloudProcessor, + # "GLBPointCloudBounds": GLBPointCloudBounds, + # "GLBPointCloudOriginAdjuster": GLBPointCloudOriginAdjuster, + # "GLBPointCloudDensityFilter": GLBPointCloudDensityFilter, + # "MaskToCoordinates": MaskToCoordinates, + # "JSONMerger": JSONMerger, + # "UESceneGenerator": UESceneGenerator, } # 节点显示名称映射 NODE_DISPLAY_NAME_MAPPINGS = { "VVLVGGTLoader": "VVL VGGT 模型加载器", "VGGTMultiInputNode": "VVL VGGT 多输入重建", - "GLBPointCloudProcessor": "VVL GLB 点云黑色点清理", - "GLBPointCloudBounds": "VVL GLB 点云包围盒计算", - "GLBPointCloudOriginAdjuster": "VVL GLB 点云原点调整", - "GLBPointCloudDensityFilter": "VVL GLB 点云密度过滤", - "MaskToCoordinates": "VVL Mask转坐标", - "JSONMerger": "VVL JSON数据合并器", - "UESceneGenerator": "VVL UE场景生成器", + "VGGTNativeFullOutputNode": "🔬 VGGT 原生完整输出", + "VGGTMaskProcessorNode": "🎭 VGGT Mask分割处理器", + # 以下节点已迁移到 ComfyUI_VVL_DataProcessing 插件 + # "GLBPointCloudProcessor": "VVL GLB 点云黑色点清理", + # "GLBPointCloudBounds": "VVL GLB 点云包围盒计算", + # "GLBPointCloudOriginAdjuster": "VVL GLB 点云原点调整", + # "GLBPointCloudDensityFilter": "VVL GLB 点云密度过滤", + # "MaskToCoordinates": "VVL Mask转坐标", + # "JSONMerger": "VVL JSON数据合并器", + # "UESceneGenerator": "VVL UE场景生成器", } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/comfyui_vggt_mask_processor.py b/comfyui_vggt_mask_processor.py new file mode 100644 index 0000000..2e56b64 --- /dev/null +++ b/comfyui_vggt_mask_processor.py @@ -0,0 +1,1617 @@ +# VGGT Mask处理节点文件 + +import os +import json +import tempfile +from typing import List, Any, Dict, Tuple +import logging +from collections import defaultdict +import time + +import cv2 +import numpy as np +import torch +from PIL import Image + +# 导入ComfyUI的路径管理 +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + folder_paths = None + FOLDER_PATHS_AVAILABLE = False + +# 导入原节点的工具函数 +try: + from .comfyui_vggt_nodes import ( + predictions_to_glb, + TRIMESH_AVAILABLE, + MATPLOTLIB_AVAILABLE, + SCIPY_AVAILABLE + ) + VGGT_UTILS_AVAILABLE = True +except ImportError: + predictions_to_glb = None + TRIMESH_AVAILABLE = False + MATPLOTLIB_AVAILABLE = False + SCIPY_AVAILABLE = False + VGGT_UTILS_AVAILABLE = False + +# 配置日志 +logger = logging.getLogger('vvl_vggt_mask_processor') +logger.setLevel(logging.DEBUG) # 设置为DEBUG级别以查看详细信息 + +# ----------------------------------------------------------------------------- +# Mask处理核心函数 +# ----------------------------------------------------------------------------- + +def process_tracks_with_masks(tracks_data: Dict, mask_sequence: List[np.ndarray], + confidence_threshold: float = 0.5) -> Dict[int, List[int]]: + """ + 使用tracks数据和mask序列进行点分割 + + Args: + tracks_data: VGGT的tracks数据，包含track_list, visibility_score, confidence_score + mask_sequence: mask图像序列，每张图像包含不同的分割ID + confidence_threshold: 置信度阈值，低于此值的点会被忽略 + + Returns: + point_labels: {point_id: [object_ids]} 每个点投票得到的物体ID列表 + """ + logger.info(f"开始处理tracks数据进行mask分割，使用固定mask阈值: 128") + + # 提取tracks数据 + track_list = tracks_data.get('track_list', []) + visibility_score = tracks_data.get('visibility_score', []) + confidence_score = tracks_data.get('confidence_score', []) + + if track_list is None or (isinstance(track_list, (list, tuple)) and len(track_list) == 0) or (isinstance(track_list, np.ndarray) and track_list.size == 0): + logger.warning("No track_list found in tracks_data") + return {} + + # 转换为numpy数组 - 更强力的tensor检测 + logger.info(f"track_list类型: {type(track_list)}, 是否有cpu方法: {hasattr(track_list, 'cpu')}") + + if hasattr(track_list, 'cpu') and hasattr(track_list, 'numpy'): + # PyTorch tensor + track_coords = track_list.cpu().numpy() + logger.info(f"从tensor转换track_coords, shape: {track_coords.shape}") + elif isinstance(track_list, np.ndarray): + track_coords = track_list + logger.info(f"使用numpy track_coords, shape: {track_coords.shape}") + elif hasattr(track_list, 'cpu'): # 可能是tensor但没有numpy方法 + track_coords = track_list.cpu().detach().numpy() + logger.info(f"强制转换tensor track_coords, shape: {track_coords.shape}") + else: + try: + track_coords = np.array(track_list) + logger.info(f"数组转换track_coords, shape: {track_coords.shape}") + except Exception as e: + logger.error(f"track_list转换失败: {e}, type: {type(track_list)}") + # 尝试递归处理列表中的tensor元素 + if isinstance(track_list, (list, tuple)) and len(track_list) > 0: + converted_list = [] + for item in track_list: + if hasattr(item, 'cpu'): + converted_list.append(item.cpu().numpy()) + else: + converted_list.append(item) + track_coords = np.array(converted_list) + logger.info(f"递归转换track_coords成功, shape: {track_coords.shape}") + else: + raise ValueError(f"无法转换track_list, type: {type(track_list)}") + + if hasattr(visibility_score, 'cpu'): + vis_scores = visibility_score.cpu().numpy() + elif isinstance(visibility_score, np.ndarray): + vis_scores = visibility_score + else: + vis_scores = np.array(visibility_score) if visibility_score is not None else None + + if hasattr(confidence_score, 'cpu'): + conf_scores = confidence_score.cpu().numpy() + elif isinstance(confidence_score, np.ndarray): + conf_scores = confidence_score + else: + conf_scores = np.array(confidence_score) if confidence_score is not None else None + + # 重塑vis_scores和conf_scores以匹配track_coords - 需要在重塑track_coords之后 + pass # 先重塑track_coords，然后再处理score + + logger.info(f"Track坐标形状: {track_coords.shape}") + logger.info(f"Mask序列长度: {len(mask_sequence)}") + + # 处理不同的track_coords格式 + if track_coords.ndim == 5: # (batch, 1, frames, points, 2) + logger.info("检测到5D格式，重塑为标准格式") + if track_coords.shape[1] == 1: # 去掉多余的维度 + track_coords = track_coords[:, 0] # (batch, frames, points, 2) + # 合并batch和points维度: (batch*points, frames, 2) + batch_size, n_frames, n_points_per_batch, coord_dim = track_coords.shape + track_coords = track_coords.reshape(-1, n_frames, coord_dim) # (batch*points, frames, 2) + # 转置为 (points, frames, 2) + track_coords = track_coords.swapaxes(0, 1).swapaxes(0, 1) + logger.info(f"重塑后形状: {track_coords.shape}") + elif track_coords.ndim == 4: # (batch, frames, points, 2) 或 (1, frames, points, 2) + if track_coords.shape[0] == 1: + track_coords = track_coords.squeeze(0) # 去掉batch维度: (frames, points, 2) + track_coords = track_coords.swapaxes(0, 1) # 转置为 (points, frames, 2) + else: + # 多个batch的情况 + batch_size, n_frames, n_points_per_batch, coord_dim = track_coords.shape + track_coords = track_coords.reshape(-1, n_frames, coord_dim) + elif track_coords.ndim == 3: # 已经是 (points, frames, 2) 或 (frames, points, 2) + if track_coords.shape[2] == 2: # (frames, points, 2) -> (points, frames, 2) + track_coords = track_coords.swapaxes(0, 1) + + # 确保最终格式是 (N_points, N_frames, 2) + if track_coords.shape[2] != 2: + logger.error(f"意外的坐标维度: {track_coords.shape}") + raise ValueError(f"期望最后一维为2（x,y坐标），但得到: {track_coords.shape}") + + n_points, n_frames, _ = track_coords.shape + logger.info(f"处理 {n_points} 个点，{n_frames} 帧") + + # 现在重塑vis_scores和conf_scores以匹配重塑后的track_coords + if vis_scores is not None: + logger.info(f"原始visibility_score形状: {vis_scores.shape}") + if vis_scores.ndim == 3 and vis_scores.shape[0] == 1: # (1, frames, points_per_batch) + vis_scores = vis_scores.squeeze(0) # (frames, points_per_batch) + vis_scores = vis_scores.T # 转置为 (points_per_batch, frames) + # 如果有多个batch，需要复制数据以匹配总点数 + if n_points > vis_scores.shape[0]: + batch_count = n_points // vis_scores.shape[0] + vis_scores = np.tile(vis_scores, (batch_count, 1)) # 复制到所有batch + elif vis_scores.ndim == 2: # (frames, points) 或 (points, frames) + if vis_scores.shape[0] == n_frames: # (frames, points) -> (points, frames) + vis_scores = vis_scores.T + # 检查是否需要扩展到匹配点数 + if vis_scores.shape[0] != n_points: + batch_count = n_points // vis_scores.shape[0] + vis_scores = np.tile(vis_scores, (batch_count, 1)) + logger.info(f"重塑后visibility_score形状: {vis_scores.shape}") + + if conf_scores is not None: + logger.info(f"原始confidence_score形状: {conf_scores.shape}") + if conf_scores.ndim == 3 and conf_scores.shape[0] == 1: # (1, frames, points_per_batch) + conf_scores = conf_scores.squeeze(0) # (frames, points_per_batch) + conf_scores = conf_scores.T # 转置为 (points_per_batch, frames) + # 如果有多个batch，需要复制数据以匹配总点数 + if n_points > conf_scores.shape[0]: + batch_count = n_points // conf_scores.shape[0] + conf_scores = np.tile(conf_scores, (batch_count, 1)) # 复制到所有batch + elif conf_scores.ndim == 2: # (frames, points) 或 (points, frames) + if conf_scores.shape[0] == n_frames: # (frames, points) -> (points, frames) + conf_scores = conf_scores.T + # 检查是否需要扩展到匹配点数 + if conf_scores.shape[0] != n_points: + batch_count = n_points // conf_scores.shape[0] + conf_scores = np.tile(conf_scores, (batch_count, 1)) + logger.info(f"重塑后confidence_score形状: {conf_scores.shape}") + + # 对每个点进行投票分割 - 添加详细调试 + point_labels = {} + + # 调试信息：检查mask图像的值范围和物体覆盖率 + logger.info("检查mask图像信息:") + for i, mask_img in enumerate(mask_sequence[:3]): # 只检查前3帧 + unique_values = np.unique(mask_img) + object_pixels = np.sum(mask_img > 128) + total_pixels = mask_img.size + coverage = object_pixels / total_pixels * 100 + logger.info(f" Frame {i}: 图像形状={mask_img.shape}, 值范围={unique_values}, 物体覆盖率={coverage:.1f}%") + + # 调试：检查更多点的坐标范围，包括所有点 + sample_coords = track_coords[:, :3] # 所有点，前3帧 + logger.info(f"所有点坐标范围 (前3帧): x=[{sample_coords[:,:,0].min():.1f}, {sample_coords[:,:,0].max():.1f}], y=[{sample_coords[:,:,1].min():.1f}, {sample_coords[:,:,1].max():.1f}]") + + # 检查前几个点在第一帧的mask值 + first_mask = mask_sequence[0] + logger.info("前20个点在第一帧的mask采样:") + for i in range(min(20, n_points)): + x, y = track_coords[i, 0] + x_int = max(0, min(int(round(x)), first_mask.shape[1] - 1)) + y_int = max(0, min(int(round(y)), first_mask.shape[0] - 1)) + mask_val = first_mask[y_int, x_int] + logger.info(f" 点{i}: ({x:.0f},{y:.0f}) -> mask={mask_val}") + if i == 9: # 分两行显示 + break + + for i in range(10, min(20, n_points)): + x, y = track_coords[i, 0] + x_int = max(0, min(int(round(x)), first_mask.shape[1] - 1)) + y_int = max(0, min(int(round(y)), first_mask.shape[0] - 1)) + mask_val = first_mask[y_int, x_int] + logger.info(f" 点{i}: ({x:.0f},{y:.0f}) -> mask={mask_val}") + + valid_votes_count = 0 + filtered_by_visibility = 0 + filtered_by_confidence = 0 + out_of_bounds_count = 0 + + for point_idx in range(n_points): + votes = defaultdict(int) # {object_id: vote_count} + total_votes = 0 + point_valid_votes = 0 + + for frame_idx in range(min(n_frames, len(mask_sequence))): + # 获取当前帧的坐标 + x, y = track_coords[point_idx, frame_idx] + + # 检查可见性和置信度 + is_visible = True + vis_score = 1.0 + conf_score = 1.0 + + if vis_scores is not None: + vis_score = vis_scores[point_idx, frame_idx] if vis_scores.ndim >= 2 else vis_scores[point_idx] + if vis_score <= 0.5: + is_visible = False + filtered_by_visibility += 1 + + if conf_scores is not None: + conf_score = conf_scores[point_idx, frame_idx] if conf_scores.ndim >= 2 else conf_scores[point_idx] + if conf_score <= confidence_threshold: + is_visible = False + filtered_by_confidence += 1 + + if not is_visible: + continue + + # 获取mask图像 + mask_img = mask_sequence[frame_idx] + h, w = mask_img.shape[:2] + + # 检查坐标是否在范围内 + if x < 0 or x >= w or y < 0 or y >= h: + out_of_bounds_count += 1 + continue + + # 确保坐标在图像范围内 + x_int = max(0, min(int(round(x)), w - 1)) + y_int = max(0, min(int(round(y)), h - 1)) + + # 获取该位置的物体ID + if mask_img.ndim == 2: + raw_id = mask_img[y_int, x_int] + else: + raw_id = mask_img[y_int, x_int, 0] # 假设第一个通道是ID + + # 处理mask值：255 -> 1, 0 -> 0 (将255映射为物体ID 1) + if raw_id > 128: # 假设超过阈值的值都是物体 + object_id = 1 # 简化为单一物体 + else: + object_id = 0 # 背景 + + # 调试信息：记录前几个点的详细投票过程 + if point_idx < 3 and frame_idx < 3: + logger.info(f" 点{point_idx} 帧{frame_idx}: 坐标=({x:.1f},{y:.1f}) -> ({x_int},{y_int}), vis={vis_score:.3f}, conf={conf_score:.3f}, raw_mask={raw_id}, object_id={object_id}") + + # 投票（包括背景ID=0） + votes[int(object_id)] += 1 # 确保使用int类型 + total_votes += 1 + point_valid_votes += 1 + + if object_id > 0: + valid_votes_count += 1 + + # 记录投票结果 + if total_votes > 0: + # 按投票数排序，取最高票的物体ID + sorted_votes = sorted(votes.items(), key=lambda x: x[1], reverse=True) + point_labels[point_idx] = [obj_id for obj_id, count in sorted_votes if count >= total_votes * 0.3] # 至少30%的票 + else: + point_labels[point_idx] = [0] # 未分类 + + # 调试：记录前几个点的投票结果 + if point_idx < 5: + logger.info(f" 点{point_idx}投票结果: 总票数={total_votes}, 有效票数={point_valid_votes}, 最终标签={point_labels[point_idx]}") + + logger.info(f"投票统计: 有效非背景票数={valid_votes_count}, 被可见性过滤={filtered_by_visibility}, 被置信度过滤={filtered_by_confidence}, 坐标越界={out_of_bounds_count}") + + # 如果没有找到物体点，尝试基于mask生成虚拟tracks + if valid_votes_count == 0: + logger.warning("所有查询点都落在背景区域，尝试基于mask生成虚拟物体点") + # 注释掉虚拟点生成，因为我们现在使用更好的完整点云分割方法 + # virtual_points = generate_virtual_object_points(mask_sequence, n_virtual_points=50) + + # if virtual_points: + # logger.info(f"生成了 {len(virtual_points)} 个虚拟物体点") + # # 为虚拟点添加到point_labels + # for i, (frame_idx, x, y, object_id) in enumerate(virtual_points): + # virtual_point_idx = n_points + i # 避免与现有点冲突 + # point_labels[virtual_point_idx] = [object_id] + # logger.info(f" 虚拟点{virtual_point_idx}: 帧{frame_idx} 坐标({x},{y}) -> 物体{object_id}") + # else: + # logger.error("无法生成虚拟物体点，mask中可能没有物体区域") + logger.info("将使用完整点云分割方法来查找物体点") + + logger.info(f"完成点分割，{len(point_labels)} 个点被分类") + + # 统计分割结果 + object_counts = defaultdict(int) + for labels in point_labels.values(): + for label in labels: + object_counts[label] += 1 + + logger.info(f"分割统计: {dict(object_counts)}") + + return point_labels + +def generate_virtual_object_points(mask_sequence: List[np.ndarray], n_virtual_points: int = 50) -> List[Tuple[int, int, int, int]]: + """ + 基于mask序列生成虚拟物体点 + + Args: + mask_sequence: mask图像序列 + n_virtual_points: 要生成的虚拟点数量 + + Returns: + List of (frame_idx, x, y, object_id) tuples + """ + virtual_points = [] + + for frame_idx, mask_img in enumerate(mask_sequence): + # 找到物体像素位置 + object_pixels = np.where(mask_img > 128) + + if len(object_pixels[0]) > 0: + # 随机采样物体像素位置 + n_samples = min(max(1, n_virtual_points // len(mask_sequence)), len(object_pixels[0])) + + if n_samples > 0: + indices = np.random.choice(len(object_pixels[0]), n_samples, replace=False) + + for idx in indices: + y = int(object_pixels[0][idx]) + x = int(object_pixels[1][idx]) + object_id = 1 # 简化为单一物体 + virtual_points.append((frame_idx, x, y, object_id)) + + # 只处理前几帧就够了 + if len(virtual_points) >= n_virtual_points: + break + + return virtual_points[:n_virtual_points] + +def create_mask_based_segmentation(raw_vggt_result: Dict, mask_sequence: List[np.ndarray], + target_object_id: int = None) -> np.ndarray: + """ + 基于mask序列对完整点云进行分割（正确的方法） + + Args: + raw_vggt_result: VGGT原生结果 + mask_sequence: mask图像序列 + target_object_id: 目标物体ID，None表示所有非背景物体 + + Returns: + segment_mask: 布尔掩码数组，指示哪些点属于目标物体 + """ + # 使用固定的初始阈值，然后通过自动检测调整 + mask_threshold = 128 + logger.info(f"开始基于mask的真正点云分割，初始mask阈值: {mask_threshold}") + + # 强制检查mask序列 + logger.info(f"🔍 强制检查mask序列: type={type(mask_sequence)}, len={len(mask_sequence) if hasattr(mask_sequence, '__len__') else 'N/A'}") + if hasattr(mask_sequence, '__len__') and len(mask_sequence) > 0: + first_mask = mask_sequence[0] + logger.info(f"🔍 第一个mask: type={type(first_mask)}") + if hasattr(first_mask, 'shape'): + logger.info(f"🔍 第一个mask形状: {first_mask.shape}") + if hasattr(first_mask, 'dtype'): + logger.info(f"🔍 第一个mask dtype: {first_mask.dtype}") + if hasattr(first_mask, 'min'): + logger.info(f"🔍 第一个mask值范围: [{first_mask.min()}, {first_mask.max()}]") + if hasattr(first_mask, 'unique'): + unique_vals = np.unique(first_mask) if hasattr(np, 'unique') else 'N/A' + logger.info(f"🔍 第一个mask唯一值: {unique_vals}") + + # 首先分析mask序列的值范围，自动调整阈值 + logger.info("分析mask序列的值分布...") + all_unique_values = set() + max_val_in_sequence = 0 + min_val_in_sequence = float('inf') + + for i, mask in enumerate(mask_sequence[:5]): # 只检查前5帧 + unique_vals = np.unique(mask) + all_unique_values.update(unique_vals) + max_val_in_sequence = max(max_val_in_sequence, mask.max()) + min_val_in_sequence = min(min_val_in_sequence, mask.min()) + + if i < 3: # 详细显示前3帧 + logger.info(f" Mask {i}: shape={mask.shape}, dtype={mask.dtype}, " + f"range=[{mask.min()}, {mask.max()}], unique_count={len(unique_vals)}") + + logger.info(f"序列总体: min={min_val_in_sequence}, max={max_val_in_sequence}, " + f"unique_values={sorted(list(all_unique_values))[:20]}...") + + # 自动调整mask阈值 + original_threshold = mask_threshold + if max_val_in_sequence <= 1.0: + # 0-1范围的mask，调整阈值 + mask_threshold = 0.5 + logger.warning(f"检测到0-1范围的mask，自动调整阈值: {original_threshold} -> {mask_threshold}") + elif max_val_in_sequence < original_threshold: + # 最大值小于阈值，使用最大值的一半 + mask_threshold = max_val_in_sequence / 2 + logger.warning(f"Mask最大值({max_val_in_sequence})小于原阈值({original_threshold})，" + f"自动调整阈值为: {mask_threshold}") + elif len(all_unique_values) == 2 and 0 in all_unique_values: + # 二值mask，使用两个值的中间值 + non_zero_vals = [v for v in all_unique_values if v > 0] + if non_zero_vals: + mask_threshold = non_zero_vals[0] / 2 + logger.info(f"检测到二值mask (0, {non_zero_vals[0]})，调整阈值为: {mask_threshold}") + + logger.info(f"最终使用的mask阈值: {mask_threshold}") + + # 获取完整点云数据 + if 'points_from_depth' in raw_vggt_result: + points = raw_vggt_result['points_from_depth'] + logger.info("使用 points_from_depth 进行分割") + elif 'points' in raw_vggt_result: + if isinstance(raw_vggt_result['points'], dict): + points = raw_vggt_result['points']['point_map'] + else: + points = raw_vggt_result['points'] + logger.info("使用 points 进行分割") + else: + raise ValueError("No point cloud data found") + + # 转换为numpy + if isinstance(points, torch.Tensor): + points_np = points.cpu().numpy() + else: + points_np = points + + # 去掉batch维度并重塑 + if points_np.ndim == 5 and points_np.shape[0] == 1: + points_np = np.squeeze(points_np, axis=0) + + # 记录原始形状以便重建 + original_shape = points_np.shape + logger.info(f"点云原始形状: {original_shape}") + + # 检查点云分辨率 + if len(original_shape) == 4: # (S, H, W, 3) + S, H, W, _ = original_shape + logger.info(f"点云分辨率: {W}x{H}, 帧数: {S}") + if len(mask_sequence) > 0: + mask_h, mask_w = mask_sequence[0].shape + logger.info(f"Mask分辨率: {mask_w}x{mask_h}") + if W != mask_w or H != mask_h: + logger.warning(f"⚠️ 点云分辨率({W}x{H})与Mask分辨率({mask_w}x{mask_h})不匹配！") + + # 重塑为 (N_points, 3) + if points_np.ndim == 4: # (S, H, W, 3) + S, H, W, _ = points_np.shape + points_flat = points_np.reshape(-1, 3) + elif points_np.ndim == 3: # (N, H, W) or similar + points_flat = points_np.reshape(-1, 3) + else: + points_flat = points_np + + # 过滤无效点 + valid_mask = ~np.isnan(points_flat).any(axis=1) & ~np.isinf(points_flat).any(axis=1) + valid_points = points_flat[valid_mask] + logger.info(f"有效点数: {len(valid_points)}/{len(points_flat)}") + + # 获取相机参数 + cameras = raw_vggt_result.get('cameras', {}) + extrinsic = cameras.get('extrinsic') + intrinsic = cameras.get('intrinsic') + + # 检查是否可以使用直接像素对应方法 + USE_PIXEL_CORRESPONDENCE = False + if len(original_shape) == 4: # (S, H, W, 3) + S, H, W, _ = original_shape + if len(mask_sequence) > 0: + mask_h, mask_w = mask_sequence[0].shape + if W == mask_w and H == mask_h: + logger.info("✅ 点云和Mask分辨率匹配，使用直接像素对应方法") + USE_PIXEL_CORRESPONDENCE = True + + # 临时强制使用简单分割来测试 + USE_SIMPLE_SEGMENTATION = False # 设为True来测试简单分割 + + if USE_PIXEL_CORRESPONDENCE: + # 使用直接像素对应方法（最简单最准确） + return create_pixel_correspondence_segmentation( + points_flat, mask_sequence, valid_mask, original_shape + ) + elif extrinsic is None or intrinsic is None or USE_SIMPLE_SEGMENTATION: + logger.warning("缺少相机参数或强制使用简单分割") + return create_simple_spatial_segmentation(valid_points, mask_sequence, valid_mask) + + # 转换相机参数 + if isinstance(extrinsic, torch.Tensor): + extrinsic = extrinsic.cpu().numpy() + if isinstance(intrinsic, torch.Tensor): + intrinsic = intrinsic.cpu().numpy() + + # 去掉batch维度 + if extrinsic.ndim == 4 and extrinsic.shape[0] == 1: + extrinsic = np.squeeze(extrinsic, axis=0) + if intrinsic.ndim == 3 and intrinsic.shape[0] == 1: + intrinsic = np.squeeze(intrinsic, axis=0) + + logger.info(f"相机外参形状: {extrinsic.shape}, 内参形状: {intrinsic.shape}") + logger.info(f"Mask序列长度: {len(mask_sequence)}") + + # 检查是否需要缩放mask（在循环之前决定） + need_resize_mask = False + if len(mask_sequence) > 0: + original_mask_h, original_mask_w = mask_sequence[0].shape + if original_mask_w == 1024 and original_mask_h == 896: + # 检查内参是否是针对518x448分辨率的 + scale_x = original_mask_w / 518.0 + scale_y = original_mask_h / 448.0 + if abs(scale_x - 2.0) < 0.1 and abs(scale_y - 2.0) < 0.1: + need_resize_mask = True + logger.info("将在投影过程中缩放mask从1024x896到518x448以匹配点云分辨率") + + # 使用更大的批次处理，加快速度 + n_points = len(valid_points) + batch_size = 500000 # 增大到50万个点，减少批次数 + n_frames = min(len(mask_sequence), extrinsic.shape[0]) + + logger.info(f"处理 {n_frames} 帧进行点云投影分割，总点数: {n_points}") + + # 初始化点标签和置信度 + point_labels = np.zeros(n_points, dtype=np.uint8) + confidence_scores = np.zeros(n_points, dtype=np.float32) + + # 添加调试统计 + total_valid_projections = 0 + total_object_hits = 0 + + # 采样处理以加速 - 如果点太多，只处理一部分 + if n_points > 1000000: # 超过100万点时采样 + sample_rate = 0.2 # 采样20%的点 + sample_indices = np.random.choice(n_points, int(n_points * sample_rate), replace=False) + sampled_points = valid_points[sample_indices] + logger.info(f"点云过大，采样处理 {len(sampled_points)} 个点 ({sample_rate*100:.0f}%)") + else: + sampled_points = valid_points + sample_indices = np.arange(n_points) + + # 分批处理采样后的点云 + total_batches = (len(sampled_points) + batch_size - 1) // batch_size + for batch_idx, batch_start in enumerate(range(0, len(sampled_points), batch_size)): + batch_end = min(batch_start + batch_size, len(sampled_points)) + batch_points = sampled_points[batch_start:batch_end] + batch_votes = np.zeros((len(batch_points), 256), dtype=np.int16) + + # 显示进度（减少日志输出） + if batch_idx % 5 == 0 or batch_idx == total_batches - 1: + progress = (batch_idx + 1) / total_batches * 100 + logger.info(f"处理批次 {batch_idx + 1}/{total_batches} ({progress:.1f}%)") + + # 第一批添加详细调试 + if batch_idx == 0: + logger.info(f"第一批点云范围: X=[{batch_points[:, 0].min():.2f}, {batch_points[:, 0].max():.2f}], " + f"Y=[{batch_points[:, 1].min():.2f}, {batch_points[:, 1].max():.2f}], " + f"Z=[{batch_points[:, 2].min():.2f}, {batch_points[:, 2].max():.2f}]") + + # 只处理关键帧以加速（每隔3帧处理一次） + frame_step = 3 + processed_frames = list(range(0, n_frames, frame_step)) + + for frame_idx in processed_frames: + if frame_idx >= len(mask_sequence): + break + + mask = mask_sequence[frame_idx] + + # 如果需要缩放mask到518x448 + if need_resize_mask: + original_h, original_w = mask.shape + if original_w == 1024 and original_h == 896: + import cv2 + mask = cv2.resize(mask, (518, 448), interpolation=cv2.INTER_NEAREST) + if batch_idx == 0 and frame_idx == 0: + logger.info(f"缩放mask从{original_w}x{original_h}到518x448以匹配点云分辨率") + + # 获取当前帧的相机参数 + if extrinsic.shape[0] > frame_idx: + extr = extrinsic[frame_idx] # (3, 4) or (4, 4) + if extr.shape[0] == 3: + # 补充为4x4矩阵 + extr_4x4 = np.eye(4) + extr_4x4[:3, :] = extr + extr = extr_4x4 + else: + extr = np.eye(4) + + # 处理内参矩阵的不同格式 + if intrinsic.ndim == 4: # (1, N, 3, 3) + intr = intrinsic[0, frame_idx] if frame_idx < intrinsic.shape[1] else intrinsic[0, 0] + elif intrinsic.ndim == 3: # (N, 3, 3) + intr = intrinsic[frame_idx] if frame_idx < intrinsic.shape[0] else intrinsic[0] + elif intrinsic.ndim == 2: # (3, 3) + intr = intrinsic + else: + logger.warning(f"内参矩阵形状异常: {intrinsic.shape}，使用默认值") + intr = np.eye(3) + intr[0, 0] = 1000 # fx + intr[1, 1] = 1000 # fy + intr[0, 2] = 512 # cx + intr[1, 2] = 448 # cy + + # 在第一帧打印内参矩阵信息（在获取mask尺寸之前） + if batch_idx == 0 and frame_idx == 0: + logger.info(f"内参矩阵:\n{intr}") + logger.info(f" fx={intr[0,0]:.1f}, fy={intr[1,1]:.1f}") + logger.info(f" cx={intr[0,2]:.1f}, cy={intr[1,2]:.1f}") + # 获取原始mask尺寸 + original_mask_h, original_mask_w = mask_sequence[0].shape + logger.info(f" 原始mask尺寸: {original_mask_w}x{original_mask_h}") + logger.info(f" 是否需要缩放mask: {need_resize_mask}") + logger.info(f"保持内参不变: fx={intr[0,0]:.1f}, fy={intr[1,1]:.1f}, " + f"cx={intr[0,2]:.1f}, cy={intr[1,2]:.1f}") + + # 投影3D点到2D图像 + # 检查外参矩阵是否接近单位矩阵 + is_identity = np.allclose(extr[:3, :3], np.eye(3), atol=0.01) + if is_identity: + # 如果外参接近单位矩阵，使用简化投影（假设点云已在相机坐标系） + projected_points = project_3d_to_2d_simple(batch_points, intr) + if batch_idx == 0 and frame_idx == 0: + logger.info("使用简化投影（点云已在相机坐标系）") + else: + # 否则使用完整投影 + projected_points = project_3d_to_2d(batch_points, extr, intr) + if batch_idx == 0 and frame_idx == 0: + logger.info("使用完整投影（需要外参变换）") + + # 获取当前mask的尺寸 + h, w = mask.shape + + # 第一批第一帧的调试信息 + if batch_idx == 0 and frame_idx == 0: + logger.info(f"第一批第一帧投影结果:") + logger.info(f" 投影点范围: X=[{projected_points[:, 0].min():.1f}, {projected_points[:, 0].max():.1f}], " + f"Y=[{projected_points[:, 1].min():.1f}, {projected_points[:, 1].max():.1f}]") + logger.info(f" 图像尺寸: {w}x{h}") + logger.info(f" 外参矩阵是否接近单位矩阵: {np.allclose(extr[:3, :3], np.eye(3), atol=0.01)}") + + # 重要：详细检查mask图像 + logger.info(f" Mask图像详细信息:") + logger.info(f" - dtype: {mask.dtype}") + logger.info(f" - shape: {mask.shape}") + logger.info(f" - 值范围: [{mask.min()}, {mask.max()}]") + logger.info(f" - 唯一值: {np.unique(mask)[:20]}...") # 只显示前20个 + logger.info(f" - 值>mask_threshold({mask_threshold})的像素数: {np.sum(mask > mask_threshold)}") + logger.info(f" - 值==255的像素数: {np.sum(mask == 255)}") + logger.info(f" - 值==1的像素数: {np.sum(mask == 1)}") + logger.info(f" - 值==0的像素数: {np.sum(mask == 0)}") + + # 检查前10个点的投影结果和对应的mask值 + for i in range(min(10, len(projected_points))): + logger.info(f" 点{i}: 3D({batch_points[i, 0]:.2f}, {batch_points[i, 1]:.2f}, {batch_points[i, 2]:.2f}) " + f"-> 2D({projected_points[i, 0]:.1f}, {projected_points[i, 1]:.1f})") + # 检查是否在图像范围内 + in_range = (0 <= projected_points[i, 0] < w) and (0 <= projected_points[i, 1] < h) + if in_range: + x_int = int(projected_points[i, 0]) + y_int = int(projected_points[i, 1]) + mask_val = mask[y_int, x_int] + # 检查周围的mask值 + y_start, y_end = max(0, y_int-1), min(h, y_int+2) + x_start, x_end = max(0, x_int-1), min(w, x_int+2) + neighbor_vals = mask[y_start:y_end, x_start:x_end] + logger.info(f" -> 在图像内({x_int},{y_int}), mask值={mask_val}, 周围3x3={neighbor_vals.flatten()}") + else: + logger.info(f" -> 超出图像范围") + + # 如果mask全是0，检查是否应该调整阈值 + if mask.max() <= 1.0: + logger.warning(f" ⚠️ 检测到mask值范围[0,1]，可能需要将mask_threshold调整为更小的值！") + logger.info(f" 建议尝试: mask_threshold=0.5 (对于0-1范围的mask)") + elif mask.max() < mask_threshold: + logger.warning(f" ⚠️ Mask最大值({mask.max()})小于阈值({mask_threshold})，所有像素都会被视为背景！") + logger.info(f" 建议将mask_threshold调整为: {mask.max() // 2}") + + # 检查有效投影（不是inf且在图像范围内） + not_inf = (projected_points[:, 0] != float('inf')) & (projected_points[:, 1] != float('inf')) + valid_u = (projected_points[:, 0] >= 0) & (projected_points[:, 0] < w) + valid_v = (projected_points[:, 1] >= 0) & (projected_points[:, 1] < h) + valid_proj = not_inf & valid_u & valid_v + + if np.any(valid_proj): + valid_indices = np.where(valid_proj)[0] + valid_coords = projected_points[valid_proj].astype(np.int32) + + # 批量查询mask值 + object_ids = mask[valid_coords[:, 1], valid_coords[:, 0]] + + # 统计 + total_valid_projections += len(valid_indices) + object_mask = object_ids > mask_threshold # 前景物体 + total_object_hits += np.sum(object_mask) + + # 更新投票 - 将mask值转换为物体ID + for idx, raw_id in zip(valid_indices, object_ids): + if raw_id > mask_threshold: # 前景 + obj_id = 1 # 简化为单一物体ID + batch_votes[idx, obj_id] += 1 + else: # 背景 + batch_votes[idx, 0] += 1 + + # 基于投票结果分配标签（批次内） + # 使用更宽松的阈值：只要有投票就算 + for i in range(len(batch_votes)): + max_votes = np.max(batch_votes[i]) + if max_votes > 0: + # 找到得票最多的类别 + label = np.argmax(batch_votes[i]) + # 映射回原始索引 + original_idx = sample_indices[batch_start + i] + # 如果是前景类别，即使只有1票也接受 + if label > 0 and batch_votes[i, label] >= 1: + point_labels[original_idx] = label + confidence_scores[original_idx] = batch_votes[i, label] / len(processed_frames) + elif batch_votes[i, 0] > batch_votes[i, label]: # 背景票数更多 + point_labels[original_idx] = 0 + confidence_scores[original_idx] = batch_votes[i, 0] / len(processed_frames) + + # 定期报告进度 + if (batch_idx + 1) % 5 == 0: + current_object_points = np.sum(point_labels > 0) + logger.info(f" 当前已找到 {current_object_points} 个物体点") + + # 统计最终结果 + logger.info(f"投影统计: 总有效投影={total_valid_projections}, 命中物体={total_object_hits}") + logger.info(f"投影命中率: {total_object_hits/total_valid_projections*100:.2f}%" if total_valid_projections > 0 else "无有效投影") + + # 创建目标物体的mask + if target_object_id is None: + # 所有非背景物体 + segment_mask = point_labels > 0 + else: + # 特定物体 + segment_mask = point_labels == target_object_id + + object_counts = np.bincount(point_labels) + logger.info(f"分割统计: 背景={object_counts[0]}, 物体1={object_counts[1] if len(object_counts) > 1 else 0}") + logger.info(f"目标物体点数: {np.sum(segment_mask)}") + + # 将结果映射回原始点云索引 + full_segment_mask = np.zeros(len(points_flat), dtype=bool) + full_segment_mask[valid_mask] = segment_mask + + return full_segment_mask + +def project_3d_to_2d_simple(points_3d: np.ndarray, intrinsic: np.ndarray) -> np.ndarray: + """简化的投影函数，假设点云已经在相机坐标系中""" + # points_3d: (N, 3) - 已经在相机坐标系中 + # intrinsic: (3, 3) - 相机内参矩阵 + + # 过滤掉Z<=0的点 + valid_z = points_3d[:, 2] > 1e-6 + projected = np.full((len(points_3d), 2), float('inf')) # 使用inf标记无效投影 + + if np.any(valid_z): + valid_points = points_3d[valid_z] + + # 应用内参矩阵 + fx = intrinsic[0, 0] + fy = intrinsic[1, 1] + cx = intrinsic[0, 2] + cy = intrinsic[1, 2] + + # 直接使用相机坐标投影（不需要归一化，因为内参已经包含了焦距） + # u = fx * X/Z + cx + # v = fy * Y/Z + cy + u = fx * valid_points[:, 0] / valid_points[:, 2] + cx + v = fy * valid_points[:, 1] / valid_points[:, 2] + cy + + # 如果Y坐标大量为负，说明Y轴方向相反 + if np.median(v) < 0: + # 翻转Y坐标并重新计算 + v = fy * (-valid_points[:, 1]) / valid_points[:, 2] + cy + logger.debug("检测到Y轴需要翻转") + + projected[valid_z, 0] = u + projected[valid_z, 1] = v + + return projected + +def project_3d_to_2d(points_3d: np.ndarray, extrinsic: np.ndarray, intrinsic: np.ndarray) -> np.ndarray: + """将3D点投影到2D图像坐标""" + # points_3d: (N, 3) + # extrinsic: (4, 4) - 世界坐标到相机坐标变换矩阵 + # intrinsic: (3, 3) - 相机内参矩阵 + + # 转换为齐次坐标 + points_homo = np.hstack([points_3d, np.ones((len(points_3d), 1))]) # (N, 4) + + # 世界坐标 -> 相机坐标 + # 注意：VGGT可能使用不同的坐标系，可能需要取逆 + # 尝试使用外参矩阵的逆 + try: + extrinsic_inv = np.linalg.inv(extrinsic) + camera_points = (extrinsic_inv @ points_homo.T).T # (N, 4) + except: + # 如果求逆失败，使用原始矩阵 + camera_points = (extrinsic @ points_homo.T).T # (N, 4) + + camera_points_3d = camera_points[:, :3] # (N, 3) + + # 相机坐标 -> 图像坐标 + # 投影到图像平面 (避免除零) + valid_z = camera_points_3d[:, 2] > 1e-6 + projected = np.zeros((len(points_3d), 2)) + + if np.any(valid_z): + valid_camera_points = camera_points_3d[valid_z] + + # 投影到图像平面 + # 注意：这里直接应用内参矩阵 + projected_homo = np.hstack([valid_camera_points, np.ones((len(valid_camera_points), 1))]) + image_coords = (intrinsic @ projected_homo[:, :3].T).T # (N, 3) + + # 归一化齐次坐标 + image_coords[:, :2] /= image_coords[:, 2:3] + projected[valid_z] = image_coords[:, :2] + + return projected + +def create_pixel_correspondence_segmentation(points_flat: np.ndarray, mask_sequence: List[np.ndarray], + valid_mask: np.ndarray, original_shape: tuple) -> np.ndarray: + """使用直接像素对应进行分割（当点云和mask分辨率匹配时）""" + logger.info("使用直接像素对应分割方法") + + if len(original_shape) != 4: + raise ValueError(f"需要4D点云形状，但得到: {original_shape}") + + S, H, W, _ = original_shape + logger.info(f"点云形状: S={S}, H={H}, W={W}") + + # 初始化分割结果 + full_segment_mask = np.zeros(len(points_flat), dtype=bool) + + # 对每一帧进行处理 + for frame_idx in range(min(S, len(mask_sequence))): + mask = mask_sequence[frame_idx] + + # 计算当前帧在扁平化数组中的索引范围 + frame_start = frame_idx * H * W + frame_end = frame_start + H * W + + # 获取当前帧的有效点索引 + frame_valid_mask = valid_mask[frame_start:frame_end] + + # 将mask扁平化 + mask_flat = mask.flatten() + + # 应用mask到有效点 + # 注意：需要考虑valid_mask的影响 + valid_indices = np.where(frame_valid_mask)[0] + for idx in valid_indices: + pixel_idx = idx # 在当前帧内的像素索引 + if pixel_idx < len(mask_flat): + if mask_flat[pixel_idx] > 127: # 使用固定阈值127 + full_segment_mask[frame_start + idx] = True + + object_count = np.sum(full_segment_mask) + total_valid = np.sum(valid_mask) + logger.info(f"像素对应分割完成: {object_count}/{total_valid} 个点被标记为物体 " + f"({object_count/total_valid*100:.1f}%)") + + return full_segment_mask + +def create_simple_spatial_segmentation(points: np.ndarray, mask_sequence: List[np.ndarray], + valid_mask: np.ndarray) -> np.ndarray: + """基于空间位置的简单分割（当缺少相机参数时）""" + logger.info("使用简单空间分割策略") + + # 计算mask的空间覆盖率 + total_coverage = 0 + for mask in mask_sequence: + coverage = np.sum(mask > 0) / mask.size + total_coverage += coverage + + avg_coverage = total_coverage / len(mask_sequence) + logger.info(f"平均mask覆盖率: {avg_coverage*100:.1f}%") + + # 基于覆盖率估算应该分割的点数 + n_object_points = int(len(points) * avg_coverage) + + # 使用Z坐标进行简单分割（假设物体在特定深度范围） + z_coords = points[:, 2] + z_sorted_indices = np.argsort(z_coords) + + # 选择中等深度的点作为物体点（避免太近或太远的点） + start_idx = len(points) // 4 + end_idx = start_idx + n_object_points + object_indices = z_sorted_indices[start_idx:end_idx] + + segment_mask = np.zeros(len(points), dtype=bool) + segment_mask[object_indices] = True + + # 映射回完整点云 + full_mask = np.zeros(len(valid_mask), dtype=bool) + full_mask[valid_mask] = segment_mask + + logger.info(f"简单分割选择了 {np.sum(segment_mask)} 个点") + return full_mask + +def create_segmented_pointcloud_new(raw_vggt_result: Dict, mask_sequence: List[np.ndarray], + target_object_id: int = None) -> Tuple[np.ndarray, np.ndarray]: + """ + 基于mask序列创建分割点云（新的正确方法） + + Args: + raw_vggt_result: VGGT原生结果 + mask_sequence: mask图像序列 + target_object_id: 目标物体ID，None表示保留所有非背景点 + + Returns: + vertices: 点云顶点坐标 (N, 3) + colors: 点云颜色 (N, 3) + """ + logger.info(f"使用新方法创建分割点云，目标物体ID: {target_object_id}，启用自动阈值检测") + + # 获取完整点云的分割mask + segment_mask = create_mask_based_segmentation(raw_vggt_result, mask_sequence, target_object_id) + + # 获取点云数据 + if 'points_from_depth' in raw_vggt_result: + points = raw_vggt_result['points_from_depth'] + logger.info("使用 points_from_depth") + elif 'points' in raw_vggt_result: + if isinstance(raw_vggt_result['points'], dict): + points = raw_vggt_result['points']['point_map'] + else: + points = raw_vggt_result['points'] + logger.info("使用 points") + else: + raise ValueError("No point cloud data found in raw_vggt_result") + + # 转换为numpy + if isinstance(points, torch.Tensor): + points_np = points.cpu().numpy() + else: + points_np = points + + # 去掉batch维度 + if points_np.ndim == 5 and points_np.shape[0] == 1: + points_np = np.squeeze(points_np, axis=0) + + # 重塑为点云 (N_points, 3) + if points_np.ndim == 4: + points_np = points_np.reshape(-1, 3) + elif points_np.ndim == 3: + points_np = points_np.reshape(-1, 3) + + # 获取颜色数据 + colors_np = None + if 'images' in raw_vggt_result: + images = raw_vggt_result['images'] + if isinstance(images, torch.Tensor): + images_np = images.cpu().numpy() + else: + images_np = images + + if images_np.ndim == 5 and images_np.shape[0] == 1: + images_np = np.squeeze(images_np, axis=0) + + if images_np.ndim == 4: + if images_np.shape[1] == 3: # (S, 3, H, W) -> (S, H, W, 3) + images_np = np.transpose(images_np, (0, 2, 3, 1)) + + colors_np = images_np.reshape(-1, 3) + if colors_np.max() <= 1.0: + colors_np = (colors_np * 255).astype(np.uint8) + else: + colors_np = colors_np.astype(np.uint8) + + # 确保segment_mask长度匹配 + if len(segment_mask) != len(points_np): + logger.warning(f"Segment mask长度({len(segment_mask)})与点云长度({len(points_np)})不匹配，调整...") + min_len = min(len(segment_mask), len(points_np)) + segment_mask = segment_mask[:min_len] + points_np = points_np[:min_len] + if colors_np is not None: + colors_np = colors_np[:min_len] + + # 应用分割mask + segmented_points = points_np[segment_mask] + + if colors_np is not None: + segmented_colors = colors_np[segment_mask] + else: + # 生成默认颜色 + if target_object_id is not None and target_object_id > 0: + # 为不同物体分配不同颜色 + colors = [ + [255, 0, 0], # 红色 + [0, 255, 0], # 绿色 + [0, 0, 255], # 蓝色 + [255, 255, 0], # 黄色 + [255, 0, 255], # 洋红 + [0, 255, 255], # 青色 + [255, 128, 0], # 橙色 + [128, 0, 255], # 紫色 + ] + color = colors[target_object_id % len(colors)] + segmented_colors = np.tile(color, (len(segmented_points), 1)).astype(np.uint8) + else: + segmented_colors = np.ones((len(segmented_points), 3), dtype=np.uint8) * 128 + + logger.info(f"新方法分割后保留 {len(segmented_points)} 个点") + return segmented_points, segmented_colors + +def create_segmented_pointcloud(raw_vggt_result: Dict, point_labels: Dict[int, List[int]], + target_object_id: int = None) -> Tuple[np.ndarray, np.ndarray]: + """ + 根据分割结果创建点云数据（兼容旧接口） + + Args: + raw_vggt_result: VGGT原生结果 + point_labels: 点分割标签（已弃用，仅为兼容性保留） + target_object_id: 目标物体ID，None表示保留所有分割的点 + + Returns: + vertices: 点云顶点坐标 (N, 3) + colors: 点云颜色 (N, 3) + """ + logger.warning("使用旧的create_segmented_pointcloud接口，建议改用create_segmented_pointcloud_new") + + # 由于旧方法问题很大，这里简单返回少量点作为兼容 + logger.info(f"创建分割点云，目标物体ID: {target_object_id}") + + # 获取点云数据 + if 'points_from_depth' in raw_vggt_result: + points = raw_vggt_result['points_from_depth'] + logger.info("使用 points_from_depth") + elif 'points' in raw_vggt_result: + if isinstance(raw_vggt_result['points'], dict): + points = raw_vggt_result['points']['point_map'] + else: + points = raw_vggt_result['points'] + logger.info("使用 points") + else: + raise ValueError("No point cloud data found in raw_vggt_result") + + # 转换为numpy + if isinstance(points, torch.Tensor): + points_np = points.cpu().numpy() + else: + points_np = points + + # 去掉batch维度 + if points_np.ndim == 5 and points_np.shape[0] == 1: + points_np = np.squeeze(points_np, axis=0) + + # 重塑为点云 (S*H*W, 3) + if points_np.ndim == 4: + points_np = points_np.reshape(-1, 3) + elif points_np.ndim == 3: + points_np = points_np.reshape(-1, 3) + + # 过滤无效点 + valid_mask = ~np.isnan(points_np).any(axis=1) & ~np.isinf(points_np).any(axis=1) + points_np = points_np[valid_mask] + + # 简单选择前1000个点作为示例 + n_sample = min(1000, len(points_np)) + indices = np.random.choice(len(points_np), n_sample, replace=False) + sample_points = points_np[indices] + + # 生成默认颜色 + if target_object_id is not None and target_object_id > 0: + colors = [ + [255, 0, 0], # 红色 + [0, 255, 0], # 绿色 + [0, 0, 255], # 蓝色 + [255, 255, 0], # 黄色 + [255, 0, 255], # 洋红 + [0, 255, 255], # 青色 + [255, 128, 0], # 橙色 + [128, 0, 255], # 紫色 + ] + color = colors[target_object_id % len(colors)] + colors_np = np.tile(color, (len(sample_points), 1)).astype(np.uint8) + else: + colors_np = np.ones((len(sample_points), 3), dtype=np.uint8) * 128 + + logger.info(f"旧方法兼容模式，保留 {len(sample_points)} 个点") + return sample_points, colors_np + +# ----------------------------------------------------------------------------- +# 主要节点实现 +# ----------------------------------------------------------------------------- + +class VGGTMaskProcessorNode: + """VGGT Mask处理节点 - 使用tracks数据和mask序列生成分割的3D模型""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "native_full_result": ("RAW_VGGT_RESULT", { + "tooltip": "来自VGGTNativeFullOutputNode的完整原生结果" + }), + "tracks_json": ("STRING", { + "tooltip": "来自VGGTNativeFullOutputNode的tracks JSON数据" + }), + "mask_sequence": ("MASK", { + "tooltip": "mask图像序列，每张图像包含分割标识（不同像素值代表不同物体）" + }), + }, + "optional": { + "target_object_id": ("INT", { + "default": 0, "min": 0, "max": 255, "step": 1, + "tooltip": "目标物体ID（0=导出所有分割物体，>0=导出指定物体）" + }), + "confidence_threshold": ("FLOAT", { + "default": 0.1, "min": 0.0, "max": 1.0, "step": 0.05, + "tooltip": "点追踪置信度阈值，低于此值的点会被忽略" + }), + "export_format": (["PLY", "GLB", "BOTH"], { + "default": "PLY", + "tooltip": "导出格式：PLY(点云)、GLB(网格)或两者都导出" + }), + "use_full_pointcloud": ("BOOLEAN", { + "default": True, + "tooltip": "是否使用完整点云分割（推荐）。如果关闭，仅使用tracks点进行分割" + }), + } + } + + RETURN_TYPES = ( + "STRING", # 分割统计信息JSON + "STRING", # PLY文件路径 + "STRING", # GLB文件路径 + "STRING", # 分割报告 + ) + RETURN_NAMES = ( + "segmentation_stats", + "segmented_ply_path", + "segmented_glb_path", + "segmentation_report", + ) + OUTPUT_TOOLTIPS = [ + "分割统计信息（JSON格式，包含每个物体的点数统计）", + "分割后的PLY点云文件路径", + "分割后的GLB 3D模型文件路径", + "详细的分割处理报告（JSON格式）" + ] + OUTPUT_NODE = True + FUNCTION = "process_mask_segmentation" + CATEGORY = "💃VVL/VGGT Mask" + + def process_mask_segmentation(self, native_full_result: Dict, tracks_json: str, mask_sequence, + target_object_id: int = 0, + confidence_threshold: float = 0.5, + export_format: str = "PLY", + use_full_pointcloud: bool = True): + """处理mask分割""" + logger.info("开始VGGT Mask分割处理") + logger.info(f"接收到tracks_json类型: {type(tracks_json)}, 长度: {len(tracks_json) if tracks_json else 0}") + + tracks_data = None + + try: + # 首先尝试从native_full_result直接获取tracks数据 + if 'tracks' in native_full_result and native_full_result['tracks']: + logger.info("直接从native_full_result获取tracks数据") + raw_tracks = native_full_result['tracks'] + logger.info(f"从native_full_result获取的tracks keys: {list(raw_tracks.keys())}") + + # 检查并转换tensor格式的数据 + converted_tracks = {} + for key, value in raw_tracks.items(): + if hasattr(value, 'cpu') and hasattr(value, 'numpy'): + # PyTorch tensor + converted_tracks[key] = value.cpu().numpy() + logger.info(f"转换tensor {key}, shape: {value.shape}") + elif isinstance(value, np.ndarray): + converted_tracks[key] = value + logger.info(f"保持numpy {key}, shape: {value.shape}") + elif isinstance(value, (list, tuple)): + # 列表或元组，可能包含tensor + try: + # 尝试转换为numpy数组（如果元素是tensor） + if len(value) > 0 and hasattr(value[0], 'cpu'): + # 列表中包含tensor + converted_list = [item.cpu().numpy() if hasattr(item, 'cpu') else item for item in value] + converted_tracks[key] = np.array(converted_list) + logger.info(f"转换tensor列表 {key}, 元素数量: {len(value)}") + else: + # 普通列表 + converted_tracks[key] = np.array(value) if isinstance(value, list) else value + logger.info(f"转换列表 {key}, 元素数量: {len(value)}") + except Exception as e: + logger.warning(f"转换列表 {key} 失败: {e}, 保持原格式") + converted_tracks[key] = value + else: + # 其他类型，可能是隐藏的tensor + try: + # 尝试作为tensor处理 + if hasattr(value, 'cpu'): + converted_tracks[key] = value.cpu().numpy() + logger.info(f"强制转换tensor {key}, type: {type(value)}") + else: + converted_tracks[key] = value + logger.info(f"保持原格式 {key}, type: {type(value)}") + except Exception as e: + logger.warning(f"处理 {key} 时出错: {e}") + converted_tracks[key] = value + + tracks_data = converted_tracks + logger.info("成功从native_full_result获取tracks数据") + + # 如果直接获取失败且有tracks_json，尝试JSON解析 + elif tracks_json and tracks_json.strip(): + logger.info("备用方案：从tracks_json解析数据") + logger.info(f"tracks_json前200字符: {tracks_json[:200]}") + parsed_data = json.loads(tracks_json) + logger.info(f"成功解析JSON，顶级keys: {list(parsed_data.keys())}") + tracks_data = parsed_data.get('tracks_data', {}) + if tracks_data: + logger.info("成功从tracks_json获取tracks数据") + + # 检查是否成功获取了tracks数据 + if not tracks_data: + raise ValueError("无法从native_full_result或tracks_json获取有效的tracks数据") + + # 处理mask序列 - 改进版本 + mask_list = [] + logger.info(f"开始处理mask序列，输入类型: {type(mask_sequence)}") + + if isinstance(mask_sequence, torch.Tensor): + mask_np = mask_sequence.cpu().numpy() + logger.info(f"Tensor mask形状: {mask_np.shape}, dtype: {mask_np.dtype}, 值范围: [{mask_np.min()}, {mask_np.max()}]") + + for i in range(mask_np.shape[0]): + mask_img = mask_np[i] + logger.info(f" 处理mask {i}: shape={mask_img.shape}, dtype={mask_img.dtype}, range=[{mask_img.min()}, {mask_img.max()}]") + + # 如果是RGB图像，转换为单通道 + if mask_img.ndim == 3: + # 检查是否所有通道都相同（灰度图的RGB表示） + if mask_img.shape[2] >= 3: + r, g, b = mask_img[:,:,0], mask_img[:,:,1], mask_img[:,:,2] + if np.allclose(r, g) and np.allclose(g, b): + mask_img = r # 使用第一个通道 + logger.info(f" 检测到灰度RGB，转换为单通道") + else: + # RGB图像，可能需要转换为灰度 + mask_img = np.mean(mask_img, axis=2) + logger.info(f" 转换RGB为灰度") + else: + mask_img = mask_img[:,:,0] + + # 处理值范围 + if mask_img.max() <= 1.0 and mask_img.dtype in [np.float32, np.float64]: + # 0-1范围的浮点mask，转换为0-255 + if np.all(np.isin(mask_img, [0.0, 1.0])): + # 严格的0-1二值mask + mask_img = (mask_img * 255).astype(np.uint8) + logger.info(f" 转换0-1二值mask为0-255") + else: + # 可能是0-1范围的多值mask + mask_img = (mask_img * 255).astype(np.uint8) + logger.info(f" 转换0-1浮点mask为0-255") + else: + # 已经是整数范围，确保类型正确 + mask_img = mask_img.astype(np.uint8) + logger.info(f" 保持整数mask，转换为uint8") + + logger.info(f" 最终mask {i}: shape={mask_img.shape}, dtype={mask_img.dtype}, range=[{mask_img.min()}, {mask_img.max()}], unique={np.unique(mask_img)}") + mask_list.append(mask_img) + + else: + # 非tensor输入，直接使用 + mask_list = list(mask_sequence) + logger.info(f"使用非tensor mask序列，长度: {len(mask_list)}") + if len(mask_list) > 0: + first_mask = mask_list[0] + logger.info(f"第一个mask: type={type(first_mask)}, shape={getattr(first_mask, 'shape', 'N/A')}") + if hasattr(first_mask, 'dtype'): + logger.info(f" dtype={first_mask.dtype}, range=[{first_mask.min()}, {first_mask.max()}]") + + logger.info(f"Mask序列处理完成，共 {len(mask_list)} 张图像") + + # 🔍 强制检查处理后的mask_list + logger.info(f"🔍 强制检查处理后的mask_list:") + for i in range(min(3, len(mask_list))): + mask = mask_list[i] + logger.info(f"🔍 mask_list[{i}]: type={type(mask)}, shape={getattr(mask, 'shape', 'N/A')}") + if hasattr(mask, 'dtype'): + logger.info(f"🔍 dtype={mask.dtype}, min={mask.min()}, max={mask.max()}") + if hasattr(mask, '__len__') and len(mask.shape) == 2: + unique_vals = np.unique(mask) + white_pixels = np.sum(mask == 255) + black_pixels = np.sum(mask == 0) + logger.info(f"🔍 unique_values={unique_vals}, white_pixels={white_pixels}, black_pixels={black_pixels}") + + # 调试tracks_data的详细信息 + logger.info(f"传递给process_tracks_with_masks的tracks_data:") + for key, value in tracks_data.items(): + logger.info(f" {key}: type={type(value)}, hasattr_cpu={hasattr(value, 'cpu')}") + if hasattr(value, 'shape'): + logger.info(f" shape: {value.shape}") + elif hasattr(value, '__len__'): + logger.info(f" length: {len(value)}") + + # 使用tracks数据和mask进行分割 + point_labels = process_tracks_with_masks( + tracks_data, mask_list, confidence_threshold + ) + + # 统计分割结果 + stats = self._generate_segmentation_stats(point_labels) + stats_json = json.dumps(stats, ensure_ascii=False, indent=2) + + # 创建输出目录 + if FOLDER_PATHS_AVAILABLE: + output_dir = os.path.join(folder_paths.get_output_directory(), "segmented_models") + else: + output_dir = os.path.join("output", "segmented_models") + os.makedirs(output_dir, exist_ok=True) + + timestamp = int(time.time()) + + # 导出文件 + ply_path = "" + glb_path = "" + + if export_format in ["PLY", "BOTH"]: + if use_full_pointcloud: + # 使用完整点云分割（推荐） + logger.info("使用完整点云分割方法") + ply_path = self._export_segmented_ply_new( + native_full_result, mask_list, target_object_id, + output_dir, timestamp + ) + else: + # 仅使用tracks点分割（可能无效） + logger.info("使用tracks点分割方法（仅使用查询点）") + if stats.get("unique_objects", 0) == 0: + logger.warning("tracks点分割未找到物体，建议启用use_full_pointcloud选项") + ply_path = self._export_segmented_ply( + native_full_result, point_labels, target_object_id, + output_dir, timestamp + ) + + if export_format in ["GLB", "BOTH"]: + glb_path = self._export_segmented_glb( + native_full_result, point_labels, target_object_id, + output_dir, timestamp + ) + + # 生成报告 + report = self._generate_segmentation_report( + stats, ply_path, glb_path, target_object_id, confidence_threshold + ) + report_json = json.dumps(report, ensure_ascii=False, indent=2) + + logger.info("VGGT Mask分割处理完成") + return (stats_json, ply_path, glb_path, report_json) + + except Exception as e: + logger.error(f"Mask分割处理失败: {e}") + import traceback + traceback.print_exc() + + error_msg = f"分割处理失败: {str(e)}" + error_json = json.dumps({"error": str(e)}) + return (error_json, "", "", error_msg) + + def _generate_segmentation_stats(self, point_labels: Dict[int, List[int]]) -> Dict: + """生成分割统计信息""" + object_counts = defaultdict(int) + total_points = len(point_labels) + + for labels in point_labels.values(): + for label in labels: + # 确保key是标准int类型 + label_int = int(label) if hasattr(label, 'item') else int(label) + object_counts[label_int] += 1 + + return { + "total_points": int(total_points), + "object_point_counts": {str(k): int(v) for k, v in object_counts.items()}, # key转为字符串 + "unique_objects": len([obj_id for obj_id in object_counts.keys() if obj_id > 0]), + "background_points": int(object_counts.get(0, 0)) + } + + def _export_segmented_ply_new(self, raw_vggt_result: Dict, mask_sequence: List[np.ndarray], + target_object_id: int, output_dir: str, timestamp: int) -> str: + """导出分割后的PLY文件（新方法）""" + try: + # 首先检查mask序列的有效性 + logger.info(f"检查mask序列: {len(mask_sequence)} 张图像") + for i in range(min(3, len(mask_sequence))): + mask = mask_sequence[i] + unique_vals = np.unique(mask) + object_ratio = np.sum(mask > 0) / mask.size * 100 + logger.info(f" Mask {i}: shape={mask.shape}, unique_values={unique_vals}, object_ratio={object_ratio:.1f}%") + + vertices, colors = create_segmented_pointcloud_new( + raw_vggt_result, mask_sequence, + target_object_id if target_object_id > 0 else None + ) + + if len(vertices) == 0: + logger.warning("分割后没有点，无法导出PLY文件") + return "" + + if target_object_id > 0: + filename = f"vggt_segmented_obj{target_object_id}_{timestamp}.ply" + else: + filename = f"vggt_segmented_all_{timestamp}.ply" + + ply_path = os.path.join(output_dir, filename) + self._write_ply_file(ply_path, vertices, colors) + + logger.info(f"新方法分割PLY已保存: {ply_path} ({len(vertices)} 个点)") + return ply_path + + except Exception as e: + logger.error(f"导出分割PLY失败: {e}") + import traceback + traceback.print_exc() + return "" + + def _export_segmented_ply(self, raw_vggt_result: Dict, point_labels: Dict[int, List[int]], + target_object_id: int, output_dir: str, timestamp: int) -> str: + """导出分割后的PLY文件（旧方法，兼容性保留）""" + try: + vertices, colors = create_segmented_pointcloud( + raw_vggt_result, point_labels, + target_object_id if target_object_id > 0 else None + ) + + if target_object_id > 0: + filename = f"vggt_segmented_obj{target_object_id}_{timestamp}.ply" + else: + filename = f"vggt_segmented_all_{timestamp}.ply" + + ply_path = os.path.join(output_dir, filename) + self._write_ply_file(ply_path, vertices, colors) + + logger.info(f"分割PLY已保存: {ply_path} ({len(vertices)} 个点)") + return ply_path + + except Exception as e: + logger.error(f"导出分割PLY失败: {e}") + return "" + + def _export_segmented_glb(self, raw_vggt_result: Dict, point_labels: Dict[int, List[int]], + target_object_id: int, output_dir: str, timestamp: int) -> str: + """导出分割后的GLB文件""" + if not VGGT_UTILS_AVAILABLE or not predictions_to_glb: + logger.warning("GLB导出功能不可用") + return "" + + try: + # 由于分割后的点数量很少，GLB导出容易失败 + # 暂时只支持PLY导出，避免复杂的3D网格生成 + logger.info("分割后点数量较少，暂时跳过GLB导出，请使用PLY文件") + return "" + + except Exception as e: + logger.error(f"导出分割GLB失败: {e}") + return "" + + def _create_segmented_predictions(self, raw_vggt_result: Dict, point_labels: Dict[int, List[int]], + target_object_id: int) -> Dict: + """创建分割版本的预测数据""" + # 创建分割mask + vertices, colors = create_segmented_pointcloud( + raw_vggt_result, point_labels, + target_object_id if target_object_id > 0 else None + ) + + # 获取原始相机参数 + cameras = raw_vggt_result.get('cameras', {}) + + # 处理相机外参 - 使用与原生节点相同的简单格式 + extrinsic = cameras.get('extrinsic', None) + if extrinsic is not None: + if isinstance(extrinsic, torch.Tensor): + extrinsic = extrinsic.cpu().numpy() + else: + # 使用单位矩阵作为默认值 (S, 3, 4) 格式 + extrinsic = np.eye(4)[None, :3, :] + + # 确保形状正确 + if extrinsic.ndim == 3: + if extrinsic.shape[-2] == 4 and extrinsic.shape[-1] == 4: + # (S, 4, 4) -> (S, 3, 4) + extrinsic = extrinsic[:, :3, :] + + # 重塑点云数据为简单的3D格式 - 参考原生节点的格式 + n_points = len(vertices) + h = int(np.sqrt(n_points)) + w = n_points // h + if h * w != n_points: + # 无法完美重塑为方形，使用接近方形的布局 + h = int(np.ceil(np.sqrt(n_points))) + # 填充到方形 + pad_size = h * h - n_points + if pad_size > 0: + # 用最后一个点填充 + last_vertex = vertices[-1:].repeat(pad_size, axis=0) + last_color = colors[-1:].repeat(pad_size, axis=0) + vertices = np.vstack([vertices, last_vertex]) + colors = np.vstack([colors, last_color]) + w = h + + # 重塑为 (1, H, W, 3) 格式（与原生节点兼容） + world_points = vertices.reshape(1, h, w, 3) + images_color = colors.reshape(1, h, w, 3) / 255.0 # 归一化到[0,1] + + # 构建与原生节点兼容的格式 + result = { + 'world_points_from_depth': world_points, + 'depth_conf': np.ones((1, h, w)), # 假设所有点置信度为1 + 'images': images_color, + 'extrinsic': extrinsic, + } + + return result + + def _write_ply_file(self, filepath: str, vertices: np.ndarray, colors: np.ndarray = None): + """写入PLY格式文件""" + try: + with open(filepath, 'w') as f: + f.write("ply\n") + f.write("format ascii 1.0\n") + f.write(f"element vertex {len(vertices)}\n") + f.write("property float x\n") + f.write("property float y\n") + f.write("property float z\n") + + if colors is not None: + f.write("property uchar red\n") + f.write("property uchar green\n") + f.write("property uchar blue\n") + + f.write("end_header\n") + + for i in range(len(vertices)): + x, y, z = vertices[i] + if colors is not None: + r, g, b = colors[i] + f.write(f"{x:.6f} {y:.6f} {z:.6f} {r} {g} {b}\n") + else: + f.write(f"{x:.6f} {y:.6f} {z:.6f}\n") + + except Exception as e: + logger.error(f"写入PLY文件失败: {e}") + raise + + def _generate_segmentation_report(self, stats: Dict, ply_path: str, glb_path: str, + target_object_id: int, confidence_threshold: float) -> Dict: + """生成详细的分割报告""" + return { + "processing_parameters": { + "target_object_id": target_object_id, + "confidence_threshold": confidence_threshold + }, + "segmentation_statistics": stats, + "output_files": { + "ply_file": ply_path, + "glb_file": glb_path + }, + "processing_time": time.time(), + "status": "success" if (ply_path or glb_path) else "failed" + } + +# ----------------------------------------------------------------------------- +# 节点注册 +# ----------------------------------------------------------------------------- + +NODE_CLASS_MAPPINGS = { + "VGGTMaskProcessorNode": VGGTMaskProcessorNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "VGGTMaskProcessorNode": "🎭 VGGT Mask Processor", +} diff --git a/comfyui_vggt_native_nodes.py b/comfyui_vggt_native_nodes.py new file mode 100644 index 0000000..4613821 --- /dev/null +++ b/comfyui_vggt_native_nodes.py @@ -0,0 +1,876 @@ +import os +import json +import tempfile +from typing import List, Any, Dict +import logging +import struct + +import cv2 +import numpy as np +import torch +from PIL import Image + +# 尝试导入 ComfyUI 的类型标记 +try: + from comfy.comfy_types import IO +except ImportError: + class IO: + VIDEO = "VIDEO" + IMAGE = "IMAGE" + +# 导入原生VGGT接口 +try: + from .vggt_native_interface import VGGTNativeInterface, VGGTImageProcessor, VGGTResultProcessor + VGGT_NATIVE_AVAILABLE = True +except ImportError: + VGGTNativeInterface = None + VGGTImageProcessor = None + VGGTResultProcessor = None + VGGT_NATIVE_AVAILABLE = False + +# 导入模型加载器 +try: + from .vggt_model_loader import VVLVGGTLoader + MODEL_LOADER_AVAILABLE = True +except ImportError: + VVLVGGTLoader = None + MODEL_LOADER_AVAILABLE = False + +# 导入ComfyUI的路径管理 +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + folder_paths = None + FOLDER_PATHS_AVAILABLE = False + +# 导入原节点的工具函数 +try: + from .comfyui_vggt_nodes import ( + _create_traj_preview, + _create_insufficient_data_image, + predictions_to_glb, + TRIMESH_AVAILABLE, + MATPLOTLIB_AVAILABLE, + SCIPY_AVAILABLE + ) + VGGT_UTILS_AVAILABLE = True +except ImportError: + _create_traj_preview = None + _create_insufficient_data_image = None + predictions_to_glb = None + TRIMESH_AVAILABLE = False + MATPLOTLIB_AVAILABLE = False + SCIPY_AVAILABLE = False + VGGT_UTILS_AVAILABLE = False + +# 配置日志 +logger = logging.getLogger('vvl_vggt_native_nodes') + +# ----------------------------------------------------------------------------- +# 原生VGGT推理函数（增强版） +# ----------------------------------------------------------------------------- + +def run_vggt_full_native_inference(images: torch.Tensor, model_instance, device, + enable_tracks: bool = True, + grid_size: int = 8, + query_region: float = 0.8) -> Dict: + """使用原生VGGT接口进行完整推理，启用所有原生输出分支""" + if not VGGT_NATIVE_AVAILABLE: + raise ImportError("VGGT Native Interface not available") + + try: + # 创建原生接口 + vggt_interface = VGGTNativeInterface(model_instance, device) + + # 确定启用的分支 + branches = ['cameras', 'depth', 'points', 'points_from_depth'] + + # 如果启用tracks，需要生成查询点 + query_points = None + if enable_tracks: + branches.append('tracks') + # 自动生成查询点：在图像指定区域均匀采样点 + B, S, C, H, W = images.shape if images.dim() == 5 else (1, *images.shape) + + # 计算查询区域 + margin_h = int(H * (1 - query_region) / 2) + margin_w = int(W * (1 - query_region) / 2) + + # 生成网格查询点 + y_coords = torch.linspace(margin_h, H - margin_h, grid_size) + x_coords = torch.linspace(margin_w, W - margin_w, grid_size) + grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing='ij') + + # 生成查询点坐标 (N_points, 2)，格式为 (x, y) + query_points = torch.stack([grid_x.flatten(), grid_y.flatten()], dim=1).to(device) + logger.info(f"Generated {len(query_points)} query points ({grid_size}x{grid_size} grid) in {query_region*100:.0f}% region") + + # 执行选择性推理以获取所有分支的输出 + results = vggt_interface.selective_inference( + images, + branches=branches, + query_points=query_points + ) + + # 添加原始图像到结果中（用于后续处理） + if images.dim() == 5: # (1, N, C, H, W) + images_formatted = images.squeeze(0).permute(0, 2, 3, 1) # (N, H, W, C) + else: # (N, C, H, W) + images_formatted = images.permute(0, 2, 3, 1) # (N, H, W, C) + results['images'] = images_formatted.cpu().numpy() + + # 添加查询点信息 + if query_points is not None: + results['query_points'] = query_points.cpu().numpy() + + logger.info(f"VGGT full native inference completed. Results keys: {list(results.keys())}") + if 'tracks' in results: + track_data = results['tracks'] + logger.info(f"Tracks data structure: {list(track_data.keys())}") + if 'track_list' in track_data: + track_shape = track_data['track_list'].shape if hasattr(track_data['track_list'], 'shape') else 'unknown' + logger.info(f"Track list shape: {track_shape}") + + return results + + except Exception as e: + logger.error(f"VGGT full native inference failed: {e}") + raise + +def preprocess_images_native(images_list: List[np.ndarray]) -> torch.Tensor: + """使用原生VGGT预处理器处理图像""" + if not VGGT_NATIVE_AVAILABLE: + raise ImportError("VGGT Native Interface not available") + + try: + return VGGTImageProcessor.preprocess_images(images_list) + except Exception as e: + logger.error(f"Native image preprocessing failed: {e}") + raise + +# ----------------------------------------------------------------------------- +# 完整原生输出节点 +# ----------------------------------------------------------------------------- + +class VGGTNativeFullOutputNode: + """VGGT 原生完整输出节点 - 输出所有原生分支数据，包括tracks""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "vggt_model": ("VVL_VGGT_MODEL", { + "tooltip": "来自VVLVGGTLoader的VGGT模型实例" + }), + "images": ("IMAGE", { + "tooltip": "图片序列输入（必需）- 支持1-200帧图像" + }), + }, + "optional": { + "enable_tracks": ("BOOLEAN", { + "default": True, + "tooltip": "是否启用点追踪分支（需要自动生成查询点）" + }), + "grid_size": ("INT", { + "default": 8, "min": 4, "max": 16, "step": 1, + "tooltip": "查询点网格大小（grid_size x grid_size个点）" + }), + "query_region": ("FLOAT", { + "default": 0.8, "min": 0.1, "max": 1.0, "step": 0.05, + "tooltip": "查询点区域比例（1.0=全图，0.8=中央80%区域）" + }), + "confidence_threshold": ("FLOAT", { + "default": 90.0, "min": 0.0, "max": 100.0, "step": 0.1, + "tooltip": "置信度阈值(%)，用于过滤3D点云中的低置信度点" + }), + "show_cameras": ("BOOLEAN", { + "default": True, + "tooltip": "是否在3D模型中显示相机位置" + }), + "mask_black_bg": ("BOOLEAN", { + "default": False, + "tooltip": "是否过滤黑色背景点" + }), + "mask_white_bg": ("BOOLEAN", { + "default": False, + "tooltip": "是否过滤白色背景点" + }), + "mask_sky": ("BOOLEAN", { + "default": False, + "tooltip": "是否过滤天空点" + }), + } + } + + RETURN_TYPES = ( + "RAW_VGGT_RESULT", # 完整原生VGGT字典 + "STRING", # 完整结果summary JSON + "STRING", # Tracks JSON数据（如果有） + "STRING", # 相机内参JSON（与原节点兼容） + "STRING", # 相机姿态JSON（与原节点兼容） + "IMAGE", # 轨迹可视化预览 + "STRING", # 3D模型GLB文件路径 + "STRING", # 原生点云PLY文件路径 + ) + RETURN_NAMES = ( + "native_full_result", + "result_summary", + "tracks_json", + "intrinsics_json", + "poses_json", + "trajectory_preview", + "model_3d_glb_path", + "pointcloud_ply_path", + ) + OUTPUT_TOOLTIPS = [ + "完整的VGGT原生推理结果（包含所有分支数据：cameras, depth, points, points_from_depth, tracks等）", + "结果摘要信息（JSON格式，包含数据形状和结构信息）", + "点追踪数据的JSON格式输出（如果启用了tracks分支）", + "相机内参数据（JSON格式，与原节点输出兼容）", + "相机位姿数据（JSON格式，与原节点输出兼容）", + "相机轨迹3D可视化预览图像", + "3D模型GLB文件路径（可连接到Preview3D或其他节点）", + "原生点云PLY文件路径（VGGT原生格式，便于外部处理）" + ] + OUTPUT_NODE = True + FUNCTION = "get_native_full_output" + CATEGORY = "💃VVL/VGGT Native" + + def get_native_full_output(self, vggt_model: Dict, images, + enable_tracks: bool = True, + grid_size: int = 8, + query_region: float = 0.8, + confidence_threshold: float = 90.0, + show_cameras: bool = True, + mask_black_bg: bool = False, + mask_white_bg: bool = False, + mask_sky: bool = False): + """获取VGGT完整原生输出""" + logger.info("开始VGGT原生完整输出") + + # 获取模型实例和设备 + model_instance = vggt_model.get("model") + device = vggt_model.get("device") + + if model_instance is None: + raise ValueError("无效的VGGT模型实例") + + logger.info(f"使用设备: {device}") + + # 处理图片序列输入 + image_list = [] + if isinstance(images, torch.Tensor): + images_np = images.cpu().numpy() + for i in range(images_np.shape[0]): + img = images_np[i] + if img.max() <= 1.0: + img = (img * 255).astype(np.uint8) + else: + img = img.astype(np.uint8) + image_list.append(img) + else: + image_list.extend(images) + + if not image_list: + raise ValueError("未能提取到有效的图像数据") + + logger.info(f"总共处理 {len(image_list)} 张图像") + + # 限制图像数量 + max_frames = 200 + if len(image_list) > max_frames: + logger.warning(f"图像数量超过限制，采样到 {max_frames} 帧") + indices = np.linspace(0, len(image_list) - 1, max_frames, dtype=int) + image_list = [image_list[i] for i in indices] + + try: + # 使用原生VGGT预处理 + processed_images = preprocess_images_native(image_list) + processed_images = processed_images.to(device) + + # 应用遮罩（如果需要） + if VGGT_NATIVE_AVAILABLE: + mask_types = [] + if mask_black_bg: + mask_types.append('black_bg') + if mask_white_bg: + mask_types.append('white_bg') + if mask_sky: + mask_types.append('sky') + + for mask_type in mask_types: + processed_images = VGGTImageProcessor.apply_masking(processed_images, mask_type) + + # 执行完整的原生推理，启用所有分支 + raw_results = run_vggt_full_native_inference( + processed_images, model_instance, device, + enable_tracks, grid_size, query_region + ) + + # 生成结果摘要 + summary = self._generate_result_summary(raw_results, { + "enable_tracks": enable_tracks, + "grid_size": grid_size, + "query_region": query_region, + "confidence_threshold": confidence_threshold, + "show_cameras": show_cameras, + "mask_black_bg": mask_black_bg, + "mask_white_bg": mask_white_bg, + "mask_sky": mask_sky + }) + summary_json = json.dumps(summary, ensure_ascii=False, indent=2) + + # 生成tracks JSON输出 + tracks_json = "" + if 'tracks' in raw_results and raw_results['tracks']: + try: + track_data = raw_results['tracks'] + logger.info(f"处理tracks数据，keys: {list(track_data.keys())}") + + # 将tensor转换为可序列化的格式 + tracks_export = {} + for key, value in track_data.items(): + logger.info(f"处理tracks key: {key}, type: {type(value)}") + if hasattr(value, 'cpu'): + # PyTorch tensor + value_np = value.cpu().numpy() + tracks_export[key] = value_np.tolist() + logger.info(f"转换tensor {key}, shape: {value_np.shape}") + elif isinstance(value, np.ndarray): + # NumPy array + tracks_export[key] = value.tolist() + logger.info(f"转换numpy {key}, shape: {value.shape}") + elif isinstance(value, (list, tuple)): + # 列表或元组，递归处理内部元素 + tracks_export[key] = self._convert_to_serializable(value) + logger.info(f"转换列表 {key}, length: {len(value)}") + else: + # 其他类型直接赋值 + tracks_export[key] = value + logger.info(f"直接赋值 {key}, type: {type(value)}") + + # 添加详细的元数据 + metadata = { + "format": "coordinates_per_frame", + "description": "Point tracking data with coordinates, visibility and confidence", + "coordinate_format": "(x, y) pixel coordinates", + "data_structure": {} + } + + # 添加每个字段的形状信息 + for key, value in tracks_export.items(): + if isinstance(value, list) and len(value) > 0: + if isinstance(value[0], list): + # 多维列表 + shape = [len(value)] + temp = value[0] + while isinstance(temp, list) and len(temp) > 0: + shape.append(len(temp)) + temp = temp[0] if isinstance(temp[0], list) else None + if temp is None: + break + metadata["data_structure"][key] = { + "shape": shape, + "description": self._get_track_field_description(key) + } + else: + metadata["data_structure"][key] = { + "length": len(value), + "description": self._get_track_field_description(key) + } + else: + metadata["data_structure"][key] = { + "type": type(value).__name__, + "description": self._get_track_field_description(key) + } + + tracks_json = json.dumps({ + "tracks_data": tracks_export, + "metadata": metadata + }, ensure_ascii=False, indent=2) + + logger.info("tracks数据序列化成功") + + except Exception as e: + logger.error(f"Failed to serialize tracks data: {e}") + import traceback + traceback.print_exc() + tracks_json = json.dumps({ + "error": f"Failed to serialize tracks: {str(e)}", + "available_keys": list(raw_results['tracks'].keys()) if 'tracks' in raw_results else [], + "debug_info": "Check logs for detailed error information" + }, ensure_ascii=False, indent=2) + else: + tracks_json = json.dumps({"message": "No tracks data available or tracks not enabled"}) + + # 生成与原节点兼容的相机参数JSON输出 + intrinsics_json = "" + poses_json = "" + if 'cameras' in raw_results: + try: + cameras_data = raw_results['cameras'] + if isinstance(cameras_data, dict): + if 'intrinsic' in cameras_data: + intrinsic = cameras_data['intrinsic'] + if hasattr(intrinsic, 'cpu'): + intrinsic = intrinsic.cpu().numpy() + intrinsics_json = json.dumps(intrinsic.tolist(), ensure_ascii=False, indent=2) + + if 'extrinsic' in cameras_data: + extrinsic = cameras_data['extrinsic'] + if hasattr(extrinsic, 'cpu'): + extrinsic = extrinsic.cpu().numpy() + poses_json = json.dumps(extrinsic.tolist(), ensure_ascii=False, indent=2) + except Exception as e: + logger.warning(f"Failed to generate camera JSON: {e}") + intrinsics_json = json.dumps({"error": f"Failed to extract intrinsics: {str(e)}"}) + poses_json = json.dumps({"error": f"Failed to extract poses: {str(e)}"}) + else: + intrinsics_json = json.dumps({"message": "No camera data available"}) + poses_json = json.dumps({"message": "No camera data available"}) + + # 生成轨迹预览 + trajectory_preview = self._create_trajectory_preview(raw_results) + + # 生成3D模型GLB文件 + model_3d_glb_path = self._generate_3d_model_official( + raw_results, + confidence_threshold, + show_cameras, + mask_black_bg, + mask_white_bg, + mask_sky + ) + + # 生成原生点云PLY文件 + pointcloud_ply_path = self._export_native_pointcloud(raw_results, confidence_threshold) + + logger.info("VGGT原生完整输出完成") + return (raw_results, summary_json, tracks_json, intrinsics_json, poses_json, + trajectory_preview, model_3d_glb_path, pointcloud_ply_path) + + except Exception as e: + logger.error(f"VGGT原生完整推理失败: {e}") + error_msg = f"VGGT原生完整推理失败: {str(e)}" + empty_result = {"error": str(e)} + error_json = json.dumps({"error": str(e)}) + error_image = self._create_error_image() + return (empty_result, error_msg, error_json, error_json, error_json, + error_image, "", "") + + def _generate_result_summary(self, results: Dict, parameters: Dict = None) -> Dict: + """生成详细的结果摘要信息""" + summary = { + "available_branches": list(results.keys()), + "data_info": {}, + "metadata": { + "generated_by": "VGGTNativeFullOutputNode", + "description": "Complete VGGT native inference results with all branches" + } + } + + # 添加处理参数到元数据 + if parameters is not None: + summary["metadata"]["processing_parameters"] = parameters + + for key, value in results.items(): + if key == "images": + if isinstance(value, np.ndarray): + summary["data_info"][key] = { + "type": "numpy.ndarray", + "shape": list(value.shape), + "dtype": str(value.dtype), + "description": "Input images in (N, H, W, C) format" + } + elif key == "cameras": + if isinstance(value, dict): + summary["data_info"][key] = { + "type": "dict", + "keys": list(value.keys()), + "description": "Camera parameters (extrinsic and intrinsic matrices)" + } + for sub_key, sub_value in value.items(): + if hasattr(sub_value, 'shape'): + summary["data_info"][key][f"{sub_key}_shape"] = list(sub_value.shape) + elif key == "depth": + if isinstance(value, dict): + summary["data_info"][key] = { + "type": "dict", + "keys": list(value.keys()), + "description": "Depth maps and confidence" + } + for sub_key, sub_value in value.items(): + if hasattr(sub_value, 'shape'): + summary["data_info"][key][f"{sub_key}_shape"] = list(sub_value.shape) + elif key == "tracks": + if isinstance(value, dict): + summary["data_info"][key] = { + "type": "dict", + "keys": list(value.keys()), + "description": "Point tracking data (coordinates, visibility, confidence)" + } + for sub_key, sub_value in value.items(): + if hasattr(sub_value, 'shape'): + summary["data_info"][key][f"{sub_key}_shape"] = list(sub_value.shape) + if sub_key == 'track_list' and hasattr(sub_value, 'shape') and len(sub_value.shape) >= 2: + summary["data_info"][key]["n_points"] = sub_value.shape[-3] if len(sub_value.shape) >= 3 else "unknown" + summary["data_info"][key]["n_frames"] = sub_value.shape[-2] if len(sub_value.shape) >= 2 else "unknown" + elif key == "query_points": + if isinstance(value, np.ndarray): + summary["data_info"][key] = { + "type": "numpy.ndarray", + "shape": list(value.shape), + "dtype": str(value.dtype), + "description": "Query points used for tracking (N_points, 2) format (x, y)", + "n_points": value.shape[0] if len(value.shape) >= 1 else 0 + } + elif hasattr(value, 'shape'): + summary["data_info"][key] = { + "type": type(value).__name__, + "shape": list(value.shape), + "description": f"Raw {key} data from VGGT" + } + elif isinstance(value, dict): + summary["data_info"][key] = { + "type": "dict", + "keys": list(value.keys()), + "description": f"Structured {key} data" + } + else: + summary["data_info"][key] = { + "type": type(value).__name__, + "description": f"Raw {key} data from VGGT" + } + + return summary + + def _create_error_image(self) -> torch.Tensor: + """创建错误提示图像""" + if VGGT_UTILS_AVAILABLE and _create_insufficient_data_image: + return _create_insufficient_data_image() + else: + # 创建简单的错误图像 + canvas = np.ones((400, 600, 3), dtype=np.float32) * 0.9 + # 这里可以添加文字，但为了简单起见直接返回灰色图像 + return torch.from_numpy(canvas).unsqueeze(0) + + def _create_trajectory_preview(self, raw_results: Dict) -> torch.Tensor: + """生成轨迹预览""" + if not VGGT_UTILS_AVAILABLE or not _create_traj_preview: + return self._create_error_image() + + try: + if 'cameras' in raw_results: + extrinsic_for_preview = raw_results['cameras']['extrinsic'] + + # 应用与GLB生成相同的维度处理逻辑 + if isinstance(extrinsic_for_preview, torch.Tensor): + if extrinsic_for_preview.ndim == 4 and extrinsic_for_preview.shape[0] == 1: + extrinsic_for_preview = extrinsic_for_preview.squeeze(0) # (1, S, 3, 4) -> (S, 3, 4) + elif isinstance(extrinsic_for_preview, np.ndarray): + if extrinsic_for_preview.ndim == 4 and extrinsic_for_preview.shape[0] == 1: + extrinsic_for_preview = np.squeeze(extrinsic_for_preview, axis=0) # (1, S, 3, 4) -> (S, 3, 4) + + logger.info(f"轨迹预览外参形状: {extrinsic_for_preview.shape}") + return _create_traj_preview(extrinsic_for_preview) + else: + return self._create_error_image() + except Exception as e: + logger.warning(f"轨迹预览生成失败: {e}") + return self._create_error_image() + + def _generate_3d_model_official(self, raw_results: Dict, confidence_threshold: float, show_cameras: bool, + mask_black_bg: bool = False, mask_white_bg: bool = False, mask_sky: bool = False) -> str: + """使用官方VGGT的predictions_to_glb函数生成高质量3D模型""" + if not VGGT_UTILS_AVAILABLE or not predictions_to_glb: + logger.warning("VGGT工具函数不可用，跳过3D模型生成") + return "" + + try: + # 创建输出目录 + if FOLDER_PATHS_AVAILABLE: + output_dir = os.path.join(folder_paths.get_output_directory(), "3d") + else: + output_dir = os.path.join("output", "3d") + os.makedirs(output_dir, exist_ok=True) + + # 生成文件名 + import time + timestamp = int(time.time()) + model_path = os.path.join(output_dir, f"vggt_native_{timestamp}.glb") + + logger.info(f"开始生成原生VGGT质量的3D模型: {model_path}") + + # 准备官方格式的预测数据 + world_points_from_depth = raw_results.get('points_from_depth') + + # depth_conf: 兼容多种返回格式 + depth_conf = None + if 'depth' in raw_results: + if isinstance(raw_results['depth'], dict): + depth_conf = raw_results['depth'].get('confidence') + else: + depth_conf = raw_results.get('depth_conf') + + # images + images_tensor = raw_results.get('images') + + # extrinsic + extrinsic_mat = None + if 'cameras' in raw_results: + if isinstance(raw_results['cameras'], dict): + extrinsic_mat = raw_results['cameras'].get('extrinsic') + else: + extrinsic_mat = raw_results.get('extrinsic') + + # 特别处理 extrinsic: 去掉 batch 维 (1, S, 3, 4) -> (S, 3, 4) + if isinstance(extrinsic_mat, torch.Tensor): + if extrinsic_mat.ndim == 4 and extrinsic_mat.shape[0] == 1: + extrinsic_mat = extrinsic_mat.squeeze(0) + elif isinstance(extrinsic_mat, np.ndarray): + if extrinsic_mat.ndim == 4 and extrinsic_mat.shape[0] == 1: + extrinsic_mat = np.squeeze(extrinsic_mat, axis=0) + + # 维度兼容处理 + images_tensor_proc = images_tensor + if isinstance(images_tensor_proc, torch.Tensor): + if images_tensor_proc.ndim == 5 and images_tensor_proc.shape[0] == 1: + images_tensor_proc = images_tensor_proc.squeeze(0) + elif isinstance(images_tensor_proc, np.ndarray): + if images_tensor_proc.ndim == 5 and images_tensor_proc.shape[0] == 1: + images_tensor_proc = np.squeeze(images_tensor_proc, axis=0) + + wpfd_proc = world_points_from_depth + if isinstance(wpfd_proc, torch.Tensor): + if wpfd_proc is not None and wpfd_proc.ndim == 5 and wpfd_proc.shape[0] == 1: + wpfd_proc = wpfd_proc.squeeze(0) + elif isinstance(wpfd_proc, np.ndarray): + if wpfd_proc is not None and wpfd_proc.ndim == 5 and wpfd_proc.shape[0] == 1: + wpfd_proc = np.squeeze(wpfd_proc, axis=0) + + predictions_formatted = { + 'world_points_from_depth': wpfd_proc, + 'depth_conf': depth_conf, + 'images': images_tensor_proc, + 'extrinsic': extrinsic_mat, + } + + # 如果没有depth-based points，尝试使用原始点云 + if predictions_formatted['world_points_from_depth'] is None: + if 'points' in raw_results: + logger.info("Using point_map as fallback for world_points") + if isinstance(raw_results['points'], dict): + predictions_formatted['world_points'] = raw_results['points']['point_map'] + predictions_formatted['world_points_conf'] = raw_results['points']['confidence'] + else: + predictions_formatted['world_points'] = raw_results['points'] + + # 确保数据格式正确（转换为numpy） + for key, value in predictions_formatted.items(): + if value is not None and isinstance(value, torch.Tensor): + predictions_formatted[key] = value.cpu().numpy() + + # 使用官方VGGT的predictions_to_glb函数 + scene_3d = predictions_to_glb( + predictions_formatted, + conf_thres=confidence_threshold, + filter_by_frames="all", + mask_black_bg=mask_black_bg, + mask_white_bg=mask_white_bg, + show_cam=show_cameras, + mask_sky=mask_sky, + target_dir=None, + prediction_mode="Depthmap and Camera Branch" + ) + + # 导出为GLB文件 + scene_3d.export(model_path) + + logger.info(f"原生VGGT质量3D模型已保存到: {model_path}") + return model_path + + except Exception as e: + logger.error(f"生成3D模型失败: {e}") + return "" + + def _export_native_pointcloud(self, raw_results: Dict, confidence_threshold: float) -> str: + """导出原生格式的点云PLY文件""" + try: + # 创建输出目录 + if FOLDER_PATHS_AVAILABLE: + output_dir = os.path.join(folder_paths.get_output_directory(), "pointclouds") + else: + output_dir = os.path.join("output", "pointclouds") + os.makedirs(output_dir, exist_ok=True) + + # 生成文件名 + import time + timestamp = int(time.time()) + ply_path = os.path.join(output_dir, f"vggt_native_{timestamp}.ply") + + # 从原生结果提取点云数据 + if 'points_from_depth' in raw_results: + points = raw_results['points_from_depth'] + logger.info("使用 points_from_depth 生成PLY") + elif 'points' in raw_results: + if isinstance(raw_results['points'], dict): + points = raw_results['points']['point_map'] + else: + points = raw_results['points'] + logger.info("使用 points 生成PLY") + else: + logger.warning("No point cloud data found in raw results") + return "" + + # 转换为numpy数组 + if isinstance(points, torch.Tensor): + points_np = points.cpu().numpy() + else: + points_np = points + + # 去掉batch维度 + if points_np.ndim == 5 and points_np.shape[0] == 1: + points_np = np.squeeze(points_np, axis=0) + + # 重塑点云数据 (S, H, W, 3) -> (S*H*W, 3) + if points_np.ndim == 4: + points_np = points_np.reshape(-1, 3) + elif points_np.ndim == 3: + points_np = points_np.reshape(-1, 3) + + # 过滤无效点 + valid_mask = ~np.isnan(points_np).any(axis=1) & ~np.isinf(points_np).any(axis=1) + points_np = points_np[valid_mask] + + if len(points_np) == 0: + logger.warning("No valid points found") + return "" + + # 获取颜色信息 + colors_rgb = None + if 'images' in raw_results: + images = raw_results['images'] + if isinstance(images, torch.Tensor): + images_np = images.cpu().numpy() + else: + images_np = images + + # 去掉batch维度 + if images_np.ndim == 5 and images_np.shape[0] == 1: + images_np = np.squeeze(images_np, axis=0) + + # 处理图像格式 (S, H, W, 3) or (S, 3, H, W) + if images_np.ndim == 4: + if images_np.shape[1] == 3: # (S, 3, H, W) -> (S, H, W, 3) + images_np = np.transpose(images_np, (0, 2, 3, 1)) + + # 重塑为点云颜色 (S*H*W, 3) + colors_rgb = images_np.reshape(-1, 3) + colors_rgb = colors_rgb[valid_mask] # 应用相同的mask + + # 确保颜色值在0-255范围 + if colors_rgb.max() <= 1.0: + colors_rgb = (colors_rgb * 255).astype(np.uint8) + else: + colors_rgb = colors_rgb.astype(np.uint8) + + # 基于置信度过滤 + if confidence_threshold > 0 and len(points_np) > 100: + # 使用距离作为置信度的简单替代 + center = np.mean(points_np, axis=0) + distances = np.linalg.norm(points_np - center, axis=1) + threshold = np.percentile(distances, confidence_threshold) + conf_mask = distances <= threshold + points_np = points_np[conf_mask] + if colors_rgb is not None: + colors_rgb = colors_rgb[conf_mask] + + # 写入PLY文件 + self._write_ply_file(ply_path, points_np, colors_rgb) + + logger.info(f"原生点云PLY已保存到: {ply_path} ({len(points_np)} 个点)") + return ply_path + + except Exception as e: + logger.error(f"导出点云PLY失败: {e}") + raise + + def _write_ply_file(self, filepath: str, vertices: np.ndarray, colors: np.ndarray = None): + """写入PLY格式点云文件""" + try: + with open(filepath, 'w') as f: + # PLY header + f.write("ply\n") + f.write("format ascii 1.0\n") + f.write(f"element vertex {len(vertices)}\n") + f.write("property float x\n") + f.write("property float y\n") + f.write("property float z\n") + + if colors is not None: + f.write("property uchar red\n") + f.write("property uchar green\n") + f.write("property uchar blue\n") + + f.write("end_header\n") + + # 写入顶点数据 + for i in range(len(vertices)): + x, y, z = vertices[i] + if colors is not None: + r, g, b = colors[i] + f.write(f"{x:.6f} {y:.6f} {z:.6f} {r} {g} {b}\n") + else: + f.write(f"{x:.6f} {y:.6f} {z:.6f}\n") + + except Exception as e: + logger.error(f"写入PLY文件失败: {e}") + raise + + def _convert_to_serializable(self, obj): + """递归转换对象为JSON可序列化格式""" + if hasattr(obj, 'cpu'): + # PyTorch tensor + return obj.cpu().numpy().tolist() + elif isinstance(obj, np.ndarray): + # NumPy array + return obj.tolist() + elif isinstance(obj, (list, tuple)): + # 列表或元组，递归处理每个元素 + return [self._convert_to_serializable(item) for item in obj] + elif isinstance(obj, dict): + # 字典，递归处理每个值 + return {key: self._convert_to_serializable(value) for key, value in obj.items()} + else: + # 其他类型直接返回 + return obj + + def _get_track_field_description(self, field_name: str) -> str: + """获取tracks字段的描述信息""" + descriptions = { + "track_list": "Point coordinates per frame - shape: (N_points, N_frames, 2) where last dim is (x, y)", + "visibility_score": "Visibility scores per point per frame - shape: (N_points, N_frames)", + "confidence_score": "Confidence scores per point per frame - shape: (N_points, N_frames)", + "coords": "Alternative coordinate format", + "vis": "Alternative visibility format", + "conf": "Alternative confidence format" + } + return descriptions.get(field_name, f"Data field: {field_name}") + +# ----------------------------------------------------------------------------- +# 节点注册 +# ----------------------------------------------------------------------------- + +NODE_CLASS_MAPPINGS = { + "VGGTNativeFullOutputNode": VGGTNativeFullOutputNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "VGGTNativeFullOutputNode": "🔬 VGGT Native Full Output", +} + +# 如果模型加载器可用，添加到映射中 +if MODEL_LOADER_AVAILABLE: + NODE_CLASS_MAPPINGS["VVLVGGTLoader"] = VVLVGGTLoader + NODE_DISPLAY_NAME_MAPPINGS["VVLVGGTLoader"] = "VVL VGGT Model Loader" \ No newline at end of file diff --git a/glb_point_cloud_processor.py b/glb_point_cloud_processor.py deleted file mode 100755 index e99fdf0..0000000 --- a/glb_point_cloud_processor.py +++ /dev/null @@ -1,2034 +0,0 @@ -import os -import json -import tempfile -import logging -import struct -import time -from typing import List, Any, Dict, Tuple - -import numpy as np -import torch - -# 导入ComfyUI的路径管理 -try: - import folder_paths - FOLDER_PATHS_AVAILABLE = True -except ImportError: - folder_paths = None - FOLDER_PATHS_AVAILABLE = False - -# 导入trimesh用于GLB文件处理 -try: - import trimesh - TRIMESH_AVAILABLE = True -except ImportError: - TRIMESH_AVAILABLE = False - trimesh = None - -# 配置日志 -logger = logging.getLogger('glb_point_cloud_processor') - -class GLBPointCloudProcessor: - """GLB点云文件处理器 - 专门用于删除黑色点和暗色点优化""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "glb_file_path": ("STRING", { - "default": "", - "tooltip": "GLB点云文件路径：输入需要处理的GLB格式点云文件路径。支持绝对路径(如C:/path/to/file.glb)或相对于ComfyUI output目录的相对路径(如pointcloud.glb)。GLB文件应包含点云数据，支持带颜色信息的点云处理" - }), - }, - "optional": { - "black_threshold": ("INT", { - "default": 30, "min": 0, "max": 255, "step": 1, - "tooltip": "黑色阈值控制：RGB三色总和小于此值的点将被删除。0=仅删除纯黑点(0,0,0)；30=删除深黑色点；60=删除较暗点；100=删除中等暗度点；255=删除所有点。调大=删除更多暗色点，点云更亮净；调小=保留更多暗色细节，但可能有噪点" - }), - "output_filename": ("STRING", { - "default": "clean_pointcloud", - "tooltip": "输出GLB文件名：处理后点云的保存文件名，系统会自动添加.glb扩展名。文件将保存到ComfyUI的output目录中。建议使用描述性名称如'clean_pointcloud'、'filtered_scan'等便于识别" - }), - } - } - - RETURN_TYPES = ( - "STRING", # 处理后的GLB文件路径 - ) - RETURN_NAMES = ( - "processed_glb_path", - ) - OUTPUT_TOOLTIPS = [ - "处理后的GLB文件完整路径", - ] - OUTPUT_NODE = True - FUNCTION = "process_point_cloud" - CATEGORY = "💃VVL/Point Cloud Cleaning" - - def process_point_cloud(self, - glb_file_path: str, - black_threshold: int = 30, - output_filename: str = "clean_pointcloud"): - """ - 处理GLB点云文件，删除黑色点和暗色点 - """ - - processing_log = [] - processing_log.append("开始GLB点云黑色点清理...") - - # 检查依赖 - if not TRIMESH_AVAILABLE: - error_msg = "trimesh库不可用，无法处理GLB文件" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("",) - - # 验证输入文件路径 - if not glb_file_path or not glb_file_path.strip(): - error_msg = "GLB文件路径为空" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("",) - - # 处理文件路径 - input_path = self._resolve_file_path(glb_file_path.strip()) - processing_log.append(f"输入文件路径: {input_path}") - - if not os.path.exists(input_path): - error_msg = f"GLB文件不存在: {input_path}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("",) - - try: - # 加载GLB文件 - processing_log.append("正在加载GLB文件...") - scene = trimesh.load(input_path) - - # 提取点云数据 - point_clouds = [] - other_geometries = [] - - if isinstance(scene, trimesh.Scene): - for name, geometry in scene.geometry.items(): - if isinstance(geometry, trimesh.PointCloud): - point_clouds.append((name, geometry)) - processing_log.append(f"发现点云: {name}, 点数: {len(geometry.vertices)}") - else: - other_geometries.append((name, geometry)) - processing_log.append(f"发现其他几何体: {name}, 类型: {type(geometry).__name__}") - elif isinstance(scene, trimesh.PointCloud): - point_clouds.append(("main_pointcloud", scene)) - processing_log.append(f"发现点云: main_pointcloud, 点数: {len(scene.vertices)}") - else: - # 尝试转换为点云 - if hasattr(scene, 'vertices'): - pc = trimesh.PointCloud(vertices=scene.vertices, colors=getattr(scene.visual, 'vertex_colors', None)) - point_clouds.append(("converted_pointcloud", pc)) - processing_log.append(f"转换为点云: converted_pointcloud, 点数: {len(pc.vertices)}") - else: - error_msg = "GLB文件中未找到点云数据" - processing_log.append(f"错误: {error_msg}") - return ("",) - - if not point_clouds: - error_msg = "GLB文件中没有点云数据" - processing_log.append(f"错误: {error_msg}") - return ("",) - - # 处理每个点云 - processed_point_clouds = [] - total_original_points = 0 - total_removed_points = 0 - - for name, point_cloud in point_clouds: - processing_log.append(f"\n处理点云: {name}") - original_count = len(point_cloud.vertices) - total_original_points += original_count - - # 获取原始点云的所有属性和状态 - original_vertices = point_cloud.vertices.copy() - colors = None - - # 获取颜色信息 - if hasattr(point_cloud.visual, 'vertex_colors') and point_cloud.visual.vertex_colors is not None: - colors = point_cloud.visual.vertex_colors.copy() - elif hasattr(point_cloud, 'colors') and point_cloud.colors is not None: - colors = point_cloud.colors.copy() - - # 初始化掩码（所有点都保留） - keep_mask = np.ones(len(original_vertices), dtype=bool) - - # 删除暗色点（仅基于颜色，不改变顶点坐标） - if colors is not None: - removed_count = self._remove_dark_points(original_vertices, colors, keep_mask, black_threshold, processing_log) - processing_log.append(f"删除黑色点: {removed_count} 个") - else: - processing_log.append("跳过黑色点过滤: 点云无颜色信息") - - # 应用掩码，保留原始的顶点坐标 - filtered_vertices = original_vertices[keep_mask] - filtered_colors = colors[keep_mask] if colors is not None else None - - remaining_count = len(filtered_vertices) - removed_count = original_count - remaining_count - total_removed_points += removed_count - - processing_log.append(f"点云 {name}: 原始 {original_count} -> 剩余 {remaining_count} (删除 {removed_count})") - - # 创建处理后的点云，保留所有原始属性 - if remaining_count > 0: - # 使用原始点云作为模板，只更新顶点和颜色 - processed_pc = point_cloud.copy() - processed_pc.vertices = filtered_vertices - if filtered_colors is not None: - processed_pc.visual.vertex_colors = filtered_colors - - processed_point_clouds.append((name, processed_pc)) - processing_log.append(f"保留点云 {name} 的所有原始属性和状态") - else: - processing_log.append(f"警告: 点云 {name} 处理后没有剩余点") - - # 直接修改原始场景，保持所有几何信息不变 - if processed_point_clouds: - # 使用原始场景作为基础，直接替换点云数据 - new_scene = scene - - # 直接修改原始场景中的点云几何体，保持变换不变 - for name, processed_pc in processed_point_clouds: - if isinstance(scene, trimesh.Scene): - # 获取原始几何体 - if name in scene.geometry: - original_geometry = scene.geometry[name] - - # 如果是点云，直接修改顶点和颜色 - if isinstance(original_geometry, trimesh.PointCloud): - original_geometry.vertices = processed_pc.vertices - if hasattr(processed_pc.visual, 'vertex_colors') and processed_pc.visual.vertex_colors is not None: - original_geometry.visual.vertex_colors = processed_pc.visual.vertex_colors - - processing_log.append(f"直接修改原始点云 {name}，保持所有变换和属性不变") - else: - # 如果不是点云类型，替换整个几何体但保持变换 - scene.delete_geometry(name) - scene.add_geometry(processed_pc, node_name=name) - processing_log.append(f"替换几何体 {name}，尝试保持变换") - else: - # 新增几何体 - scene.add_geometry(processed_pc, node_name=name) - processing_log.append(f"添加新点云 {name}") - else: - # 单个几何体的情况 - new_scene = processed_pc - processing_log.append("单个几何体，直接使用处理后的点云") - - processing_log.append("直接修改原始场景，所有变换矩阵和几何属性完全保持不变") - - # 生成输出路径 - output_path = self._generate_output_path(output_filename) - processing_log.append(f"输出文件路径: {output_path}") - - # 保存处理后的GLB文件 - processing_log.append("正在保存处理后的GLB文件...") - new_scene.export(output_path) - - # 验证文件是否保存成功 - if os.path.exists(output_path): - file_size = os.path.getsize(output_path) - processing_log.append(f"GLB文件保存成功，文件大小: {file_size} bytes") - - # 验证输出文件的一致性 - try: - # 加载输出文件进行验证 - verification_scene = trimesh.load(output_path) - processing_log.append("文件验证: 输出文件可正常加载") - - # 检查是否保留了场景结构 - if isinstance(scene, trimesh.Scene) and isinstance(verification_scene, trimesh.Scene): - processing_log.append(f"场景结构验证: 原始 {len(scene.geometry)} 个对象 -> 输出 {len(verification_scene.geometry)} 个对象") - - except Exception as e: - processing_log.append(f"文件验证警告: {str(e)}") - - processing_log.append("GLB点云黑色点清理完成! 已保留原始模型的大小、朝向和几何属性") - return (output_path,) - else: - error_msg = "GLB文件保存失败" - processing_log.append(f"错误: {error_msg}") - return ("",) - else: - error_msg = "处理后没有剩余的点云数据" - processing_log.append(f"错误: {error_msg}") - return ("",) - - except Exception as e: - error_msg = f"处理GLB文件时发生错误: {str(e)}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - import traceback - traceback.print_exc() - return ("",) - - def _resolve_file_path(self, file_path: str) -> str: - """解析文件路径，支持绝对路径和相对路径""" - if os.path.isabs(file_path): - return file_path - - # 尝试相对于ComfyUI输出目录 - if FOLDER_PATHS_AVAILABLE: - output_dir = folder_paths.get_output_directory() - candidate_path = os.path.join(output_dir, file_path) - if os.path.exists(candidate_path): - return candidate_path - - # 尝试相对于当前工作目录 - if os.path.exists(file_path): - return os.path.abspath(file_path) - - # 返回原始路径（让后续检查处理错误） - return file_path - - def _generate_output_path(self, filename: str) -> str: - """生成输出文件路径""" - if FOLDER_PATHS_AVAILABLE: - output_dir = folder_paths.get_output_directory() - else: - output_dir = "output" - - # 确保输出目录存在 - os.makedirs(output_dir, exist_ok=True) - - # 处理文件名，确保是.glb扩展名 - if not filename.lower().endswith('.glb'): - filename = f"{filename}.glb" - - return os.path.join(output_dir, filename) - - def _remove_dark_points(self, vertices: np.ndarray, colors: np.ndarray, keep_mask: np.ndarray, - black_threshold: int, processing_log: List[str]) -> int: - """删除暗色点""" - removed_count = 0 - - # 确保颜色数据格式正确 - if colors.shape[1] == 4: # RGBA - rgb_colors = colors[:, :3] - else: # RGB - rgb_colors = colors - - # 如果颜色值在0-1范围内，转换为0-255 - if rgb_colors.max() <= 1.0: - rgb_colors = (rgb_colors * 255).astype(np.uint8) - else: - rgb_colors = rgb_colors.astype(np.uint8) - - # 计算RGB总和 - rgb_sum = rgb_colors.sum(axis=1) - - # 应用黑色阈值过滤 - if black_threshold > 0: - dark_mask = rgb_sum >= black_threshold - before_count = keep_mask.sum() - keep_mask &= dark_mask - after_count = keep_mask.sum() - removed_by_black = before_count - after_count - removed_count += removed_by_black - processing_log.append(f" 按RGB总和阈值({black_threshold})删除: {removed_by_black} 个点") - - return removed_count - - -class GLBPointCloudBounds: - """GLB点云包围盒计算器 - 计算包围盒并生成带可视化的预览文件（原模型数据保持不变）""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "glb_file_path": ("STRING", { - "default": "", - "tooltip": "GLB点云文件路径：输入需要计算包围盒的GLB格式点云文件。支持绝对路径或相对于ComfyUI output目录的相对路径。注意：此节点在原模型基础上添加预览信息，原模型数据保持不变" - }), - }, - "optional": { - "bounding_box_type": (["axis_aligned", "oriented"], { - "default": "axis_aligned", - "tooltip": "包围盒类型：axis_aligned=轴对齐包围盒(AABB)，计算快速；oriented=有向包围盒(OBB)，体积最小" - }), - "units": (["meters", "centimeters", "millimeters"], { - "default": "meters", - "tooltip": "输出单位：指定包围盒尺寸的输出单位" - }), - "add_bounding_box_visualization": ("BOOLEAN", { - "default": True, - "tooltip": "添加包围盒可视化：是否在输出的点云文件中添加红色包围盒线框点云，便于直接预览验证" - }), - "add_coordinate_axes": ("BOOLEAN", { - "default": False, - "tooltip": "添加坐标轴：是否添加RGB坐标轴(X=红色，Y=绿色，Z=蓝色)到输出点云中" - }), - "wireframe_density": ("INT", { - "default": 150, "min": 20, "max": 200, "step": 10, - "tooltip": "线框密度：每条包围盒边的点数，越高线框越清晰但点数越多。推荐50-100获得清晰效果" - }), - "enhance_visibility": ("BOOLEAN", { - "default": True, - "tooltip": "增强可见性：添加顶点高亮和面中心标记，让包围盒更容易识别" - }), - "output_filename": ("STRING", { - "default": "pointcloud_with_bounds", - "tooltip": "输出文件名：带包围盒可视化的点云文件名，系统会自动添加.glb扩展名" - }), - } - } - - RETURN_TYPES = ( - "STRING", # 输出的GLB文件路径 - "STRING", # JSON格式的scale数组 - ) - RETURN_NAMES = ( - "output_glb_path", - "scale_json", - ) - OUTPUT_TOOLTIPS = [ - "带预览的GLB文件路径 - 包含原始点云+包围盒可视化+坐标轴的完整点云文件", - "纯净的scale数组JSON: {\"scale\": [长, 宽, 高]} - 仅包含scale信息，可直接用于3D引擎缩放", - ] - OUTPUT_NODE = True - FUNCTION = "calculate_bounds_and_visualize" - CATEGORY = "💃VVL/Point Cloud Analysis" - - def calculate_bounds_and_visualize(self, - glb_file_path: str, - bounding_box_type: str = "axis_aligned", - units: str = "meters", - add_bounding_box_visualization: bool = True, - add_coordinate_axes: bool = True, - wireframe_density: int = 50, - enhance_visibility: bool = True, - output_filename: str = "pointcloud_with_bounds"): - """ - 计算GLB点云文件的包围盒尺寸并生成带可视化的点云文件 - """ - - processing_log = [] - processing_log.append("开始GLB点云包围盒计算和可视化生成...") - - # 检查依赖 - if not TRIMESH_AVAILABLE: - error_msg = "trimesh库不可用，无法处理GLB文件" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 验证输入文件路径 - if not glb_file_path or not glb_file_path.strip(): - error_msg = "GLB文件路径为空" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 处理文件路径 - input_path = self._resolve_file_path(glb_file_path.strip()) - self._current_input_path = input_path # 保存当前输入路径供后续使用 - processing_log.append(f"输入文件路径: {input_path}") - - if not os.path.exists(input_path): - error_msg = f"GLB文件不存在: {input_path}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("", "") - - try: - # 加载GLB文件 - processing_log.append("正在加载GLB文件...") - scene = trimesh.load(input_path) - - # 收集所有点云数据，考虑变换矩阵 - all_vertices = [] - point_cloud_count = 0 - - if isinstance(scene, trimesh.Scene): - for node_name in scene.graph.nodes: - if node_name in scene.geometry: - geometry = scene.geometry[node_name] - transform_matrix = scene.graph[node_name][0] # 获取变换矩阵 - - if hasattr(geometry, 'vertices') and geometry.vertices is not None: - # 应用变换矩阵获取真实世界坐标 - if transform_matrix is not None and not np.allclose(transform_matrix, np.eye(4)): - world_vertices = trimesh.transformations.transform_points(geometry.vertices, transform_matrix) - processing_log.append(f"发现几何体: {node_name}, 点数: {len(geometry.vertices)}, 已应用变换矩阵") - else: - world_vertices = geometry.vertices.copy() - processing_log.append(f"发现几何体: {node_name}, 点数: {len(geometry.vertices)}, 无变换") - - all_vertices.append(world_vertices) - point_cloud_count += 1 - elif isinstance(scene, trimesh.PointCloud): - all_vertices.append(scene.vertices) - point_cloud_count = 1 - processing_log.append(f"发现点云: main_pointcloud, 点数: {len(scene.vertices)}") - elif hasattr(scene, 'vertices') and scene.vertices is not None: - all_vertices.append(scene.vertices) - point_cloud_count = 1 - processing_log.append(f"发现几何体(作为点云处理): main_geometry, 点数: {len(scene.vertices)}") - else: - error_msg = "GLB文件中未找到点云或几何体数据" - processing_log.append(f"错误: {error_msg}") - return ("", "") - - if not all_vertices: - error_msg = "GLB文件中没有可用的顶点数据" - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 合并所有顶点（已经是世界坐标） - combined_vertices = np.vstack(all_vertices) - total_points = len(combined_vertices) - processing_log.append(f"合并了 {point_cloud_count} 个几何体，总点数: {total_points} (已应用变换矩阵)") - - # 计算包围盒 - processing_log.append(f"计算包围盒类型: {bounding_box_type}") - - if bounding_box_type == "axis_aligned": - # 轴对齐包围盒 (AABB) - min_point = np.min(combined_vertices, axis=0) - max_point = np.max(combined_vertices, axis=0) - extents = max_point - min_point - center = (min_point + max_point) / 2 - - processing_log.append(f"AABB计算完成:") - processing_log.append(f" 最小点: [{min_point[0]:.6f}, {min_point[1]:.6f}, {min_point[2]:.6f}]") - processing_log.append(f" 最大点: [{max_point[0]:.6f}, {max_point[1]:.6f}, {max_point[2]:.6f}]") - processing_log.append(f" 中心点: [{center[0]:.6f}, {center[1]:.6f}, {center[2]:.6f}]") - processing_log.append(f" 尺寸: [{extents[0]:.6f}, {extents[1]:.6f}, {extents[2]:.6f}]") - processing_log.append(f" 体积: {np.prod(extents):.6f}") - - else: # oriented - # 有向包围盒 (OBB) - try: - to_origin, obb_extents = trimesh.bounds.oriented_bounds(combined_vertices) - - # 计算OBB的中心点 - obb_center = -to_origin[:3, 3] # 变换矩阵的平移部分的负值 - center = obb_center # 设置center变量 - extents = obb_extents - - processing_log.append(f"OBB计算完成:") - processing_log.append(f" 中心点: [{center[0]:.6f}, {center[1]:.6f}, {center[2]:.6f}]") - processing_log.append(f" 尺寸: [{extents[0]:.6f}, {extents[1]:.6f}, {extents[2]:.6f}]") - processing_log.append(f" 体积: {np.prod(extents):.6f}") - - except Exception as e: - processing_log.append(f"OBB计算失败，回退到AABB: {str(e)}") - # 回退到AABB - min_point = np.min(combined_vertices, axis=0) - max_point = np.max(combined_vertices, axis=0) - extents = max_point - min_point - center = (min_point + max_point) / 2 - - - - # 应用单位转换 - unit_scale = {"meters": 1.0, "centimeters": 100.0, "millimeters": 1000.0} - scale_factor = unit_scale.get(units, 1.0) - - scaled_extents = extents * scale_factor - - # 生成scale JSON (按照 [长, 宽, 高] 的顺序，通常是 [X, Y, Z]) - scale_array = [float(scaled_extents[0]), float(scaled_extents[1]), float(scaled_extents[2])] - scale_json = { - "name": output_filename, - "scale": scale_array - } - - - - processing_log.append(f"") - processing_log.append(f"输出结果 (单位: {units}):") - processing_log.append(f" Scale数组: [{scale_array[0]:.6f}, {scale_array[1]:.6f}, {scale_array[2]:.6f}]") - processing_log.append(f" 长度(X): {scale_array[0]:.6f} {units}") - processing_log.append(f" 宽度(Y): {scale_array[1]:.6f} {units}") - processing_log.append(f" 高度(Z): {scale_array[2]:.6f} {units}") - - # 生成带可视化的点云文件（原模型数据不变，仅添加预览） - output_glb_path = "" - - try: - processing_log.append("正在生成带包围盒可视化的点云文件...") - output_glb_path = self._generate_visualization_pointcloud( - combined_vertices, extents, center, bounding_box_type, - add_bounding_box_visualization, add_coordinate_axes, - wireframe_density, enhance_visibility, output_filename, processing_log - ) - processing_log.append("可视化点云文件生成完成!") - except Exception as e: - processing_log.append(f"可视化点云生成失败: {str(e)}") - import traceback - traceback.print_exc() - - processing_log.append("GLB点云包围盒计算和可视化完成!") - - return ( - output_glb_path, - json.dumps(scale_json, indent=2), - ) - - except Exception as e: - error_msg = f"计算包围盒时发生错误: {str(e)}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - import traceback - traceback.print_exc() - return ("", "") - - def _resolve_file_path(self, file_path: str) -> str: - """解析文件路径，支持绝对路径和相对路径""" - if os.path.isabs(file_path): - return file_path - - # 尝试相对于ComfyUI输出目录 - if FOLDER_PATHS_AVAILABLE: - output_dir = folder_paths.get_output_directory() - candidate_path = os.path.join(output_dir, file_path) - if os.path.exists(candidate_path): - return candidate_path - - # 尝试相对于当前工作目录 - if os.path.exists(file_path): - return os.path.abspath(file_path) - - # 返回原始路径（让后续检查处理错误） - return file_path - - def _get_current_input_path(self): - """获取当前输入文件路径""" - return getattr(self, '_current_input_path', None) - - def _normalize_colors(self, colors): - """标准化颜色数组为RGBA格式，0-255范围""" - if colors is None: - return None - - colors = np.array(colors) - - # 如果是0-1范围，转换为0-255 - if colors.max() <= 1.0: - colors = (colors * 255).astype(np.uint8) - else: - colors = colors.astype(np.uint8) - - # 确保是RGBA格式 - if colors.shape[1] == 3: # RGB -> RGBA - alpha_channel = np.full((colors.shape[0], 1), 255, dtype=np.uint8) - colors = np.hstack([colors, alpha_channel]) - elif colors.shape[1] != 4: - # 如果不是3或4通道，创建白色RGBA - colors = np.tile([255, 255, 255, 255], (colors.shape[0], 1)).astype(np.uint8) - - return colors - - def _generate_visualization_pointcloud(self, combined_vertices, extents, center, bounds_type, - add_bounding_box_visualization, add_coordinate_axes, - wireframe_density, enhance_visibility, output_filename, processing_log): - """生成包含原始点云、包围盒线框和坐标轴的可视化点云文件""" - try: - # 直接使用已加载的原始场景数据，避免重新加载可能改变朝向 - input_path = self._get_current_input_path() - - # 提取原始点云数据（完全保持原始状态） - all_vertices = [] - all_colors = [] - original_point_count = 0 - - # 使用之前已经加载的combined_vertices作为基础，重新从原始场景提取完整数据 - original_scene = trimesh.load(input_path) - - # 创建一个新的场景，完全保持原始场景的结构和变换 - visualization_scene = trimesh.Scene() - - # 保持原始场景的完整结构 - if isinstance(original_scene, trimesh.Scene): - # 完全复制原始场景的所有几何体和变换 - for node_name in original_scene.graph.nodes: - if node_name in original_scene.geometry: - geometry = original_scene.geometry[node_name] - transform_matrix = original_scene.graph[node_name][0] # 获取变换矩阵 - - # 完全复制几何体，保持所有属性 - copied_geometry = geometry.copy() - - # 添加到新场景，保持原始的变换矩阵 - visualization_scene.add_geometry(copied_geometry, node_name=node_name, transform=transform_matrix) - - # 提取顶点用于统计（注意：这里用于显示统计，真实包围盒计算已在前面完成） - if hasattr(geometry, 'vertices') and geometry.vertices is not None: - # 这里不需要再次应用变换，因为包围盒计算已经在combined_vertices中处理过了 - all_vertices.append(geometry.vertices) - - # 保持原始颜色 - if isinstance(geometry, trimesh.PointCloud): - if hasattr(geometry.visual, 'vertex_colors') and geometry.visual.vertex_colors is not None: - colors = geometry.visual.vertex_colors.copy() - elif hasattr(geometry, 'colors') and geometry.colors is not None: - colors = geometry.colors.copy() - else: - colors = np.tile([255, 255, 255, 255], (len(geometry.vertices), 1)) - else: - if hasattr(geometry.visual, 'vertex_colors') and geometry.visual.vertex_colors is not None: - colors = geometry.visual.vertex_colors.copy() - else: - colors = np.tile([200, 200, 200, 255], (len(geometry.vertices), 1)) - - # 确保颜色格式一致 (RGBA, 0-255) - colors = self._normalize_colors(colors) - all_colors.append(colors) - original_point_count += len(geometry.vertices) - - processing_log.append(f"完全保持原始几何体: {node_name}, 变换矩阵和所有属性已保留") - - elif isinstance(original_scene, trimesh.PointCloud): - # 单个点云，直接复制 - visualization_scene = original_scene.copy() - all_vertices.append(original_scene.vertices) - - if hasattr(original_scene.visual, 'vertex_colors') and original_scene.visual.vertex_colors is not None: - colors = original_scene.visual.vertex_colors.copy() - elif hasattr(original_scene, 'colors') and original_scene.colors is not None: - colors = original_scene.colors.copy() - else: - colors = np.tile([255, 255, 255, 255], (len(original_scene.vertices), 1)) - - colors = self._normalize_colors(colors) - all_colors.append(colors) - original_point_count = len(original_scene.vertices) - processing_log.append(f"完全保持原始点云: {original_point_count:,} 个点") - - else: - # 单个几何体，转换为场景保持结构 - visualization_scene.add_geometry(original_scene.copy(), node_name="main_geometry") - if hasattr(original_scene, 'vertices') and original_scene.vertices is not None: - all_vertices.append(original_scene.vertices) - - if hasattr(original_scene.visual, 'vertex_colors') and original_scene.visual.vertex_colors is not None: - colors = original_scene.visual.vertex_colors.copy() - else: - colors = np.tile([200, 200, 200, 255], (len(original_scene.vertices), 1)) - - colors = self._normalize_colors(colors) - all_colors.append(colors) - original_point_count = len(original_scene.vertices) - processing_log.append(f"完全保持原始几何体: {original_point_count:,} 个点") - - processing_log.append(f"原始模型结构完全保留: {original_point_count:,} 个点，所有变换矩阵和朝向不变") - - # 添加包围盒线框点云到场景 - if add_bounding_box_visualization: - box_vertices, box_colors = self._create_bounding_box_pointcloud( - extents, center, bounds_type, wireframe_density, enhance_visibility, processing_log - ) - if len(box_vertices) > 0: - # 创建包围盒点云并添加到场景 - bounding_box_pointcloud = trimesh.PointCloud(vertices=box_vertices, colors=box_colors) - visualization_scene.add_geometry(bounding_box_pointcloud, node_name="bounding_box_visualization") - processing_log.append(f"包围盒线框: {len(box_vertices):,} 个点") - - # 添加坐标轴点云到场景 - if add_coordinate_axes: - axes_vertices, axes_colors = self._create_coordinate_axes_pointcloud( - extents, center, wireframe_density, processing_log - ) - if len(axes_vertices) > 0: - # 创建坐标轴点云并添加到场景 - coordinate_axes_pointcloud = trimesh.PointCloud(vertices=axes_vertices, colors=axes_colors) - visualization_scene.add_geometry(coordinate_axes_pointcloud, node_name="coordinate_axes") - processing_log.append(f"坐标轴: {len(axes_vertices):,} 个点") - - # 统计总点数 - total_visualization_points = original_point_count - box_point_count = 0 - axes_point_count = 0 - - if add_bounding_box_visualization: - box_point_count = len(box_vertices) if 'box_vertices' in locals() and len(box_vertices) > 0 else 0 - total_visualization_points += box_point_count - - if add_coordinate_axes: - axes_point_count = len(axes_vertices) if 'axes_vertices' in locals() and len(axes_vertices) > 0 else 0 - total_visualization_points += axes_point_count - - processing_log.append(f"总点数: {total_visualization_points:,} 个点") - - # 生成输出路径 - output_path = self._generate_output_path(output_filename) - - # 保存场景，完全保持原始模型的结构和朝向 - visualization_scene.export(output_path) - - processing_log.append(f"可视化点云已保存: {output_path}") - return output_path - - except Exception as e: - processing_log.append(f"可视化点云生成失败: {str(e)}") - raise e - - def _create_bounding_box_pointcloud(self, extents, center, bounds_type, wireframe_density, enhance_visibility, processing_log): - """创建包围盒线框的点云表示""" - try: - # 计算包围盒的8个顶点 - min_point = center - extents/2 - max_point = center + extents/2 - - box_vertices = np.array([ - [min_point[0], min_point[1], min_point[2]], # 0: min corner - [max_point[0], min_point[1], min_point[2]], # 1 - [max_point[0], max_point[1], min_point[2]], # 2 - [min_point[0], max_point[1], min_point[2]], # 3 - [min_point[0], min_point[1], max_point[2]], # 4 - [max_point[0], min_point[1], max_point[2]], # 5 - [max_point[0], max_point[1], max_point[2]], # 6: max corner - [min_point[0], max_point[1], max_point[2]], # 7 - ]) - - # 定义包围盒的12条边 - edges = [ - (0, 1), (1, 2), (2, 3), (3, 0), # 底面 - (4, 5), (5, 6), (6, 7), (7, 4), # 顶面 - (0, 4), (1, 5), (2, 6), (3, 7), # 垂直边 - ] - - wireframe_points = [] - - # 1. 生成线框边上的点（更密集） - for start_idx, end_idx in edges: - start_point = box_vertices[start_idx] - end_point = box_vertices[end_idx] - - # 在每条边上生成密集的点 - for i in range(wireframe_density): - t = i / (wireframe_density - 1) if wireframe_density > 1 else 0 - point = start_point + t * (end_point - start_point) - wireframe_points.append(point) - - # 2. 添加顶点高亮和面中心标记（可选增强可见性） - if enhance_visibility: - # 顶点高亮（每个顶点周围添加小点云） - vertex_highlight_radius = np.min(extents) * 0.01 # 顶点高亮半径 - highlight_density = max(5, wireframe_density // 10) # 顶点周围的点数 - - for vertex in box_vertices: - # 在每个顶点周围创建小的点云球 - for i in range(highlight_density): - # 生成随机方向 - theta = np.random.uniform(0, 2 * np.pi) - phi = np.random.uniform(0, np.pi) - r = np.random.uniform(0, vertex_highlight_radius) - - # 球坐标转换为笛卡尔坐标 - x = vertex[0] + r * np.sin(phi) * np.cos(theta) - y = vertex[1] + r * np.sin(phi) * np.sin(theta) - z = vertex[2] + r * np.cos(phi) - - wireframe_points.append([x, y, z]) - - # 面中心点标记 - face_centers = [ - # 底面中心 - (box_vertices[0] + box_vertices[1] + box_vertices[2] + box_vertices[3]) / 4, - # 顶面中心 - (box_vertices[4] + box_vertices[5] + box_vertices[6] + box_vertices[7]) / 4, - # 前面中心 - (box_vertices[0] + box_vertices[1] + box_vertices[4] + box_vertices[5]) / 4, - # 后面中心 - (box_vertices[2] + box_vertices[3] + box_vertices[6] + box_vertices[7]) / 4, - # 左面中心 - (box_vertices[0] + box_vertices[3] + box_vertices[4] + box_vertices[7]) / 4, - # 右面中心 - (box_vertices[1] + box_vertices[2] + box_vertices[5] + box_vertices[6]) / 4, - ] - - # 在每个面中心添加标记点 - face_mark_density = max(3, wireframe_density // 20) - face_mark_radius = np.min(extents) * 0.005 - - for face_center in face_centers: - for i in range(face_mark_density): - # 在面中心周围添加小的点云 - theta = np.random.uniform(0, 2 * np.pi) - phi = np.random.uniform(0, np.pi) - r = np.random.uniform(0, face_mark_radius) - - x = face_center[0] + r * np.sin(phi) * np.cos(theta) - y = face_center[1] + r * np.sin(phi) * np.sin(theta) - z = face_center[2] + r * np.cos(phi) - - wireframe_points.append([x, y, z]) - - wireframe_vertices = np.array(wireframe_points) - - # 为包围盒线框设置红色 - wireframe_colors = np.tile([255, 0, 0, 255], (len(wireframe_vertices), 1)) - - processing_log.append(f"包围盒线框生成: {len(wireframe_vertices):,} 个点 (密度={wireframe_density})") - - return wireframe_vertices, wireframe_colors - - except Exception as e: - processing_log.append(f"创建包围盒点云失败: {str(e)}") - return np.array([]), np.array([]) - - def _create_coordinate_axes_pointcloud(self, extents, origin_center, wireframe_density, processing_log): - """创建坐标轴的点云表示""" - try: - axis_length = np.max(extents) * 0.4 - axis_thickness = np.min(extents) * 0.005 # 坐标轴厚度 - axes_points = [] - axes_colors = [] - - # 计算坐标轴密度 - axis_line_density = wireframe_density - axis_thickness_points = max(3, wireframe_density // 15) # 厚度方向的点数 - - # X轴 - 红色（粗线条）- 从点云中心开始 - for i in range(axis_line_density): - t = i / (axis_line_density - 1) if axis_line_density > 1 else 0 - base_point = [origin_center[0] + t * axis_length, origin_center[1], origin_center[2]] - - # 主轴线 - axes_points.append(base_point) - axes_colors.append([255, 0, 0, 255]) - - # 增加厚度（在YZ平面上添加点） - for j in range(axis_thickness_points): - for k in range(axis_thickness_points): - offset_y = (j - axis_thickness_points//2) * axis_thickness / axis_thickness_points - offset_z = (k - axis_thickness_points//2) * axis_thickness / axis_thickness_points - thick_point = [base_point[0], base_point[1] + offset_y, base_point[2] + offset_z] - axes_points.append(thick_point) - axes_colors.append([255, 0, 0, 255]) - - # X轴箭头头部 - arrow_length = axis_length * 0.1 - arrow_base = axis_length * 0.9 - for i in range(axis_line_density // 2): - t = i / (axis_line_density // 2 - 1) if axis_line_density > 2 else 0 - # 箭头点 - arrow_x = origin_center[0] + arrow_base + t * arrow_length - arrow_offset = (1 - t) * axis_thickness * 2 - - axes_points.append([arrow_x, origin_center[1] + arrow_offset, origin_center[2]]) - axes_points.append([arrow_x, origin_center[1] - arrow_offset, origin_center[2]]) - axes_points.append([arrow_x, origin_center[1], origin_center[2] + arrow_offset]) - axes_points.append([arrow_x, origin_center[1], origin_center[2] - arrow_offset]) - axes_colors.extend([[255, 0, 0, 255]] * 4) - - # Y轴 - 绿色（粗线条）- 从点云中心开始 - for i in range(axis_line_density): - t = i / (axis_line_density - 1) if axis_line_density > 1 else 0 - base_point = [origin_center[0], origin_center[1] + t * axis_length, origin_center[2]] - - # 主轴线 - axes_points.append(base_point) - axes_colors.append([0, 255, 0, 255]) - - # 增加厚度（在XZ平面上添加点） - for j in range(axis_thickness_points): - for k in range(axis_thickness_points): - offset_x = (j - axis_thickness_points//2) * axis_thickness / axis_thickness_points - offset_z = (k - axis_thickness_points//2) * axis_thickness / axis_thickness_points - thick_point = [base_point[0] + offset_x, base_point[1], base_point[2] + offset_z] - axes_points.append(thick_point) - axes_colors.append([0, 255, 0, 255]) - - # Y轴箭头头部 - for i in range(axis_line_density // 2): - t = i / (axis_line_density // 2 - 1) if axis_line_density > 2 else 0 - arrow_y = origin_center[1] + arrow_base + t * arrow_length - arrow_offset = (1 - t) * axis_thickness * 2 - - axes_points.append([origin_center[0] + arrow_offset, arrow_y, origin_center[2]]) - axes_points.append([origin_center[0] - arrow_offset, arrow_y, origin_center[2]]) - axes_points.append([origin_center[0], arrow_y, origin_center[2] + arrow_offset]) - axes_points.append([origin_center[0], arrow_y, origin_center[2] - arrow_offset]) - axes_colors.extend([[0, 255, 0, 255]] * 4) - - # Z轴 - 蓝色（粗线条）- 从点云中心开始 - for i in range(axis_line_density): - t = i / (axis_line_density - 1) if axis_line_density > 1 else 0 - base_point = [origin_center[0], origin_center[1], origin_center[2] + t * axis_length] - - # 主轴线 - axes_points.append(base_point) - axes_colors.append([0, 0, 255, 255]) - - # 增加厚度（在XY平面上添加点） - for j in range(axis_thickness_points): - for k in range(axis_thickness_points): - offset_x = (j - axis_thickness_points//2) * axis_thickness / axis_thickness_points - offset_y = (k - axis_thickness_points//2) * axis_thickness / axis_thickness_points - thick_point = [base_point[0] + offset_x, base_point[1] + offset_y, base_point[2]] - axes_points.append(thick_point) - axes_colors.append([0, 0, 255, 255]) - - # Z轴箭头头部 - for i in range(axis_line_density // 2): - t = i / (axis_line_density // 2 - 1) if axis_line_density > 2 else 0 - arrow_z = origin_center[2] + arrow_base + t * arrow_length - arrow_offset = (1 - t) * axis_thickness * 2 - - axes_points.append([origin_center[0] + arrow_offset, origin_center[1], arrow_z]) - axes_points.append([origin_center[0] - arrow_offset, origin_center[1], arrow_z]) - axes_points.append([origin_center[0], origin_center[1] + arrow_offset, arrow_z]) - axes_points.append([origin_center[0], origin_center[1] - arrow_offset, arrow_z]) - axes_colors.extend([[0, 0, 255, 255]] * 4) - - axes_vertices = np.array(axes_points) - axes_colors_array = np.array(axes_colors) - - processing_log.append(f"坐标轴生成: {len(axes_vertices):,} 个点 (长度={axis_length:.3f}, 原点=[{origin_center[0]:.3f}, {origin_center[1]:.3f}, {origin_center[2]:.3f}])") - - return axes_vertices, axes_colors_array - - except Exception as e: - processing_log.append(f"创建坐标轴点云失败: {str(e)}") - return np.array([]), np.array([]) - - def _generate_output_path(self, output_filename): - """生成输出文件路径""" - # 确保文件名有正确的扩展名 - if not output_filename.lower().endswith('.glb'): - output_filename += '.glb' - - # 生成输出路径 - output_dir = folder_paths.get_output_directory() if FOLDER_PATHS_AVAILABLE else "output" - os.makedirs(output_dir, exist_ok=True) - - # 添加时间戳避免文件名冲突 - import time - timestamp = str(int(time.time())) - name_parts = output_filename.rsplit('.', 1) - if len(name_parts) == 2: - timestamped_filename = f"{name_parts[0]}_{timestamp}.{name_parts[1]}" - else: - timestamped_filename = f"{output_filename}_{timestamp}" - - return os.path.join(output_dir, timestamped_filename) - - -class GLBPointCloudOriginAdjuster: - """GLB点云原点调整计算器 - 计算原点调整信息并生成带坐标轴的预览文件（原模型数据保持不变）""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "glb_file_path": ("STRING", { - "default": "", - "tooltip": "GLB点云文件路径：输入需要调整原点的GLB格式点云文件。注意：此节点在原模型基础上添加坐标轴预览，原模型数据保持不变" - }), - }, - "optional": { - "origin_mode": (["center", "bottom_center"], { - "default": "bottom_center", - "tooltip": "原点模式：center=点云几何中心；bottom_center=点云底部中心(脚底)" - }), - "output_units": (["meters", "centimeters", "millimeters"], { - "default": "centimeters", - "tooltip": "输出单位：变换信息的输出单位" - }), - "add_coordinate_axes": ("BOOLEAN", { - "default": True, - "tooltip": "添加坐标轴：是否添加RGB坐标轴(X=红色，Y=绿色，Z=蓝色)到输出点云中，显示调整后的原点位置" - }), - "wireframe_density": ("INT", { - "default": 100, "min": 20, "max": 200, "step": 10, - "tooltip": "坐标轴密度：每条坐标轴的点数，越高坐标轴越清晰但点数越多。推荐80-120获得清晰效果" - }), - "output_filename": ("STRING", { - "default": "adjusted_pointcloud", - "tooltip": "输出文件名：调整后的点云文件名，系统会自动添加.glb扩展名" - }), - } - } - - RETURN_TYPES = ( - "STRING", # 处理后的GLB文件路径 - "STRING", # 变换信息JSON (position) - ) - RETURN_NAMES = ( - "adjusted_glb_path", - "transform_info", - ) - OUTPUT_TOOLTIPS = [ - "带预览的GLB文件路径 - 包含原始点云+坐标轴预览的完整点云文件", - "纯净的position信息JSON: {\"position\": [x, y, z]} - 仅包含位置信息，可直接用于UE等引擎", - ] - OUTPUT_NODE = True - FUNCTION = "adjust_origin" - CATEGORY = "💃VVL/Point Cloud Transform" - - def adjust_origin(self, - glb_file_path: str, - origin_mode: str = "bottom_center", - output_units: str = "centimeters", - add_coordinate_axes: bool = False, - wireframe_density: int = 100, - output_filename: str = "adjusted_pointcloud"): - """ - 调整GLB点云的原点位置和旋转 - """ - - processing_log = [] - processing_log.append("开始GLB点云原点调整...") - - # 检查依赖 - if not TRIMESH_AVAILABLE: - error_msg = "trimesh库不可用，无法处理GLB文件" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 验证输入文件路径 - if not glb_file_path or not glb_file_path.strip(): - error_msg = "GLB文件路径为空" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 处理文件路径 - input_path = self._resolve_file_path(glb_file_path.strip()) - processing_log.append(f"输入文件路径: {input_path}") - - if not os.path.exists(input_path): - error_msg = f"GLB文件不存在: {input_path}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("", "") - - try: - # 加载GLB文件 - processing_log.append("正在加载GLB文件...") - scene = trimesh.load(input_path) - - # 收集所有几何体和点云，考虑变换矩阵 - geometries_to_transform = [] - all_vertices = [] - total_points = 0 - - if isinstance(scene, trimesh.Scene): - for node_name in scene.graph.nodes: - if node_name in scene.geometry: - geometry = scene.geometry[node_name] - transform_matrix = scene.graph[node_name][0] # 获取变换矩阵 - geometries_to_transform.append((node_name, geometry)) - - if hasattr(geometry, 'vertices') and geometry.vertices is not None: - # 应用变换矩阵获取真实世界坐标 - if transform_matrix is not None and not np.allclose(transform_matrix, np.eye(4)): - world_vertices = trimesh.transformations.transform_points(geometry.vertices, transform_matrix) - processing_log.append(f"发现几何体: {node_name}, 顶点数: {len(geometry.vertices)}, 已应用变换矩阵") - else: - world_vertices = geometry.vertices.copy() - processing_log.append(f"发现几何体: {node_name}, 顶点数: {len(geometry.vertices)}, 无变换") - - all_vertices.append(world_vertices) - total_points += len(geometry.vertices) - else: - # 单个几何体 - geometries_to_transform.append(("main_geometry", scene)) - if hasattr(scene, 'vertices') and scene.vertices is not None: - all_vertices.append(scene.vertices) - total_points = len(scene.vertices) - processing_log.append(f"发现几何体: main_geometry, 顶点数: {total_points}") - - if not all_vertices: - error_msg = "GLB文件中没有可用的顶点数据" - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 合并所有顶点计算包围盒（已经是世界坐标） - combined_vertices = np.vstack(all_vertices) - processing_log.append(f"总顶点数: {total_points:,} (已应用变换矩阵)") - - # 计算原始包围盒 - original_min = np.min(combined_vertices, axis=0) - original_max = np.max(combined_vertices, axis=0) - original_center = (original_min + original_max) / 2 - original_size = original_max - original_min - - processing_log.append(f"原始包围盒:") - processing_log.append(f" 最小点: [{original_min[0]:.6f}, {original_min[1]:.6f}, {original_min[2]:.6f}]") - processing_log.append(f" 最大点: [{original_max[0]:.6f}, {original_max[1]:.6f}, {original_max[2]:.6f}]") - processing_log.append(f" 中心点: [{original_center[0]:.6f}, {original_center[1]:.6f}, {original_center[2]:.6f}]") - processing_log.append(f" 尺寸: [{original_size[0]:.6f}, {original_size[1]:.6f}, {original_size[2]:.6f}]") - - # 确定新的原点位置 - if origin_mode == "center": - new_origin = original_center.copy() - processing_log.append("原点模式: 几何中心") - else: # bottom_center - new_origin = np.array([original_center[0], original_center[1], original_min[2]]) - processing_log.append("原点模式: 底部中心(脚底)") - - # 计算平移向量 - translation = -new_origin - processing_log.append(f"平移向量: [{translation[0]:.6f}, {translation[1]:.6f}, {translation[2]:.6f}]") - - # 创建平移变换矩阵（仅平移，无旋转） - transform_matrix = np.eye(4) - transform_matrix[:3, 3] = translation - - # 应用变换到所有几何体 - new_scene = trimesh.Scene() - transformed_vertices_list = [] - - for name, geometry in geometries_to_transform: - # 复制几何体 - new_geometry = geometry.copy() - - # 应用变换 - new_geometry.apply_transform(transform_matrix) - - # 添加到新场景 - new_scene.add_geometry(new_geometry, node_name=name) - - # 收集变换后的顶点用于统计 - if hasattr(new_geometry, 'vertices') and new_geometry.vertices is not None: - transformed_vertices_list.append(new_geometry.vertices) - - processing_log.append(f"已变换几何体: {name}") - - # 添加坐标轴预览（在变换后的原点位置，即(0,0,0)） - if add_coordinate_axes: - try: - # 计算变换后的包围盒来确定坐标轴长度 - if transformed_vertices_list: - all_transformed_vertices = np.vstack(transformed_vertices_list) - transformed_extents = np.max(all_transformed_vertices, axis=0) - np.min(all_transformed_vertices, axis=0) - axis_length = np.max(transformed_extents) * 0.4 - else: - axis_length = 1.0 # 默认长度 - - # 在新的原点(0,0,0)处生成坐标轴 - axes_vertices, axes_colors = self._create_coordinate_axes_pointcloud_at_origin( - axis_length, wireframe_density, processing_log - ) - - if len(axes_vertices) > 0: - # 创建坐标轴点云 - axes_pointcloud = trimesh.PointCloud(vertices=axes_vertices, colors=axes_colors) - new_scene.add_geometry(axes_pointcloud, node_name="coordinate_axes") - processing_log.append(f"坐标轴预览: {len(axes_vertices):,} 个点 (显示新的原点位置)") - - except Exception as e: - processing_log.append(f"坐标轴生成失败: {str(e)}") - - # 计算变换后的包围盒 - if transformed_vertices_list: - transformed_vertices = np.vstack(transformed_vertices_list) - new_min = np.min(transformed_vertices, axis=0) - new_max = np.max(transformed_vertices, axis=0) - new_center = (new_min + new_max) / 2 - new_size = new_max - new_min - - processing_log.append(f"变换后包围盒:") - processing_log.append(f" 最小点: [{new_min[0]:.6f}, {new_min[1]:.6f}, {new_min[2]:.6f}]") - processing_log.append(f" 最大点: [{new_max[0]:.6f}, {new_max[1]:.6f}, {new_max[2]:.6f}]") - processing_log.append(f" 中心点: [{new_center[0]:.6f}, {new_center[1]:.6f}, {new_center[2]:.6f}]") - processing_log.append(f" 尺寸: [{new_size[0]:.6f}, {new_size[1]:.6f}, {new_size[2]:.6f}]") - - # 生成带预览的输出文件（原模型数据不变，仅添加坐标轴预览） - output_path = self._generate_output_path(output_filename) - processing_log.append(f"输出文件路径: {output_path}") - - # 创建包含原始点云和坐标轴的预览场景，完全保持原始结构 - original_scene = trimesh.load(input_path) # 重新加载原始场景 - - # 创建预览场景，完全复制原始场景的结构 - if isinstance(original_scene, trimesh.Scene): - preview_scene = trimesh.Scene() - # 完全复制原始场景的所有几何体和变换矩阵 - for node_name in original_scene.graph.nodes: - if node_name in original_scene.geometry: - geometry = original_scene.geometry[node_name] - transform_matrix = original_scene.graph[node_name][0] # 获取变换矩阵 - - # 完全复制几何体，保持所有属性 - copied_geometry = geometry.copy() - - # 添加到预览场景，保持原始的变换矩阵 - preview_scene.add_geometry(copied_geometry, node_name=node_name, transform=transform_matrix) - processing_log.append(f"完全保持原始几何体: {node_name}, 变换矩阵和朝向已保留") - else: - # 单个几何体的情况 - preview_scene = trimesh.Scene() - preview_scene.add_geometry(original_scene.copy(), node_name="main_geometry") - processing_log.append("完全保持原始几何体: main_geometry, 所有属性已保留") - - # 添加坐标轴预览（如果需要） - if add_coordinate_axes: - try: - # 使用原始包围盒来确定坐标轴长度 - axis_length = np.max(original_size) * 0.4 - - # 在计算出的新原点位置生成坐标轴 - axes_vertices, axes_colors = self._create_coordinate_axes_pointcloud_at_position( - new_origin, axis_length, wireframe_density, processing_log - ) - - if len(axes_vertices) > 0: - # 创建坐标轴点云并添加到场景 - axes_pointcloud = trimesh.PointCloud(vertices=axes_vertices, colors=axes_colors) - preview_scene.add_geometry(axes_pointcloud, node_name="coordinate_axes_preview") - processing_log.append(f"坐标轴预览: {len(axes_vertices):,} 个点 (显示建议的新原点位置)") - - except Exception as e: - processing_log.append(f"坐标轴生成失败: {str(e)}") - - # 保存预览文件 - processing_log.append("正在保存带预览的GLB文件...") - preview_scene.export(output_path) - - if os.path.exists(output_path): - file_size = os.path.getsize(output_path) - processing_log.append(f"预览GLB文件保存成功，文件大小: {file_size} bytes") - else: - error_msg = "预览GLB文件保存失败" - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 应用单位转换 - unit_scale = {"meters": 1.0, "centimeters": 100.0, "millimeters": 1000.0} - scale_factor = unit_scale.get(output_units, 1.0) - - # 生成变换信息（UE格式） - final_position = new_origin * scale_factor # 相对于原始坐标系的位置 - transform_info = { - "name": output_filename, - "position": [float(final_position[0]), float(final_position[1]), float(final_position[2])] - } - - processing_log.append("") - processing_log.append(f"变换信息 (单位: {output_units}):") - processing_log.append(f" Position: [{transform_info['position'][0]:.2f}, {transform_info['position'][1]:.2f}, {transform_info['position'][2]:.2f}]") - processing_log.append("原点调整预览生成完成!") - - return ( - output_path, - json.dumps(transform_info, indent=2), - ) - - except Exception as e: - error_msg = f"调整原点时发生错误: {str(e)}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - import traceback - traceback.print_exc() - return ("", "") - - def _resolve_file_path(self, file_path: str) -> str: - """解析文件路径，支持绝对路径和相对路径""" - if os.path.isabs(file_path): - return file_path - - # 尝试相对于ComfyUI输出目录 - if FOLDER_PATHS_AVAILABLE: - output_dir = folder_paths.get_output_directory() - candidate_path = os.path.join(output_dir, file_path) - if os.path.exists(candidate_path): - return candidate_path - - # 尝试相对于当前工作目录 - if os.path.exists(file_path): - return os.path.abspath(file_path) - - # 返回原始路径（让后续检查处理错误） - return file_path - - def _generate_output_path(self, filename: str) -> str: - """生成输出文件路径""" - # 确保文件名有正确的扩展名 - if not filename.lower().endswith('.glb'): - filename += '.glb' - - # 生成输出路径 - output_dir = folder_paths.get_output_directory() if FOLDER_PATHS_AVAILABLE else "output" - os.makedirs(output_dir, exist_ok=True) - - # 添加时间戳避免文件名冲突 - timestamp = str(int(time.time())) - name_parts = filename.rsplit('.', 1) - if len(name_parts) == 2: - timestamped_filename = f"{name_parts[0]}_{timestamp}.{name_parts[1]}" - else: - timestamped_filename = f"{filename}_{timestamp}" - - return os.path.join(output_dir, timestamped_filename) - - - - def _create_coordinate_axes_pointcloud_at_position(self, position, axis_length, wireframe_density, processing_log): - """在指定位置创建坐标轴的点云表示""" - try: - center = np.array(position) # 指定的位置 - axis_thickness = axis_length * 0.01 # 坐标轴厚度 - axes_points = [] - axes_colors = [] - - # 计算坐标轴密度 - axis_line_density = wireframe_density # 使用参数控制密度 - axis_thickness_points = max(3, axis_line_density // 15) # 厚度方向的点数 - - # X轴 - 红色（从指定位置开始） - for i in range(axis_line_density): - t = i / (axis_line_density - 1) if axis_line_density > 1 else 0 - base_point = [center[0] + t * axis_length, center[1], center[2]] - - # 主轴线 - axes_points.append(base_point) - axes_colors.append([255, 0, 0, 255]) - - # 增加厚度（在YZ平面上添加点） - for j in range(axis_thickness_points): - for k in range(axis_thickness_points): - offset_y = (j - axis_thickness_points//2) * axis_thickness / axis_thickness_points - offset_z = (k - axis_thickness_points//2) * axis_thickness / axis_thickness_points - thick_point = [base_point[0], base_point[1] + offset_y, base_point[2] + offset_z] - axes_points.append(thick_point) - axes_colors.append([255, 0, 0, 255]) - - # X轴箭头头部 - arrow_length = axis_length * 0.1 - arrow_base = axis_length * 0.9 - for i in range(axis_line_density // 2): - t = i / (axis_line_density // 2 - 1) if axis_line_density > 2 else 0 - arrow_x = center[0] + arrow_base + t * arrow_length - arrow_offset = (1 - t) * axis_thickness * 2 - - axes_points.append([arrow_x, center[1] + arrow_offset, center[2]]) - axes_points.append([arrow_x, center[1] - arrow_offset, center[2]]) - axes_points.append([arrow_x, center[1], center[2] + arrow_offset]) - axes_points.append([arrow_x, center[1], center[2] - arrow_offset]) - axes_colors.extend([[255, 0, 0, 255]] * 4) - - # Y轴 - 绿色（从指定位置开始） - for i in range(axis_line_density): - t = i / (axis_line_density - 1) if axis_line_density > 1 else 0 - base_point = [center[0], center[1] + t * axis_length, center[2]] - - # 主轴线 - axes_points.append(base_point) - axes_colors.append([0, 255, 0, 255]) - - # 增加厚度（在XZ平面上添加点） - for j in range(axis_thickness_points): - for k in range(axis_thickness_points): - offset_x = (j - axis_thickness_points//2) * axis_thickness / axis_thickness_points - offset_z = (k - axis_thickness_points//2) * axis_thickness / axis_thickness_points - thick_point = [base_point[0] + offset_x, base_point[1], base_point[2] + offset_z] - axes_points.append(thick_point) - axes_colors.append([0, 255, 0, 255]) - - # Y轴箭头头部 - for i in range(axis_line_density // 2): - t = i / (axis_line_density // 2 - 1) if axis_line_density > 2 else 0 - arrow_y = center[1] + arrow_base + t * arrow_length - arrow_offset = (1 - t) * axis_thickness * 2 - - axes_points.append([center[0] + arrow_offset, arrow_y, center[2]]) - axes_points.append([center[0] - arrow_offset, arrow_y, center[2]]) - axes_points.append([center[0], arrow_y, center[2] + arrow_offset]) - axes_points.append([center[0], arrow_y, center[2] - arrow_offset]) - axes_colors.extend([[0, 255, 0, 255]] * 4) - - # Z轴 - 蓝色（从指定位置开始） - for i in range(axis_line_density): - t = i / (axis_line_density - 1) if axis_line_density > 1 else 0 - base_point = [center[0], center[1], center[2] + t * axis_length] - - # 主轴线 - axes_points.append(base_point) - axes_colors.append([0, 0, 255, 255]) - - # 增加厚度（在XY平面上添加点） - for j in range(axis_thickness_points): - for k in range(axis_thickness_points): - offset_x = (j - axis_thickness_points//2) * axis_thickness / axis_thickness_points - offset_y = (k - axis_thickness_points//2) * axis_thickness / axis_thickness_points - thick_point = [base_point[0] + offset_x, base_point[1] + offset_y, base_point[2]] - axes_points.append(thick_point) - axes_colors.append([0, 0, 255, 255]) - - # Z轴箭头头部 - for i in range(axis_line_density // 2): - t = i / (axis_line_density // 2 - 1) if axis_line_density > 2 else 0 - arrow_z = center[2] + arrow_base + t * arrow_length - arrow_offset = (1 - t) * axis_thickness * 2 - - axes_points.append([center[0] + arrow_offset, center[1], arrow_z]) - axes_points.append([center[0] - arrow_offset, center[1], arrow_z]) - axes_points.append([center[0], center[1] + arrow_offset, arrow_z]) - axes_points.append([center[0], center[1] - arrow_offset, arrow_z]) - axes_colors.extend([[0, 0, 255, 255]] * 4) - - axes_vertices = np.array(axes_points) - axes_colors_array = np.array(axes_colors) - - processing_log.append(f"坐标轴生成: {len(axes_vertices):,} 个点 (长度={axis_length:.3f}, 密度={wireframe_density}, 位置=[{center[0]:.3f}, {center[1]:.3f}, {center[2]:.3f}])") - - return axes_vertices, axes_colors_array - - except Exception as e: - processing_log.append(f"创建坐标轴点云失败: {str(e)}") - return np.array([]), np.array([]) - - def _create_coordinate_axes_pointcloud_at_origin(self, axis_length, wireframe_density, processing_log): - """在原点(0,0,0)创建坐标轴的点云表示""" - try: - origin_center = np.array([0.0, 0.0, 0.0]) # 新的原点位置 - axis_thickness = axis_length * 0.01 # 坐标轴厚度 - axes_points = [] - axes_colors = [] - - # 计算坐标轴密度 - axis_line_density = wireframe_density # 使用参数控制密度 - axis_thickness_points = max(3, axis_line_density // 15) # 厚度方向的点数 - - # X轴 - 红色（从原点开始） - for i in range(axis_line_density): - t = i / (axis_line_density - 1) if axis_line_density > 1 else 0 - base_point = [origin_center[0] + t * axis_length, origin_center[1], origin_center[2]] - - # 主轴线 - axes_points.append(base_point) - axes_colors.append([255, 0, 0, 255]) - - # 增加厚度（在YZ平面上添加点） - for j in range(axis_thickness_points): - for k in range(axis_thickness_points): - offset_y = (j - axis_thickness_points//2) * axis_thickness / axis_thickness_points - offset_z = (k - axis_thickness_points//2) * axis_thickness / axis_thickness_points - thick_point = [base_point[0], base_point[1] + offset_y, base_point[2] + offset_z] - axes_points.append(thick_point) - axes_colors.append([255, 0, 0, 255]) - - # X轴箭头头部 - arrow_length = axis_length * 0.1 - arrow_base = axis_length * 0.9 - for i in range(axis_line_density // 2): - t = i / (axis_line_density // 2 - 1) if axis_line_density > 2 else 0 - arrow_x = origin_center[0] + arrow_base + t * arrow_length - arrow_offset = (1 - t) * axis_thickness * 2 - - axes_points.append([arrow_x, origin_center[1] + arrow_offset, origin_center[2]]) - axes_points.append([arrow_x, origin_center[1] - arrow_offset, origin_center[2]]) - axes_points.append([arrow_x, origin_center[1], origin_center[2] + arrow_offset]) - axes_points.append([arrow_x, origin_center[1], origin_center[2] - arrow_offset]) - axes_colors.extend([[255, 0, 0, 255]] * 4) - - # Y轴 - 绿色（从原点开始） - for i in range(axis_line_density): - t = i / (axis_line_density - 1) if axis_line_density > 1 else 0 - base_point = [origin_center[0], origin_center[1] + t * axis_length, origin_center[2]] - - # 主轴线 - axes_points.append(base_point) - axes_colors.append([0, 255, 0, 255]) - - # 增加厚度（在XZ平面上添加点） - for j in range(axis_thickness_points): - for k in range(axis_thickness_points): - offset_x = (j - axis_thickness_points//2) * axis_thickness / axis_thickness_points - offset_z = (k - axis_thickness_points//2) * axis_thickness / axis_thickness_points - thick_point = [base_point[0] + offset_x, base_point[1], base_point[2] + offset_z] - axes_points.append(thick_point) - axes_colors.append([0, 255, 0, 255]) - - # Y轴箭头头部 - for i in range(axis_line_density // 2): - t = i / (axis_line_density // 2 - 1) if axis_line_density > 2 else 0 - arrow_y = origin_center[1] + arrow_base + t * arrow_length - arrow_offset = (1 - t) * axis_thickness * 2 - - axes_points.append([origin_center[0] + arrow_offset, arrow_y, origin_center[2]]) - axes_points.append([origin_center[0] - arrow_offset, arrow_y, origin_center[2]]) - axes_points.append([origin_center[0], arrow_y, origin_center[2] + arrow_offset]) - axes_points.append([origin_center[0], arrow_y, origin_center[2] - arrow_offset]) - axes_colors.extend([[0, 255, 0, 255]] * 4) - - # Z轴 - 蓝色（从原点开始） - for i in range(axis_line_density): - t = i / (axis_line_density - 1) if axis_line_density > 1 else 0 - base_point = [origin_center[0], origin_center[1], origin_center[2] + t * axis_length] - - # 主轴线 - axes_points.append(base_point) - axes_colors.append([0, 0, 255, 255]) - - # 增加厚度（在XY平面上添加点） - for j in range(axis_thickness_points): - for k in range(axis_thickness_points): - offset_x = (j - axis_thickness_points//2) * axis_thickness / axis_thickness_points - offset_y = (k - axis_thickness_points//2) * axis_thickness / axis_thickness_points - thick_point = [base_point[0] + offset_x, base_point[1] + offset_y, base_point[2]] - axes_points.append(thick_point) - axes_colors.append([0, 0, 255, 255]) - - # Z轴箭头头部 - for i in range(axis_line_density // 2): - t = i / (axis_line_density // 2 - 1) if axis_line_density > 2 else 0 - arrow_z = origin_center[2] + arrow_base + t * arrow_length - arrow_offset = (1 - t) * axis_thickness * 2 - - axes_points.append([origin_center[0] + arrow_offset, origin_center[1], arrow_z]) - axes_points.append([origin_center[0] - arrow_offset, origin_center[1], arrow_z]) - axes_points.append([origin_center[0], origin_center[1] + arrow_offset, arrow_z]) - axes_points.append([origin_center[0], origin_center[1] - arrow_offset, arrow_z]) - axes_colors.extend([[0, 0, 255, 255]] * 4) - - axes_vertices = np.array(axes_points) - axes_colors_array = np.array(axes_colors) - - processing_log.append(f"坐标轴生成: {len(axes_vertices):,} 个点 (长度={axis_length:.3f}, 密度={wireframe_density}, 原点=[0.000, 0.000, 0.000])") - - return axes_vertices, axes_colors_array - - except Exception as e: - processing_log.append(f"创建坐标轴点云失败: {str(e)}") - return np.array([]), np.array([]) - - -class GLBPointCloudDensityFilter: - """GLB点云密度过滤器 - 根据局部密度删除稀疏区域，保留密度最高的核心区域""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "glb_file_path": ("STRING", { - "default": "", - "tooltip": "GLB点云文件路径：输入需要进行密度过滤的GLB格式点云文件。支持绝对路径或相对于ComfyUI output目录的相对路径" - }), - }, - "optional": { - "density_threshold": ("FLOAT", { - "default": 0.3, "min": 0.01, "max": 0.99, "step": 0.01, - "tooltip": "密度阈值：保留密度高于此比例的点。0.3=保留密度前70%的点；0.5=保留密度前50%的点。值越大删除越多稀疏点" - }), - "neighborhood_radius": ("FLOAT", { - "default": 0.1, "min": 0.01, "max": 1.0, "step": 0.01, - "tooltip": "邻域半径：计算密度时的搜索半径。较小值=精细密度检测；较大值=平滑密度检测。建议0.05-0.2" - }), - "min_neighbors": ("INT", { - "default": 5, "min": 1, "max": 50, "step": 1, - "tooltip": "最小邻居数：一个点在邻域内至少需要的邻居点数量才不被视为稀疏点。值越大要求密度越高" - }), - "preserve_core_percentage": ("FLOAT", { - "default": 0.8, "min": 0.1, "max": 1.0, "step": 0.05, - "tooltip": "核心区域保留比例：无论密度如何，都会保留密度最高的这部分点。0.8=保证保留最密集的80%区域" - }), - "use_adaptive_radius": ("BOOLEAN", { - "default": True, - "tooltip": "自适应半径：根据点云整体尺寸自动调整邻域半径，确保在不同尺寸的点云上都有良好效果" - }), - "output_filename": ("STRING", { - "default": "density_filtered_pointcloud", - "tooltip": "输出GLB文件名：密度过滤后点云的保存文件名，系统会自动添加.glb扩展名" - }), - } - } - - RETURN_TYPES = ( - "STRING", # 处理后的GLB文件路径 - "STRING", # 过滤统计信息JSON - ) - RETURN_NAMES = ( - "filtered_glb_path", - "filter_stats", - ) - OUTPUT_TOOLTIPS = [ - "密度过滤后的GLB文件完整路径", - "过滤统计信息JSON：包含原始点数、保留点数、删除点数等统计数据", - ] - OUTPUT_NODE = True - FUNCTION = "filter_by_density" - CATEGORY = "💃VVL/Point Cloud Filtering" - - def filter_by_density(self, - glb_file_path: str, - density_threshold: float = 0.3, - neighborhood_radius: float = 0.1, - min_neighbors: int = 5, - preserve_core_percentage: float = 0.8, - use_adaptive_radius: bool = True, - output_filename: str = "density_filtered_pointcloud"): - """ - 根据局部密度过滤GLB点云文件，删除稀疏区域 - """ - - processing_log = [] - processing_log.append("开始GLB点云密度过滤...") - - # 检查依赖 - if not TRIMESH_AVAILABLE: - error_msg = "trimesh库不可用，无法处理GLB文件" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 验证输入文件路径 - if not glb_file_path or not glb_file_path.strip(): - error_msg = "GLB文件路径为空" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 处理文件路径 - input_path = self._resolve_file_path(glb_file_path.strip()) - processing_log.append(f"输入文件路径: {input_path}") - - if not os.path.exists(input_path): - error_msg = f"GLB文件不存在: {input_path}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - return ("", "") - - try: - # 加载GLB文件 - processing_log.append("正在加载GLB文件...") - scene = trimesh.load(input_path) - - # 提取点云数据 - point_clouds = [] - other_geometries = [] - - if isinstance(scene, trimesh.Scene): - for name, geometry in scene.geometry.items(): - if isinstance(geometry, trimesh.PointCloud): - point_clouds.append((name, geometry)) - processing_log.append(f"发现点云: {name}, 点数: {len(geometry.vertices)}") - else: - other_geometries.append((name, geometry)) - elif isinstance(scene, trimesh.PointCloud): - point_clouds.append(("main_pointcloud", scene)) - processing_log.append(f"发现点云: main_pointcloud, 点数: {len(scene.vertices)}") - else: - # 尝试转换为点云 - if hasattr(scene, 'vertices'): - pc = trimesh.PointCloud(vertices=scene.vertices, colors=getattr(scene.visual, 'vertex_colors', None)) - point_clouds.append(("converted_pointcloud", pc)) - processing_log.append(f"转换为点云: converted_pointcloud, 点数: {len(pc.vertices)}") - else: - error_msg = "GLB文件中未找到点云数据" - processing_log.append(f"错误: {error_msg}") - return ("", "") - - if not point_clouds: - error_msg = "GLB文件中没有点云数据" - processing_log.append(f"错误: {error_msg}") - return ("", "") - - # 处理每个点云 - filtered_point_clouds = [] - total_original_points = 0 - total_filtered_points = 0 - - for name, point_cloud in point_clouds: - processing_log.append(f"\n处理点云: {name}") - original_count = len(point_cloud.vertices) - total_original_points += original_count - - # 获取原始点云的所有属性和状态 - original_vertices = point_cloud.vertices.copy() - colors = None - - # 获取颜色信息 - if hasattr(point_cloud.visual, 'vertex_colors') and point_cloud.visual.vertex_colors is not None: - colors = point_cloud.visual.vertex_colors.copy() - elif hasattr(point_cloud, 'colors') and point_cloud.colors is not None: - colors = point_cloud.colors.copy() - - # 应用密度过滤（仅基于密度，不改变顶点坐标） - keep_mask, filter_info = self._apply_density_filter( - original_vertices, colors, density_threshold, neighborhood_radius, min_neighbors, - preserve_core_percentage, use_adaptive_radius, processing_log - ) - - # 应用掩码，保留原始的顶点坐标 - filtered_vertices = original_vertices[keep_mask] - filtered_colors = colors[keep_mask] if colors is not None else None - - filtered_count = len(filtered_vertices) - total_filtered_points += filtered_count - - processing_log.append(f"点云 {name}: 原始 {original_count} -> 保留 {filtered_count} (删除 {original_count - filtered_count})") - - # 创建过滤后的点云，保留所有原始属性 - if filtered_count > 0: - # 使用原始点云作为模板，只更新顶点和颜色 - filtered_pc = point_cloud.copy() - filtered_pc.vertices = filtered_vertices - if filtered_colors is not None: - filtered_pc.visual.vertex_colors = filtered_colors - - filtered_point_clouds.append((name, filtered_pc)) - processing_log.append(f"保留点云 {name} 的所有原始属性和状态") - else: - processing_log.append(f"警告: 点云 {name} 过滤后没有剩余点") - - # 直接修改原始场景，保持所有几何信息不变 - if filtered_point_clouds: - # 使用原始场景作为基础，直接替换点云数据 - new_scene = scene - - # 直接修改原始场景中的点云几何体，保持变换不变 - for name, filtered_pc in filtered_point_clouds: - if isinstance(scene, trimesh.Scene): - # 获取原始几何体 - if name in scene.geometry: - original_geometry = scene.geometry[name] - - # 如果是点云，直接修改顶点和颜色 - if isinstance(original_geometry, trimesh.PointCloud): - original_geometry.vertices = filtered_pc.vertices - if hasattr(filtered_pc.visual, 'vertex_colors') and filtered_pc.visual.vertex_colors is not None: - original_geometry.visual.vertex_colors = filtered_pc.visual.vertex_colors - - processing_log.append(f"直接修改原始点云 {name}，保持所有变换和属性不变") - else: - # 如果不是点云类型，替换整个几何体但保持变换 - scene.delete_geometry(name) - scene.add_geometry(filtered_pc, node_name=name) - processing_log.append(f"替换几何体 {name}，尝试保持变换") - else: - # 新增几何体 - scene.add_geometry(filtered_pc, node_name=name) - processing_log.append(f"添加新点云 {name}") - else: - # 单个几何体的情况 - new_scene = filtered_pc - processing_log.append("单个几何体，直接使用过滤后的点云") - - processing_log.append("直接修改原始场景，所有变换矩阵和几何属性完全保持不变") - - # 生成输出路径 - output_path = self._generate_output_path(output_filename) - processing_log.append(f"输出文件路径: {output_path}") - - # 保存过滤后的GLB文件 - processing_log.append("正在保存过滤后的GLB文件...") - new_scene.export(output_path) - - # 验证文件是否保存成功 - if os.path.exists(output_path): - file_size = os.path.getsize(output_path) - processing_log.append(f"GLB文件保存成功，文件大小: {file_size} bytes") - - # 生成统计信息 - filter_stats = { - "original_points": total_original_points, - "filtered_points": total_filtered_points, - "removed_points": total_original_points - total_filtered_points, - "retention_rate": total_filtered_points / total_original_points if total_original_points > 0 else 0, - "parameters": { - "density_threshold": density_threshold, - "neighborhood_radius": neighborhood_radius, - "min_neighbors": min_neighbors, - "preserve_core_percentage": preserve_core_percentage, - "use_adaptive_radius": use_adaptive_radius - } - } - - stats_json = json.dumps(filter_stats, indent=2) - - processing_log.append("GLB点云密度过滤完成! 已保留原始模型的大小、朝向和几何属性") - processing_log.append(f"统计: 原始{total_original_points}点 -> 保留{total_filtered_points}点 (保留率: {filter_stats['retention_rate']:.1%})") - - return (output_path, stats_json) - else: - error_msg = "GLB文件保存失败" - processing_log.append(f"错误: {error_msg}") - return ("", "") - else: - error_msg = "过滤后没有剩余的点云数据" - processing_log.append(f"错误: {error_msg}") - return ("", "") - - except Exception as e: - error_msg = f"处理GLB文件时发生错误: {str(e)}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - import traceback - traceback.print_exc() - return ("", "") - - def _apply_density_filter(self, vertices, colors, density_threshold, neighborhood_radius, - min_neighbors, preserve_core_percentage, use_adaptive_radius, processing_log): - """ - 应用密度过滤算法，返回保留点的掩码 - """ - - n_points = len(vertices) - processing_log.append(f"开始密度分析，总点数: {n_points}") - - # 自适应半径调整 - if use_adaptive_radius: - # 计算点云的包围盒对角线长度 - bbox_min = np.min(vertices, axis=0) - bbox_max = np.max(vertices, axis=0) - bbox_diagonal = np.linalg.norm(bbox_max - bbox_min) - - # 根据包围盒大小调整半径 - adaptive_radius = neighborhood_radius * bbox_diagonal / 10.0 # 可调整的比例因子 - processing_log.append(f"自适应半径: {neighborhood_radius} -> {adaptive_radius:.4f} (基于包围盒对角线 {bbox_diagonal:.4f})") - neighborhood_radius = adaptive_radius - - # 使用KDTree进行高效的邻域搜索 - try: - from scipy.spatial import cKDTree - kdtree = cKDTree(vertices) - processing_log.append("使用scipy.spatial.cKDTree进行邻域搜索") - except ImportError: - processing_log.append("scipy不可用，使用简化的距离计算") - kdtree = None - - # 计算每个点的局部密度 - densities = np.zeros(n_points) - - if kdtree is not None: - # 高效的KDTree搜索 - for i in range(n_points): - # 查找邻域内的点 - neighbor_indices = kdtree.query_ball_point(vertices[i], neighborhood_radius) - densities[i] = len(neighbor_indices) - 1 # 减去自己 - else: - # 简化的距离计算（适用于小点云） - for i in range(n_points): - distances = np.linalg.norm(vertices - vertices[i], axis=1) - neighbor_count = np.sum(distances <= neighborhood_radius) - 1 # 减去自己 - densities[i] = neighbor_count - - processing_log.append(f"密度统计: 最小={densities.min():.1f}, 最大={densities.max():.1f}, 平均={densities.mean():.1f}") - - # 应用多层过滤策略 - keep_mask = np.ones(n_points, dtype=bool) - - # 1. 基于最小邻居数的过滤 - min_neighbor_mask = densities >= min_neighbors - removed_by_min_neighbors = np.sum(~min_neighbor_mask) - keep_mask &= min_neighbor_mask - processing_log.append(f"最小邻居数过滤: 删除 {removed_by_min_neighbors} 个点") - - # 2. 基于密度阈值的过滤 - if density_threshold > 0 and np.sum(keep_mask) > 0: - valid_densities = densities[keep_mask] - density_percentile = np.percentile(valid_densities, (1 - density_threshold) * 100) - density_mask = densities >= density_percentile - - # 保留当前mask中的点，再应用密度过滤 - temp_mask = keep_mask.copy() - keep_mask[temp_mask] &= density_mask[temp_mask] - - removed_by_density = np.sum(temp_mask) - np.sum(keep_mask) - processing_log.append(f"密度阈值过滤: 删除 {removed_by_density} 个点 (阈值密度: {density_percentile:.1f})") - - # 3. 核心区域保护 - if preserve_core_percentage < 1.0 and np.sum(keep_mask) > 0: - # 计算需要保留的核心点数 - target_core_points = int(n_points * preserve_core_percentage) - current_points = np.sum(keep_mask) - - if current_points < target_core_points: - # 当前保留的点数少于核心要求，需要添加高密度点 - need_more_points = target_core_points - current_points - - # 找到密度最高的点 - excluded_indices = np.where(~keep_mask)[0] - if len(excluded_indices) > 0: - excluded_densities = densities[excluded_indices] - # 按密度排序，选择密度最高的点 - sorted_indices = excluded_indices[np.argsort(-excluded_densities)] - points_to_restore = min(need_more_points, len(sorted_indices)) - - keep_mask[sorted_indices[:points_to_restore]] = True - processing_log.append(f"核心区域保护: 恢复 {points_to_restore} 个高密度点") - - # 统计信息 - original_count = n_points - filtered_count = np.sum(keep_mask) - removed_count = original_count - filtered_count - - filter_info = { - "original_points": original_count, - "filtered_points": filtered_count, - "removed_points": removed_count, - "retention_rate": filtered_count / original_count if original_count > 0 else 0, - "density_stats": { - "min": float(densities.min()), - "max": float(densities.max()), - "mean": float(densities.mean()), - "std": float(densities.std()) - } - } - - processing_log.append(f"密度过滤完成: 保留 {filtered_count}/{original_count} 点 (保留率: {filter_info['retention_rate']:.1%})") - - return keep_mask, filter_info - - def _resolve_file_path(self, file_path: str) -> str: - """解析文件路径，支持绝对路径和相对路径""" - if os.path.isabs(file_path): - return file_path - - # 尝试相对于ComfyUI输出目录 - if FOLDER_PATHS_AVAILABLE: - output_dir = folder_paths.get_output_directory() - candidate_path = os.path.join(output_dir, file_path) - if os.path.exists(candidate_path): - return candidate_path - - # 尝试相对于当前工作目录 - if os.path.exists(file_path): - return os.path.abspath(file_path) - - # 返回原始路径（让后续检查处理错误） - return file_path - - def _generate_output_path(self, filename: str) -> str: - """生成输出文件路径""" - if FOLDER_PATHS_AVAILABLE: - output_dir = folder_paths.get_output_directory() - else: - output_dir = "output" - - # 确保输出目录存在 - os.makedirs(output_dir, exist_ok=True) - - # 处理文件名，确保是.glb扩展名 - if not filename.lower().endswith('.glb'): - filename = f"{filename}.glb" - - return os.path.join(output_dir, filename) - - -# ----------------------------------------------------------------------------- -# 节点注册 - 更新注册信息 -# ----------------------------------------------------------------------------- - -NODE_CLASS_MAPPINGS = { - "GLBPointCloudProcessor": GLBPointCloudProcessor, - "GLBPointCloudBounds": GLBPointCloudBounds, - "GLBPointCloudOriginAdjuster": GLBPointCloudOriginAdjuster, - "GLBPointCloudDensityFilter": GLBPointCloudDensityFilter, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "GLBPointCloudProcessor": "VVL GLB Point Cloud Processor", - "GLBPointCloudBounds": "VVL GLB Point Cloud Bounds Visualizer", - "GLBPointCloudOriginAdjuster": "VVL GLB Point Cloud Origin Adjuster", - "GLBPointCloudDensityFilter": "VVL GLB Point Cloud Density Filter", -} - -# # 添加节点信息，帮助ComfyUI更好地识别 -# WEB_DIRECTORY = "./web" -# __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/json_merger.py b/json_merger.py deleted file mode 100644 index e41b144..0000000 --- a/json_merger.py +++ /dev/null @@ -1,494 +0,0 @@ -import json -import logging -from typing import Any, Dict, Union - -# 配置日志 -logger = logging.getLogger('json_merger') - -class JSONMerger: - """JSON合并器 - 用于合并GLBPointCloudBounds和GLBPointCloudOriginAdjuster的输出数据""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "json_data_1": ("STRING", { - "tooltip": "第一个JSON数据：来自GLBPointCloudBounds或GLBPointCloudOriginAdjuster的输出。支持包含name、scale、position等字段的JSON格式数据" - }), - "json_data_2": ("STRING", { - "tooltip": "第二个JSON数据：来自GLBPointCloudBounds或GLBPointCloudOriginAdjuster的输出。支持包含name、scale、position等字段的JSON格式数据" - }), - }, - "optional": { - "json_data_3": ("STRING", { - "tooltip": "第三个JSON数据（可选）：额外的JSON数据，用于三路合并。留空则只合并前两个数据" - }), - "merge_mode": (["smart_merge", "array_concat"], { - "default": "smart_merge", - "tooltip": "合并模式：smart_merge=智能合并（相同name合并字段，不同name用逗号连接）；array_concat=数组连接（所有数据放入数组）" - }), - } - } - - RETURN_TYPES = ( - "STRING", # 合并后的JSON字符串 - ) - RETURN_NAMES = ( - "merged_json", - ) - OUTPUT_TOOLTIPS = [ - "合并后的JSON数据：根据name字段智能合并或连接多个JSON对象的结果", - ] - OUTPUT_NODE = True - FUNCTION = "merge_json_data" - CATEGORY = "💃VVL/JSON Processing" - - def merge_json_data(self, - json_data_1: str, - json_data_2: str, - json_data_3: str = "", - merge_mode: str = "smart_merge"): - """ - 合并多个JSON数据 - """ - - processing_log = [] - processing_log.append("开始JSON数据合并...") - - try: - # 解析输入的JSON数据 - json_objects = [] - - # 解析第一个JSON - if json_data_1.strip(): - try: - obj1 = json.loads(json_data_1.strip()) - json_objects.append(obj1) - processing_log.append(f"解析JSON数据1成功: {self._get_object_summary(obj1)}") - except json.JSONDecodeError as e: - error_msg = f"JSON数据1解析失败: {str(e)}" - processing_log.append(f"错误: {error_msg}") - return (json.dumps({"error": error_msg}, indent=2),) - - # 解析第二个JSON - if json_data_2.strip(): - try: - obj2 = json.loads(json_data_2.strip()) - json_objects.append(obj2) - processing_log.append(f"解析JSON数据2成功: {self._get_object_summary(obj2)}") - except json.JSONDecodeError as e: - error_msg = f"JSON数据2解析失败: {str(e)}" - processing_log.append(f"错误: {error_msg}") - return (json.dumps({"error": error_msg}, indent=2),) - - # 解析第三个JSON（可选） - if json_data_3.strip(): - try: - obj3 = json.loads(json_data_3.strip()) - json_objects.append(obj3) - processing_log.append(f"解析JSON数据3成功: {self._get_object_summary(obj3)}") - except json.JSONDecodeError as e: - error_msg = f"JSON数据3解析失败: {str(e)}" - processing_log.append(f"错误: {error_msg}") - return (json.dumps({"error": error_msg}, indent=2),) - - if not json_objects: - error_msg = "没有有效的JSON数据输入" - processing_log.append(f"错误: {error_msg}") - return (json.dumps({"error": error_msg}, indent=2),) - - if len(json_objects) == 1: - processing_log.append("只有一个有效JSON对象，直接返回") - result = json_objects[0] - else: - # 根据合并模式处理 - if merge_mode == "smart_merge": - result = self._smart_merge(json_objects, processing_log) - else: # array_concat - result = self._array_concat(json_objects, processing_log) - - # 生成最终的JSON字符串 - merged_json = json.dumps(result, indent=2, ensure_ascii=False) - - processing_log.append("JSON数据合并完成!") - processing_log.append(f"最终结果: {self._get_object_summary(result)}") - - return (merged_json,) - - except Exception as e: - error_msg = f"合并JSON数据时发生错误: {str(e)}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - import traceback - traceback.print_exc() - return (json.dumps({"error": error_msg}, indent=2),) - - def _smart_merge(self, json_objects: list, processing_log: list) -> Union[Dict, list]: - """ - 智能合并：相同name的对象合并字段，不同name的对象用数组形式返回 - """ - processing_log.append("使用智能合并模式...") - - # 按name字段分组 - name_groups = {} - unnamed_objects = [] - - for i, obj in enumerate(json_objects): - if isinstance(obj, dict): - name = obj.get("name") - if name: - if name not in name_groups: - name_groups[name] = [] - name_groups[name].append(obj) - processing_log.append(f"对象{i+1}归入组 '{name}'") - else: - unnamed_objects.append(obj) - processing_log.append(f"对象{i+1}没有name字段，归入未命名组") - else: - unnamed_objects.append(obj) - processing_log.append(f"对象{i+1}不是字典类型，归入未命名组") - - # 合并结果 - merged_results = [] - - # 处理有name的组 - for name, group_objects in name_groups.items(): - if len(group_objects) == 1: - # 只有一个对象，直接使用 - merged_results.append(group_objects[0]) - processing_log.append(f"组 '{name}' 只有一个对象，直接使用") - else: - # 多个对象，合并字段 - merged_obj = {"name": name} - for obj in group_objects: - for key, value in obj.items(): - if key != "name": # 跳过name字段 - if key not in merged_obj: - merged_obj[key] = value - else: - # 如果字段已存在，优先使用非空值 - if not merged_obj[key] and value: - merged_obj[key] = value - - merged_results.append(merged_obj) - processing_log.append(f"组 '{name}' 合并了 {len(group_objects)} 个对象") - - # 添加未命名对象 - merged_results.extend(unnamed_objects) - if unnamed_objects: - processing_log.append(f"添加了 {len(unnamed_objects)} 个未命名对象") - - # 返回结果 - if len(merged_results) == 1: - processing_log.append("合并结果为单个对象") - return merged_results[0] - else: - processing_log.append(f"合并结果为 {len(merged_results)} 个对象的数组") - return merged_results - - def _array_concat(self, json_objects: list, processing_log: list) -> list: - """ - 数组连接：将所有对象放入一个数组 - """ - processing_log.append("使用数组连接模式...") - processing_log.append(f"将 {len(json_objects)} 个对象连接为数组") - return json_objects - - def _get_object_summary(self, obj: Any) -> str: - """ - 获取对象的简要描述 - """ - if isinstance(obj, dict): - keys = list(obj.keys()) - name = obj.get("name", "未命名") - return f"字典对象 name='{name}', 字段=[{', '.join(keys)}]" - elif isinstance(obj, list): - return f"数组对象，包含 {len(obj)} 个元素" - else: - return f"其他类型: {type(obj).__name__}" - -class UESceneGenerator: - """UE场景生成器 - 将JSONMerger的输出转换为UE场景配置格式""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "merged_json": ("*", { - "tooltip": "合并后的JSON数据：来自JSONMerger的输出，包含name、position、scale等字段的JSON数据，支持STRING/LIST/DICT类型" - }), - }, - "optional": { - "camera_position": ("STRING", { - "default": "[-1000, 0, 100]", - "tooltip": "相机位置：UE场景中相机的世界坐标位置 [X, Y, Z]，默认为[-1000, 0, 100]" - }), - "camera_rotation": ("STRING", { - "default": "[-12, 0, 0]", - "tooltip": "相机旋转：UE场景中相机的旋转角度 [Pitch, Yaw, Roll]，默认为[-12, 0, 0]" - }), - "camera_fov": ("FLOAT", { - "default": 58.53, "min": 1.0, "max": 179.0, "step": 0.1, - "tooltip": "相机视野角度：UE场景中相机的视野角度(FOV)，默认为58.53度。注意：这是一个数值参数，不要连接数组输入！" - }), - "default_object_rotation": ("STRING", { - "default": "[0, 0, 0]", - "tooltip": "默认物体旋转：当输入数据中没有rotation字段时使用的默认旋转值 [Pitch, Yaw, Roll]" - }), - "scale_multiplier": ("FLOAT", { - "default": 1.0, "min": 0.001, "max": 1000.0, "step": 0.1, - "tooltip": "缩放倍数：对所有scale值进行统一缩放的倍数，用于调整物体大小到合适的UE单位。注意：这是一个数值参数，不要连接数组输入！" - }), - } - } - - RETURN_TYPES = ( - "STRING", # UE场景配置JSON - ) - RETURN_NAMES = ( - "ue_scene_json", - ) - OUTPUT_TOOLTIPS = [ - "UE场景配置JSON：符合UE导入格式的完整场景配置，包含相机设置和物体列表", - ] - OUTPUT_NODE = True - FUNCTION = "generate_ue_scene" - CATEGORY = "💃VVL/UE Integration" - - def generate_ue_scene(self, - merged_json, - camera_position: str = "[-1000, 0, 100]", - camera_rotation: str = "[-12, 0, 0]", - camera_fov = 58.53, - default_object_rotation: str = "[0, 0, 0]", - scale_multiplier = 1.0): - """ - 将合并的JSON数据转换为UE场景配置格式 - """ - - processing_log = [] - processing_log.append("开始生成UE场景配置...") - - # 参数验证和修复 - try: - # 修复可能错误连接的参数 - if isinstance(camera_fov, str): - processing_log.append(f"警告: camera_fov参数类型错误(收到字符串'{camera_fov}')，使用默认值58.53") - camera_fov = 58.53 - elif not isinstance(camera_fov, (int, float)): - processing_log.append(f"警告: camera_fov参数类型错误(收到{type(camera_fov).__name__})，使用默认值58.53") - camera_fov = 58.53 - - if isinstance(scale_multiplier, str): - processing_log.append(f"警告: scale_multiplier参数类型错误(收到字符串'{scale_multiplier}')，使用默认值1.0") - scale_multiplier = 1.0 - elif not isinstance(scale_multiplier, (int, float)): - processing_log.append(f"警告: scale_multiplier参数类型错误(收到{type(scale_multiplier).__name__})，使用默认值1.0") - scale_multiplier = 1.0 - - # 确保参数在合理范围内 - camera_fov = max(1.0, min(179.0, float(camera_fov))) - scale_multiplier = max(0.001, min(1000.0, float(scale_multiplier))) - - processing_log.append(f"参数验证完成: camera_fov={camera_fov}, scale_multiplier={scale_multiplier}") - - except Exception as e: - processing_log.append(f"参数验证失败: {str(e)}，使用默认值") - camera_fov = 58.53 - scale_multiplier = 1.0 - - try: - # 处理输入的合并JSON数据（可能是字符串或直接的数据） - if isinstance(merged_json, str): - # 字符串输入，需要解析 - if not merged_json.strip(): - error_msg = "合并JSON数据为空" - processing_log.append(f"错误: {error_msg}") - return (json.dumps({"error": error_msg}, indent=2),) - - try: - merged_data = json.loads(merged_json.strip()) - processing_log.append("字符串JSON数据解析成功") - except json.JSONDecodeError as e: - error_msg = f"JSON数据解析失败: {str(e)}" - processing_log.append(f"错误: {error_msg}") - return (json.dumps({"error": error_msg}, indent=2),) - elif isinstance(merged_json, (list, dict)): - # 直接的数据结构输入 - merged_data = merged_json - processing_log.append(f"直接数据输入: {type(merged_json).__name__}") - - # 详细分析输入结构 - if isinstance(merged_json, list): - processing_log.append(f"列表包含 {len(merged_json)} 个元素:") - for i, item in enumerate(merged_json): - if isinstance(item, dict): - keys = list(item.keys()) - processing_log.append(f" 元素{i}: dict with keys {keys}") - elif isinstance(item, str): - preview = item[:100] + "..." if len(item) > 100 else item - processing_log.append(f" 元素{i}: string '{preview}'") - else: - processing_log.append(f" 元素{i}: {type(item).__name__}") - - else: - error_msg = f"不支持的输入数据类型: {type(merged_json).__name__}" - processing_log.append(f"错误: {error_msg}") - return (json.dumps({"error": error_msg}, indent=2),) - - # 解析相机参数 - try: - cam_pos = json.loads(camera_position) - cam_rot = json.loads(camera_rotation) - default_rot = json.loads(default_object_rotation) - - if not isinstance(cam_pos, list) or len(cam_pos) != 3: - raise ValueError("相机位置必须是包含3个数值的数组") - if not isinstance(cam_rot, list) or len(cam_rot) != 3: - raise ValueError("相机旋转必须是包含3个数值的数组") - if not isinstance(default_rot, list) or len(default_rot) != 3: - raise ValueError("默认旋转必须是包含3个数值的数组") - - processing_log.append(f"相机参数解析成功: 位置{cam_pos}, 旋转{cam_rot}, FOV{camera_fov}") - - except (json.JSONDecodeError, ValueError) as e: - error_msg = f"相机参数解析失败: {str(e)}" - processing_log.append(f"错误: {error_msg}") - return (json.dumps({"error": error_msg}, indent=2),) - - # 创建UE场景配置基础结构 - ue_scene = { - "camera": { - "position": cam_pos, - "rotation": cam_rot, - "fov": camera_fov - }, - "objects": [] - } - - # 处理物体数据 - objects_processed = 0 - - def process_data_item(data_item, item_name=""): - """递归处理数据项，支持嵌套结构""" - nonlocal objects_processed - - if isinstance(data_item, dict): - # 单个物体字典 - ue_object = self._convert_to_ue_object( - data_item, default_rot, scale_multiplier, processing_log - ) - if ue_object: - ue_scene["objects"].append(ue_object) - objects_processed += 1 - - elif isinstance(data_item, list): - # 数组，递归处理每个元素 - for i, sub_item in enumerate(data_item): - sub_name = f"{item_name}[{i}]" if item_name else f"item_{i}" - process_data_item(sub_item, sub_name) - - elif isinstance(data_item, str): - # 字符串，尝试解析为JSON - try: - parsed_item = json.loads(data_item) - processing_log.append(f"解析嵌套JSON字符串: {item_name}") - process_data_item(parsed_item, f"{item_name}_parsed") - except json.JSONDecodeError: - processing_log.append(f"跳过无法解析的字符串: {item_name} = '{data_item[:50]}...'") - - else: - processing_log.append(f"跳过不支持的数据类型 {item_name}: {type(data_item).__name__}") - - # 开始处理数据 - if isinstance(merged_data, (dict, list)): - processing_log.append(f"开始处理输入数据类型: {type(merged_data).__name__}") - process_data_item(merged_data, "root") - else: - error_msg = f"不支持的合并数据类型: {type(merged_data).__name__}" - processing_log.append(f"错误: {error_msg}") - return (json.dumps({"error": error_msg}, indent=2),) - - processing_log.append(f"成功处理 {objects_processed} 个物体") - - # 生成最终的UE场景JSON - ue_scene_json = json.dumps(ue_scene, indent=2, ensure_ascii=False) - - processing_log.append("UE场景配置生成完成!") - processing_log.append(f"场景包含: 1个相机 + {len(ue_scene['objects'])} 个物体") - - return (ue_scene_json,) - - except Exception as e: - error_msg = f"生成UE场景配置时发生错误: {str(e)}" - logger.error(error_msg) - processing_log.append(f"错误: {error_msg}") - import traceback - traceback.print_exc() - return (json.dumps({"error": error_msg}, indent=2),) - - def _convert_to_ue_object(self, obj_data: dict, default_rotation: list, - scale_multiplier: float, processing_log: list) -> dict: - """ - 将单个对象数据转换为UE物体格式 - """ - try: - # 获取物体名称 - name = obj_data.get("name", "未命名物体") - - # 获取位置信息 - position = obj_data.get("position", [0, 0, 0]) - if not isinstance(position, list) or len(position) != 3: - processing_log.append(f"物体 '{name}' 位置格式错误，使用默认值 [0, 0, 0]") - position = [0, 0, 0] - - # 获取旋转信息 - rotation = obj_data.get("rotation", default_rotation) - if not isinstance(rotation, list) or len(rotation) != 3: - processing_log.append(f"物体 '{name}' 旋转格式错误，使用默认值 {default_rotation}") - rotation = default_rotation.copy() - - # 获取缩放信息 - scale = obj_data.get("scale", [1, 1, 1]) - if not isinstance(scale, list) or len(scale) != 3: - processing_log.append(f"物体 '{name}' 缩放格式错误，使用默认值 [1, 1, 1]") - scale = [1, 1, 1] - - # 应用缩放倍数 - if scale_multiplier != 1.0: - scale = [s * scale_multiplier for s in scale] - processing_log.append(f"物体 '{name}' 应用缩放倍数 {scale_multiplier}") - - # 创建UE物体 - ue_object = { - "name": name, - "position": [float(position[0]), float(position[1]), float(position[2])], - "rotation": [float(rotation[0]), float(rotation[1]), float(rotation[2])], - "scale": [float(scale[0]), float(scale[1]), float(scale[2])] - } - - processing_log.append(f"转换物体 '{name}': 位置{position}, 旋转{rotation}, 缩放{scale}") - - return ue_object - - except Exception as e: - processing_log.append(f"转换物体失败: {str(e)}") - return None - - -# ----------------------------------------------------------------------------- -# 节点注册 -# ----------------------------------------------------------------------------- - -NODE_CLASS_MAPPINGS = { - "JSONMerger": JSONMerger, - "UESceneGenerator": UESceneGenerator, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "JSONMerger": "VVL JSON Merger", - "UESceneGenerator": "VVL UE Scene Generator", -} - -# 添加节点信息，帮助ComfyUI更好地识别 -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/mask_to_coordinates.py b/mask_to_coordinates.py deleted file mode 100644 index 1bf82ef..0000000 --- a/mask_to_coordinates.py +++ /dev/null @@ -1,328 +0,0 @@ -""" -Mask To Coordinates Node -从 mask 生成 coordinates_positive 和 coordinates_negative 坐标信息 -""" - -import torch -import numpy as np -import json -import cv2 -from typing import Tuple, List, Optional - - -class MaskToCoordinates: - """ - 从 mask 生成坐标信息的节点 - 白色区域 (值 >= threshold) 生成 coordinates_positive - 黑色区域 (值 < threshold) 生成 coordinates_negative - """ - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "mask": ("MASK", { - "tooltip": "输入的遮罩图像。白色区域将生成正样本坐标，黑色区域将生成负样本坐标。" - }), - "sample_method": (["random", "grid", "contour"], { - "default": "random", - "tooltip": "坐标采样方法：\n• random - 随机采样，在有效区域内随机选择点\n• grid - 网格采样，按网格模式均匀分布采样点\n• contour - 轮廓采样，沿着区域轮廓边缘采样" - }), - "positive_points": ("INT", { - "default": 10, - "min": 1, - "max": 1000, - "step": 1, - "tooltip": "从白色区域(值≥阈值)生成的正样本坐标点数量。这些坐标将作为 coordinates_positive 输出到 SAM2 分割节点。" - }), - "negative_points": ("INT", { - "default": 10, - "min": 0, - "max": 1000, - "step": 1, - "tooltip": "从黑色区域(值<阈值)生成的负样本坐标点数量。这些坐标将作为 coordinates_negative 输出到 SAM2 分割节点。设为0则不生成负样本点。" - }), - "threshold": ("FLOAT", { - "default": 0.5, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "二值化阈值，用于区分正负样本区域。像素值≥此阈值的区域被视为白色(正样本)，<此阈值的区域被视为黑色(负样本)。" - }), - }, - "optional": { - "min_distance": ("INT", { - "default": 10, - "min": 1, - "max": 100, - "step": 1, - "tooltip": "采样点之间的最小距离(像素)。用于避免点过于密集，确保采样点分布更均匀。仅在random和contour采样方法中有效。" - }), - "edge_margin": ("INT", { - "default": 5, - "min": 0, - "max": 50, - "step": 1, - "tooltip": "图像边缘留白像素数。在图像边缘指定像素范围内不会生成采样点，避免边界效应影响分割效果。设为0则不留白。" - }), - } - } - - RETURN_TYPES = ("STRING", "STRING", "IMAGE") - RETURN_NAMES = ("coordinates_positive", "coordinates_negative", "visualization") - FUNCTION = "mask_to_coordinates" - CATEGORY = "VVL" - DESCRIPTION = "从 mask 生成正负坐标信息，用于 SAM2 分割" - - def mask_to_coordinates( - self, - mask: torch.Tensor, - sample_method: str = "random", - positive_points: int = 10, - negative_points: int = 10, - threshold: float = 0.5, - min_distance: int = 10, - edge_margin: int = 5 - ) -> Tuple[str, str, torch.Tensor]: - """ - 从 mask 生成坐标信息 - - Args: - mask: 输入的 mask 张量 - sample_method: 采样方法 ("random", "grid", "contour") - positive_points: 正样本点数量 - negative_points: 负样本点数量 - threshold: 二值化阈值 - min_distance: 点之间最小距离 - edge_margin: 边缘留白 - - Returns: - coordinates_positive: 正坐标 JSON 字符串 - coordinates_negative: 负坐标 JSON 字符串 - visualization: 可视化图像 - """ - - # 处理 mask 张量 - if mask.dim() == 3: - mask = mask[0] # 取第一帧 - elif mask.dim() == 4: - mask = mask[0, 0] # 取第一帧第一通道 - - # 转换为 numpy 数组 - mask_np = mask.cpu().numpy() - height, width = mask_np.shape - - # 二值化处理 - positive_mask = (mask_np >= threshold).astype(np.uint8) - negative_mask = (mask_np < threshold).astype(np.uint8) - - # 应用边缘留白 - if edge_margin > 0: - positive_mask = self._apply_edge_margin(positive_mask, edge_margin) - negative_mask = self._apply_edge_margin(negative_mask, edge_margin) - - # 生成坐标点 - positive_coords = self._sample_points( - positive_mask, positive_points, sample_method, min_distance - ) - negative_coords = self._sample_points( - negative_mask, negative_points, sample_method, min_distance - ) - - # 转换为 JSON 格式 - coordinates_positive = json.dumps([ - {"x": int(x), "y": int(y)} for x, y in positive_coords - ]) - coordinates_negative = json.dumps([ - {"x": int(x), "y": int(y)} for x, y in negative_coords - ]) if negative_coords else json.dumps([]) - - # 生成可视化图像 - visualization = self._create_visualization( - mask_np, positive_coords, negative_coords, threshold - ) - - return coordinates_positive, coordinates_negative, visualization - - def _apply_edge_margin(self, mask: np.ndarray, margin: int) -> np.ndarray: - """应用边缘留白""" - if margin <= 0: - return mask - - result = mask.copy() - result[:margin, :] = 0 - result[-margin:, :] = 0 - result[:, :margin] = 0 - result[:, -margin:] = 0 - return result - - def _sample_points( - self, - mask: np.ndarray, - num_points: int, - method: str, - min_distance: int - ) -> List[Tuple[int, int]]: - """从 mask 中采样点""" - - if num_points <= 0: - return [] - - # 找到所有有效像素位置 - valid_positions = np.where(mask > 0) - if len(valid_positions[0]) == 0: - return [] - - valid_coords = list(zip(valid_positions[1], valid_positions[0])) # (x, y) - - if len(valid_coords) == 0: - return [] - - if method == "random": - return self._random_sampling(valid_coords, num_points, min_distance) - elif method == "grid": - return self._grid_sampling(mask, num_points) - elif method == "contour": - return self._contour_sampling(mask, num_points, min_distance) - else: - return self._random_sampling(valid_coords, num_points, min_distance) - - def _random_sampling( - self, - valid_coords: List[Tuple[int, int]], - num_points: int, - min_distance: int - ) -> List[Tuple[int, int]]: - """随机采样""" - if not valid_coords: - return [] - - selected_points = [] - available_coords = valid_coords.copy() - - for _ in range(min(num_points, len(available_coords))): - if not available_coords: - break - - # 随机选择一个点 - idx = np.random.randint(len(available_coords)) - point = available_coords[idx] - selected_points.append(point) - - # 移除距离太近的点 - if min_distance > 0: - available_coords = [ - coord for coord in available_coords - if np.sqrt((coord[0] - point[0])**2 + (coord[1] - point[1])**2) >= min_distance - ] - else: - available_coords.remove(point) - - return selected_points - - def _grid_sampling(self, mask: np.ndarray, num_points: int) -> List[Tuple[int, int]]: - """网格采样""" - height, width = mask.shape - - # 计算网格大小 - grid_size = int(np.sqrt(num_points)) - if grid_size == 0: - grid_size = 1 - - points = [] - step_y = height // (grid_size + 1) - step_x = width // (grid_size + 1) - - for i in range(1, grid_size + 1): - for j in range(1, grid_size + 1): - y = i * step_y - x = j * step_x - - # 检查该位置是否有效 - if y < height and x < width and mask[y, x] > 0: - points.append((x, y)) - - if len(points) >= num_points: - break - if len(points) >= num_points: - break - - return points - - def _contour_sampling( - self, - mask: np.ndarray, - num_points: int, - min_distance: int - ) -> List[Tuple[int, int]]: - """轮廓采样""" - # 查找轮廓 - contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - - if not contours: - return [] - - # 选择最大的轮廓 - largest_contour = max(contours, key=cv2.contourArea) - - if len(largest_contour) < num_points: - # 如果轮廓点数不够，返回所有轮廓点 - return [(int(point[0][0]), int(point[0][1])) for point in largest_contour] - - # 均匀采样轮廓点 - step = len(largest_contour) // num_points - sampled_points = [] - - for i in range(0, len(largest_contour), step): - if len(sampled_points) >= num_points: - break - point = largest_contour[i][0] - sampled_points.append((int(point[0]), int(point[1]))) - - return sampled_points - - def _create_visualization( - self, - mask: np.ndarray, - positive_coords: List[Tuple[int, int]], - negative_coords: List[Tuple[int, int]], - threshold: float - ) -> torch.Tensor: - """创建可视化图像""" - height, width = mask.shape - - # 创建彩色图像 - vis_image = np.zeros((height, width, 3), dtype=np.uint8) - - # 绘制 mask 背景 - mask_norm = (mask * 255).astype(np.uint8) - vis_image[:, :, 0] = mask_norm # 红色通道显示 mask - vis_image[:, :, 1] = mask_norm # 绿色通道显示 mask - vis_image[:, :, 2] = mask_norm # 蓝色通道显示 mask - - # 绘制正样本点 (绿色) - for x, y in positive_coords: - cv2.circle(vis_image, (x, y), 3, (0, 255, 0), -1) - cv2.circle(vis_image, (x, y), 5, (0, 255, 0), 1) - - # 绘制负样本点 (红色) - for x, y in negative_coords: - cv2.circle(vis_image, (x, y), 3, (0, 0, 255), -1) - cv2.circle(vis_image, (x, y), 5, (0, 0, 255), 1) - - # 转换为 torch 张量 - vis_tensor = torch.from_numpy(vis_image).float() / 255.0 - vis_tensor = vis_tensor.unsqueeze(0) # 添加 batch 维度 - - return vis_tensor - - -# 节点类映射 -NODE_CLASS_MAPPINGS = { - "MaskToCoordinates": MaskToCoordinates, -} - -# 节点显示名称映射 -NODE_DISPLAY_NAME_MAPPINGS = { - "MaskToCoordinates": "VVL Mask转坐标", -} \ No newline at end of file diff --git a/test_path_format.py b/test_path_format.py deleted file mode 100644 index 0418195..0000000 --- a/test_path_format.py +++ /dev/null @@ -1,41 +0,0 @@ -#!/usr/bin/env python3 -""" -测试ComfyUI的带注释文件路径格式 -""" - -import os -import sys - -# 添加ComfyUI路径 -sys.path.append(os.path.join(os.path.dirname(__file__), '../../..')) - -try: - import folder_paths - - def test_annotated_path(): - """测试带注释的文件路径格式""" - # 测试我们使用的路径格式 - test_path = "3d/vggt_3d_model_conf50.0_camTrue.glb [output]" - - print(f"测试路径: {test_path}") - - # 解析带注释的路径 - name, base_dir = folder_paths.annotated_filepath(test_path) - print(f"解析结果: name={name}, base_dir={base_dir}") - - # 获取完整路径 - full_path = folder_paths.get_annotated_filepath(test_path) - print(f"完整路径: {full_path}") - - # 验证这是否与我们预期的output/3d目录匹配 - expected_dir = os.path.join(folder_paths.get_output_directory(), "3d") - print(f"期望目录: {expected_dir}") - - return True - - if __name__ == "__main__": - test_annotated_path() - -except ImportError as e: - print(f"无法导入folder_paths: {e}") - print("请确保在ComfyUI环境中运行此测试") \ No newline at end of file