update
This commit is contained in:
wTechArtist
@@ -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
|
||||
+27
-18
@@ -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']
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
File diff suppressed because it is too large
Load Diff
-494
@@ -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"]
|
||||
@@ -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转坐标",
|
||||
}
|
||||
@@ -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环境中运行此测试")
|
||||
Reference in New Issue
Block a user