336 lines
14 KiB
Python
Executable File
336 lines
14 KiB
Python
Executable File
import os
|
|
import json
|
|
import tempfile
|
|
import logging
|
|
import struct
|
|
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.glb",
|
|
"tooltip": "输出GLB文件名：处理后点云的保存文件名，自动添加.glb扩展名。文件将保存到ComfyUI的output目录中。建议使用描述性名称如'clean_pointcloud.glb'、'filtered_scan.glb'等便于识别"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"STRING", # 处理后的GLB文件路径
|
|
"STRING", # 处理统计信息（JSON格式）
|
|
"STRING", # 处理日志信息
|
|
)
|
|
RETURN_NAMES = (
|
|
"processed_glb_path",
|
|
"processing_stats",
|
|
"processing_log",
|
|
)
|
|
OUTPUT_TOOLTIPS = [
|
|
"处理后的GLB文件完整路径",
|
|
"处理统计信息（JSON格式）- 包含删除的点数、剩余点数等",
|
|
"详细的处理日志信息"
|
|
]
|
|
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"):
|
|
"""
|
|
处理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 "", json.dumps({"error": error_msg}), "\n".join(processing_log)
|
|
|
|
# 验证输入文件路径
|
|
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 "", json.dumps({"error": error_msg}), "\n".join(processing_log)
|
|
|
|
# 处理文件路径
|
|
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 "", json.dumps({"error": error_msg}), "\n".join(processing_log)
|
|
|
|
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 "", json.dumps({"error": error_msg}), "\n".join(processing_log)
|
|
|
|
if not point_clouds:
|
|
error_msg = "GLB文件中没有点云数据"
|
|
processing_log.append(f"错误: {error_msg}")
|
|
return "", json.dumps({"error": error_msg}), "\n".join(processing_log)
|
|
|
|
# 处理每个点云
|
|
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
|
|
|
|
# 获取顶点和颜色
|
|
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(vertices), dtype=bool)
|
|
|
|
# 删除暗色点（默认启用）
|
|
if colors is not None:
|
|
removed_count = self._remove_dark_points(vertices, colors, keep_mask, black_threshold, processing_log)
|
|
processing_log.append(f"删除黑色点: {removed_count} 个")
|
|
else:
|
|
processing_log.append("跳过黑色点过滤: 点云无颜色信息")
|
|
|
|
# 应用掩码
|
|
filtered_vertices = 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 = trimesh.PointCloud(vertices=filtered_vertices, colors=filtered_colors)
|
|
processed_point_clouds.append((name, processed_pc))
|
|
else:
|
|
processing_log.append(f"警告: 点云 {name} 处理后没有剩余点")
|
|
|
|
# 创建新的场景
|
|
if processed_point_clouds:
|
|
new_scene = trimesh.Scene()
|
|
|
|
# 添加处理后的点云
|
|
for name, pc in processed_point_clouds:
|
|
new_scene.add_geometry(pc, node_name=name)
|
|
|
|
# 添加其他几何体（如相机模型）
|
|
for name, geometry in other_geometries:
|
|
new_scene.add_geometry(geometry, node_name=name)
|
|
processing_log.append(f"保留几何体: {name}")
|
|
|
|
# 生成输出路径
|
|
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")
|
|
|
|
# 生成统计信息
|
|
stats = {
|
|
"original_points": total_original_points,
|
|
"remaining_points": total_original_points - total_removed_points,
|
|
"removed_points": total_removed_points,
|
|
"removal_percentage": (total_removed_points / total_original_points * 100) if total_original_points > 0 else 0,
|
|
"point_clouds_processed": len(point_clouds),
|
|
"other_geometries_preserved": len(other_geometries),
|
|
"output_file_size": file_size,
|
|
"settings": {
|
|
"black_threshold": black_threshold,
|
|
"remove_dark_points": True
|
|
}
|
|
}
|
|
|
|
processing_log.append("GLB点云黑色点清理完成!")
|
|
return output_path, json.dumps(stats, indent=2), "\n".join(processing_log)
|
|
else:
|
|
error_msg = "GLB文件保存失败"
|
|
processing_log.append(f"错误: {error_msg}")
|
|
return "", json.dumps({"error": error_msg}), "\n".join(processing_log)
|
|
else:
|
|
error_msg = "处理后没有剩余的点云数据"
|
|
processing_log.append(f"错误: {error_msg}")
|
|
return "", json.dumps({"error": error_msg}), "\n".join(processing_log)
|
|
|
|
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 "", json.dumps({"error": error_msg}), "\n".join(processing_log)
|
|
|
|
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
|
|
|
|
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# 节点注册 - 独立注册，ComfyUI会自动发现此文件
|
|
# -----------------------------------------------------------------------------
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GLBPointCloudProcessor": GLBPointCloudProcessor,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GLBPointCloudProcessor": "VVL GLB Point Cloud Processor",
|
|
}
|
|
|
|
# # 添加节点信息，帮助ComfyUI更好地识别
|
|
# WEB_DIRECTORY = "./web"
|
|
# __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] |