325 lines
14 KiB
Python
325 lines
14 KiB
Python
"""
|
||
VGGT Native Interface
|
||
原生VGGT功能接口封装 - 基于官方API
|
||
"""
|
||
|
||
import torch
|
||
import numpy as np
|
||
from typing import Dict, Optional, Tuple, List
|
||
import logging
|
||
|
||
logger = logging.getLogger('vggt_native_interface')
|
||
|
||
class VGGTNativeInterface:
|
||
"""VGGT原生功能接口封装 - 基于官方API"""
|
||
|
||
def __init__(self, model: torch.nn.Module, device: torch.device):
|
||
self.model = model # 原生VGGT模型实例
|
||
self.device = device
|
||
|
||
def full_inference_pipeline(self, images: torch.Tensor, query_points: torch.Tensor = None) -> Dict:
|
||
"""完整的VGGT推理管道 - 基于官方最新API"""
|
||
with torch.no_grad():
|
||
with torch.cuda.amp.autocast(dtype=torch.float16):
|
||
# 确保输入是tensor格式，并添加batch维度如果需要
|
||
if isinstance(images, np.ndarray):
|
||
images = torch.from_numpy(images).to(self.device)
|
||
|
||
if images.dim() == 4: # (N, C, H, W)
|
||
images = images[None] # (1, N, C, H, W)
|
||
|
||
# 使用官方的forward方法，获得标准的predictions字典
|
||
predictions = self.model.forward(images, query_points)
|
||
|
||
# 添加原始图像数据（官方predictions_to_glb需要）
|
||
# 转换图像格式：从(1, N, C, H, W)到(N, H, W, C)
|
||
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)
|
||
predictions['images'] = images_formatted
|
||
|
||
# 直接返回标准的VGGT输出格式，添加相机解码结果
|
||
if 'pose_enc' in predictions:
|
||
try:
|
||
extrinsic, intrinsic = self._decode_camera_params(
|
||
predictions['pose_enc'], images.shape[-2:]
|
||
)
|
||
predictions['cameras'] = {'extrinsic': extrinsic, 'intrinsic': intrinsic}
|
||
except Exception as e:
|
||
logger.warning(f"Failed to decode camera parameters: {e}")
|
||
|
||
# 从深度图构建3D点云（如果可能）
|
||
if 'depth' in predictions and 'cameras' in predictions:
|
||
try:
|
||
point_map_unprojected = self._unproject_depth_to_points(
|
||
predictions['depth'].squeeze(0),
|
||
predictions['cameras']['extrinsic'].squeeze(0),
|
||
predictions['cameras']['intrinsic'].squeeze(0)
|
||
)
|
||
predictions['points_from_depth'] = point_map_unprojected
|
||
except Exception as e:
|
||
logger.warning(f"Failed to unproject depth to points: {e}")
|
||
|
||
return predictions
|
||
|
||
def selective_inference(self, images: torch.Tensor, branches: List[str], **kwargs) -> Dict:
|
||
"""选择性推理：只执行指定的分支"""
|
||
with torch.no_grad():
|
||
with torch.cuda.amp.autocast(dtype=torch.float16):
|
||
# 添加batch维度如果需要
|
||
if images.dim() == 4: # (N, C, H, W)
|
||
images = images[None] # (1, N, C, H, W)
|
||
|
||
# 1. 特征聚合 (必须步骤)
|
||
aggregated_tokens_list, ps_idx = self.model.aggregator(images)
|
||
|
||
results = {}
|
||
|
||
# 2. 根据需要执行不同分支
|
||
if 'cameras' in branches:
|
||
pose_enc = self.model.camera_head(aggregated_tokens_list)[-1]
|
||
extrinsic, intrinsic = self._decode_camera_params(pose_enc, images.shape[-2:])
|
||
results['cameras'] = {'extrinsic': extrinsic, 'intrinsic': intrinsic}
|
||
|
||
if 'depth' in branches:
|
||
depth_map, depth_conf = self.model.depth_head(aggregated_tokens_list, images, ps_idx)
|
||
results['depth'] = {'depth_map': depth_map, 'confidence': depth_conf}
|
||
|
||
if 'points' in branches:
|
||
point_map, point_conf = self.model.point_head(aggregated_tokens_list, images, ps_idx)
|
||
results['points'] = {'point_map': point_map, 'confidence': point_conf}
|
||
|
||
if 'points_from_depth' in branches and 'cameras' in results and 'depth' in results:
|
||
point_map_unprojected = self._unproject_depth_to_points(
|
||
results['depth']['depth_map'].squeeze(0),
|
||
results['cameras']['extrinsic'].squeeze(0),
|
||
results['cameras']['intrinsic'].squeeze(0)
|
||
)
|
||
results['points_from_depth'] = point_map_unprojected
|
||
|
||
if 'tracks' in branches and 'query_points' in kwargs:
|
||
track_list, vis_score, conf_score = self.model.track_head(
|
||
aggregated_tokens_list, images, ps_idx,
|
||
query_points=kwargs['query_points'][None]
|
||
)
|
||
results['tracks'] = {
|
||
'track_list': track_list,
|
||
'visibility_score': vis_score,
|
||
'confidence_score': conf_score
|
||
}
|
||
|
||
return results
|
||
|
||
def _decode_camera_params(self, pose_enc: torch.Tensor, image_shape: Tuple[int, int]) -> Tuple[torch.Tensor, torch.Tensor]:
|
||
"""解码相机参数 - 使用VGGT官方工具函数"""
|
||
try:
|
||
from vggt.utils.pose_enc import pose_encoding_to_extri_intri
|
||
return pose_encoding_to_extri_intri(pose_enc, image_shape)
|
||
except ImportError:
|
||
logger.error("VGGT pose_enc utils not available")
|
||
raise
|
||
|
||
def _unproject_depth_to_points(self, depth_map: torch.Tensor, extrinsic: torch.Tensor,
|
||
intrinsic: torch.Tensor) -> torch.Tensor:
|
||
"""从深度图反投影生成3D点 - 使用VGGT官方工具函数"""
|
||
try:
|
||
from vggt.utils.geometry import unproject_depth_map_to_point_map
|
||
return unproject_depth_map_to_point_map(depth_map, extrinsic, intrinsic)
|
||
except ImportError:
|
||
logger.error("VGGT geometry utils not available")
|
||
raise
|
||
|
||
|
||
class VGGTImageProcessor:
|
||
"""VGGT图像预处理器 - 基于官方load_and_preprocess_images函数"""
|
||
|
||
@staticmethod
|
||
def preprocess_images(images: List[np.ndarray], target_size: Tuple[int, int] = (336, 336)) -> torch.Tensor:
|
||
"""
|
||
图像预处理规范 - 遵循VGGT官方标准：
|
||
1. 尺寸调整到336x336（VGGT标准输入尺寸）
|
||
2. 归一化到[0,1]范围
|
||
3. 转换为torch.Tensor格式
|
||
4. 通道顺序RGB
|
||
5. 支持masking（设置像素值为0或1来屏蔽不需要的区域）
|
||
"""
|
||
try:
|
||
# 使用VGGT官方预处理函数
|
||
from vggt.utils.load_fn import load_and_preprocess_images
|
||
|
||
# 如果输入是numpy数组列表，需要先保存为临时文件
|
||
import tempfile
|
||
import os
|
||
from PIL import Image
|
||
|
||
temp_paths = []
|
||
temp_dir = tempfile.mkdtemp()
|
||
|
||
try:
|
||
# 保存图像到临时文件
|
||
for i, img_array in enumerate(images):
|
||
if isinstance(img_array, np.ndarray):
|
||
# 确保是RGB格式
|
||
if img_array.shape[-1] == 3:
|
||
img_pil = Image.fromarray(img_array.astype(np.uint8))
|
||
else:
|
||
img_pil = Image.fromarray(img_array.astype(np.uint8)).convert('RGB')
|
||
else:
|
||
img_pil = img_array
|
||
|
||
temp_path = os.path.join(temp_dir, f"temp_{i:04d}.jpg")
|
||
img_pil.save(temp_path)
|
||
temp_paths.append(temp_path)
|
||
|
||
# 使用官方预处理函数
|
||
processed_images = load_and_preprocess_images(temp_paths)
|
||
|
||
return processed_images
|
||
|
||
finally:
|
||
# 清理临时文件
|
||
for temp_path in temp_paths:
|
||
if os.path.exists(temp_path):
|
||
os.remove(temp_path)
|
||
os.rmdir(temp_dir)
|
||
|
||
except ImportError:
|
||
# 如果官方函数不可用，使用备用方案
|
||
return VGGTImageProcessor._fallback_preprocess(images, target_size)
|
||
|
||
@staticmethod
|
||
def _fallback_preprocess(images: List[np.ndarray], target_size: Tuple[int, int]) -> torch.Tensor:
|
||
"""备用图像预处理方案"""
|
||
import torch
|
||
import torchvision.transforms as transforms
|
||
from PIL import Image
|
||
|
||
transform = transforms.Compose([
|
||
transforms.Resize(target_size),
|
||
transforms.ToTensor(), # 自动归一化到[0,1]
|
||
])
|
||
|
||
processed = []
|
||
for img_array in images:
|
||
if isinstance(img_array, np.ndarray):
|
||
# 转换为PIL Image
|
||
if img_array.dtype != np.uint8:
|
||
img_array = (img_array * 255).astype(np.uint8)
|
||
pil_image = Image.fromarray(img_array).convert('RGB')
|
||
else:
|
||
pil_image = img_array
|
||
|
||
tensor_img = transform(pil_image)
|
||
processed.append(tensor_img)
|
||
|
||
return torch.stack(processed)
|
||
|
||
@staticmethod
|
||
def apply_masking(images: torch.Tensor, mask_type: str = None) -> torch.Tensor:
|
||
"""
|
||
应用遮罩 - VGGT支持简单的像素遮罩
|
||
mask_type: 'black_bg', 'white_bg', 'sky' 等
|
||
"""
|
||
if mask_type is None:
|
||
return images
|
||
|
||
masked_images = images.clone()
|
||
|
||
if mask_type == 'black_bg':
|
||
# 将黑色像素（接近0的像素）设置为0
|
||
black_mask = (images.sum(dim=1, keepdim=True) < 0.1) # RGB总和小于0.1
|
||
masked_images[black_mask.expand_as(images)] = 0.0
|
||
|
||
elif mask_type == 'white_bg':
|
||
# 将白色像素（接近1的像素）设置为1
|
||
white_mask = (images.sum(dim=1, keepdim=True) > 2.7) # RGB总和大于2.7
|
||
masked_images[white_mask.expand_as(images)] = 1.0
|
||
|
||
# 可以添加更多遮罩类型
|
||
|
||
return masked_images
|
||
|
||
|
||
class VGGTResultProcessor:
|
||
"""VGGT结果后处理器"""
|
||
|
||
@staticmethod
|
||
def format_camera_results(extrinsic: torch.Tensor, intrinsic: torch.Tensor) -> Dict:
|
||
"""相机参数结果格式化，同时保留原生Tensor输出"""
|
||
# 确保输入是tensor格式
|
||
if isinstance(extrinsic, np.ndarray):
|
||
extrinsic = torch.from_numpy(extrinsic)
|
||
if isinstance(intrinsic, np.ndarray):
|
||
intrinsic = torch.from_numpy(intrinsic)
|
||
|
||
return {
|
||
"raw": {
|
||
"extrinsic": extrinsic, # 原生Tensor
|
||
"intrinsic": intrinsic,
|
||
},
|
||
"json": {
|
||
"extrinsic_matrices": extrinsic.cpu().numpy().tolist(),
|
||
"intrinsic_matrices": intrinsic.cpu().numpy().tolist(),
|
||
"format": "opencv_convention",
|
||
"coordinate_system": "camera_from_world",
|
||
}
|
||
}
|
||
|
||
@staticmethod
|
||
def format_depth_results(depth_map: torch.Tensor, confidence: torch.Tensor) -> Dict:
|
||
"""深度结果格式化，同时保留原生Tensor输出"""
|
||
# 确保输入是tensor格式
|
||
if isinstance(depth_map, np.ndarray):
|
||
depth_map = torch.from_numpy(depth_map)
|
||
if isinstance(confidence, np.ndarray):
|
||
confidence = torch.from_numpy(confidence)
|
||
|
||
return {
|
||
"raw": {
|
||
"depth_map": depth_map,
|
||
"confidence": confidence,
|
||
},
|
||
"json": {
|
||
"depth_shape": list(depth_map.shape),
|
||
"depth_range": [float(depth_map.min()), float(depth_map.max())],
|
||
"confidence_range": [float(confidence.min()), float(confidence.max())],
|
||
}
|
||
}
|
||
|
||
@staticmethod
|
||
def format_point_cloud(points: torch.Tensor, colors: torch.Tensor = None, confidence: torch.Tensor = None) -> Dict:
|
||
"""点云结果格式化，同时保留原生Tensor输出"""
|
||
# 确保输入是tensor格式
|
||
if isinstance(points, np.ndarray):
|
||
points = torch.from_numpy(points)
|
||
if colors is not None and isinstance(colors, np.ndarray):
|
||
colors = torch.from_numpy(colors)
|
||
if confidence is not None and isinstance(confidence, np.ndarray):
|
||
confidence = torch.from_numpy(confidence)
|
||
|
||
result = {
|
||
"raw": {
|
||
"points": points,
|
||
"colors": colors,
|
||
"confidence": confidence,
|
||
},
|
||
"json": {
|
||
"num_points": int(points.shape[0]) if points.dim() > 1 else int(points.numel() // 3),
|
||
"point_shape": list(points.shape),
|
||
"bounds": {
|
||
"min": points.min(dim=0)[0].cpu().numpy().tolist() if points.numel() > 0 else [0, 0, 0],
|
||
"max": points.max(dim=0)[0].cpu().numpy().tolist() if points.numel() > 0 else [0, 0, 0],
|
||
}
|
||
}
|
||
}
|
||
|
||
if colors is not None:
|
||
result["json"]["has_colors"] = True
|
||
result["json"]["color_shape"] = list(colors.shape)
|
||
|
||
if confidence is not None:
|
||
result["json"]["has_confidence"] = True
|
||
result["json"]["confidence_range"] = [float(confidence.min()), float(confidence.max())]
|
||
|
||
return result |