diff --git a/MASK_TO_COORDINATES_README.md b/MASK_TO_COORDINATES_README.md new file mode 100644 index 0000000..0a9a514 --- /dev/null +++ b/MASK_TO_COORDINATES_README.md @@ -0,0 +1,109 @@ +# 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 a7a4c7f..53169bf 100644 --- a/__init__.py +++ b/__init__.py @@ -6,6 +6,8 @@ ComfyUI VVL Video Camera Advanced from .comfyui_vggt_nodes import VGGTMultiInputNode from .vggt_model_loader import VVLVGGTLoader from .glb_point_cloud_processor import * +from .mask_to_coordinates import MaskToCoordinates +from .json_merger import JSONMerger, UESceneGenerator # 节点类映射 NODE_CLASS_MAPPINGS = { @@ -14,6 +16,10 @@ NODE_CLASS_MAPPINGS = { "GLBPointCloudProcessor": GLBPointCloudProcessor, "GLBPointCloudBounds": GLBPointCloudBounds, "GLBPointCloudOriginAdjuster": GLBPointCloudOriginAdjuster, + "GLBPointCloudDensityFilter": GLBPointCloudDensityFilter, + "MaskToCoordinates": MaskToCoordinates, + "JSONMerger": JSONMerger, + "UESceneGenerator": UESceneGenerator, } # 节点显示名称映射 @@ -23,6 +29,10 @@ NODE_DISPLAY_NAME_MAPPINGS = { "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_nodes.py b/comfyui_vggt_nodes.py index 5d80f91..ac0174c 100644 --- a/comfyui_vggt_nodes.py +++ b/comfyui_vggt_nodes.py @@ -529,20 +529,61 @@ def _matrices_to_json(intrinsic, extrinsic, source_type="video") -> (str, str): json.dumps(poses_data, ensure_ascii=False, indent=2), ) -def _create_traj_preview(extrinsic: torch.Tensor) -> torch.Tensor: +def _create_traj_preview(extrinsic) -> torch.Tensor: """根据相机外参创建3D轨迹可视化（使用matplotlib 3D绘图）。""" - ext = extrinsic.cpu().numpy() # (N,3,4) + # 支持torch.Tensor和numpy.ndarray输入 + if isinstance(extrinsic, torch.Tensor): + ext = extrinsic.cpu().numpy() + elif isinstance(extrinsic, np.ndarray): + ext = extrinsic + else: + print(f"VGGT: 不支持的外参数据类型: {type(extrinsic)}") + return _create_insufficient_data_image() + + print(f"VGGT: 轨迹预览处理外参矩阵形状: {ext.shape}") + positions = [] orientations = [] - for mat in ext: - R = mat[:3, :3] - t = mat[:3, 3] - pos = -R.T @ t # 相机在世界坐标系中的位置 + # 检查矩阵形状并相应处理 + if len(ext.shape) == 3 and ext.shape[-1] == 4: # (N, 3, 4) 格式 + print(f"VGGT: 使用完整外参矩阵 (N, 3, 4)") + for mat in ext: + R = mat[:3, :3] + t = mat[:3, 3] + pos = -R.T @ t # 相机在世界坐标系中的位置 + positions.append(pos) + # 提取相机朝向（Z轴方向） + forward = -R[:, 2] # 相机朝向（Z轴负方向） + orientations.append(forward) + elif len(ext.shape) == 3 and ext.shape[-1] == 3: # (N, 3, 3) 格式，只有旋转矩阵 + print(f"VGGT: 外参矩阵只包含旋转信息 {ext.shape}，使用索引作为位置") + for i, mat in enumerate(ext): + R = mat[:3, :3] + # 没有平移信息，使用索引作为位置 + pos = np.array([i * 1.0, 0.0, 0.0]) # 简单的线性排列 + positions.append(pos) + # 提取相机朝向（Z轴方向） + forward = -R[:, 2] # 相机朝向（Z轴负方向） + orientations.append(forward) + elif len(ext.shape) == 2 and ext.shape == (3, 4): # 单个 (3, 4) 矩阵 + print(f"VGGT: 单个外参矩阵 (3, 4)") + R = ext[:3, :3] + t = ext[:3, 3] + pos = -R.T @ t positions.append(pos) - # 提取相机朝向（Z轴方向） - forward = -R[:, 2] # 相机朝向（Z轴负方向） + forward = -R[:, 2] orientations.append(forward) + elif len(ext.shape) == 2 and ext.shape == (3, 3): # 单个 (3, 3) 矩阵 + print(f"VGGT: 单个旋转矩阵 (3, 3)") + R = ext[:3, :3] + pos = np.array([0.0, 0.0, 0.0]) # 原点位置 + positions.append(pos) + forward = -R[:, 2] + orientations.append(forward) + else: + print(f"VGGT: 不支持的外参矩阵形状: {ext.shape}") + return _create_insufficient_data_image() positions = np.array(positions) orientations = np.array(orientations) @@ -1265,9 +1306,22 @@ class VGGTMultiInputNode: trajectory_preview = _create_insufficient_data_image() # 默认图像 if 'cameras' in raw_results: try: - trajectory_preview = _create_traj_preview(raw_results['cameras']['extrinsic']) + 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}") + trajectory_preview = _create_traj_preview(extrinsic_for_preview) except Exception as e: logger.warning(f"轨迹预览生成失败: {e}") + import traceback + traceback.print_exc() # 生成3D模型文件 model_3d_path = "" diff --git a/glb_point_cloud_processor.py b/glb_point_cloud_processor.py index 5a8f0ee..e99fdf0 100755 --- a/glb_point_cloud_processor.py +++ b/glb_point_cloud_processor.py @@ -46,8 +46,8 @@ class GLBPointCloudProcessor: "tooltip": "黑色阈值控制：RGB三色总和小于此值的点将被删除。0=仅删除纯黑点(0,0,0)；30=删除深黑色点；60=删除较暗点；100=删除中等暗度点；255=删除所有点。调大=删除更多暗色点，点云更亮净；调小=保留更多暗色细节，但可能有噪点" }), "output_filename": ("STRING", { - "default": "clean_pointcloud.glb", - "tooltip": "输出GLB文件名：处理后点云的保存文件名，自动添加.glb扩展名。文件将保存到ComfyUI的output目录中。建议使用描述性名称如'clean_pointcloud.glb'、'filtered_scan.glb'等便于识别" + "default": "clean_pointcloud", + "tooltip": "输出GLB文件名：处理后点云的保存文件名，系统会自动添加.glb扩展名。文件将保存到ComfyUI的output目录中。建议使用描述性名称如'clean_pointcloud'、'filtered_scan'等便于识别" }), } } @@ -68,7 +68,7 @@ class GLBPointCloudProcessor: def process_point_cloud(self, glb_file_path: str, black_threshold: int = 30, - output_filename: str = "clean_pointcloud.glb"): + output_filename: str = "clean_pointcloud"): """ 处理GLB点云文件，删除黑色点和暗色点 """ @@ -375,8 +375,8 @@ class GLBPointCloudBounds: "tooltip": "增强可见性：添加顶点高亮和面中心标记，让包围盒更容易识别" }), "output_filename": ("STRING", { - "default": "pointcloud_with_bounds.glb", - "tooltip": "输出文件名：带包围盒可视化的点云文件名" + "default": "pointcloud_with_bounds", + "tooltip": "输出文件名：带包围盒可视化的点云文件名，系统会自动添加.glb扩展名" }), } } @@ -405,7 +405,7 @@ class GLBPointCloudBounds: add_coordinate_axes: bool = True, wireframe_density: int = 50, enhance_visibility: bool = True, - output_filename: str = "pointcloud_with_bounds.glb"): + output_filename: str = "pointcloud_with_bounds"): """ 计算GLB点云文件的包围盒尺寸并生成带可视化的点云文件 """ @@ -538,6 +538,7 @@ class GLBPointCloudBounds: # 生成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 } @@ -1044,8 +1045,8 @@ class GLBPointCloudOriginAdjuster: "tooltip": "坐标轴密度：每条坐标轴的点数，越高坐标轴越清晰但点数越多。推荐80-120获得清晰效果" }), "output_filename": ("STRING", { - "default": "adjusted_pointcloud.glb", - "tooltip": "输出文件名：调整后的点云文件名" + "default": "adjusted_pointcloud", + "tooltip": "输出文件名：调整后的点云文件名，系统会自动添加.glb扩展名" }), } } @@ -1072,7 +1073,7 @@ class GLBPointCloudOriginAdjuster: output_units: str = "centimeters", add_coordinate_axes: bool = False, wireframe_density: int = 100, - output_filename: str = "adjusted_pointcloud.glb"): + output_filename: str = "adjusted_pointcloud"): """ 调整GLB点云的原点位置和旋转 """ @@ -1303,6 +1304,7 @@ class GLBPointCloudOriginAdjuster: # 生成变换信息（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])] } @@ -1596,6 +1598,419 @@ class GLBPointCloudOriginAdjuster: 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) + + # ----------------------------------------------------------------------------- # 节点注册 - 更新注册信息 # ----------------------------------------------------------------------------- @@ -1604,14 +2019,16 @@ 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 +# __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/json_merger.py b/json_merger.py new file mode 100644 index 0000000..e41b144 --- /dev/null +++ b/json_merger.py @@ -0,0 +1,494 @@ +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 new file mode 100644 index 0000000..1bf82ef --- /dev/null +++ b/mask_to_coordinates.py @@ -0,0 +1,328 @@ +""" +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