510 lines
17 KiB
Python
510 lines
17 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
"""
|
|
GLB 3D export utilities for GCT predictions.
|
|
"""
|
|
|
|
import os
|
|
import copy
|
|
from typing import Optional, Tuple
|
|
|
|
import numpy as np
|
|
import cv2
|
|
import matplotlib
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
from lingbot_map.vis.sky_segmentation import (
|
|
_SKYSEG_INPUT_SIZE,
|
|
_SKYSEG_SOFT_THRESHOLD,
|
|
_mask_to_float,
|
|
_mask_to_uint8,
|
|
_result_map_to_non_sky_conf,
|
|
)
|
|
|
|
try:
|
|
import trimesh
|
|
except ImportError:
|
|
trimesh = None
|
|
print("trimesh not found. GLB export will not work.")
|
|
|
|
|
|
def predictions_to_glb(
|
|
predictions: dict,
|
|
conf_thres: float = 50.0,
|
|
filter_by_frames: str = "all",
|
|
mask_black_bg: bool = False,
|
|
mask_white_bg: bool = False,
|
|
show_cam: bool = True,
|
|
mask_sky: bool = False,
|
|
target_dir: Optional[str] = None,
|
|
prediction_mode: str = "Predicted Pointmap",
|
|
) -> "trimesh.Scene":
|
|
"""
|
|
Converts GCT predictions to a 3D scene represented as a GLB file.
|
|
|
|
Args:
|
|
predictions: Dictionary containing model predictions with keys:
|
|
- world_points: 3D point coordinates (S, H, W, 3)
|
|
- world_points_conf: Confidence scores (S, H, W)
|
|
- images: Input images (S, H, W, 3) or (S, 3, H, W)
|
|
- extrinsic: Camera extrinsic matrices (S, 3, 4)
|
|
conf_thres: Percentage of low-confidence points to filter out
|
|
filter_by_frames: Frame filter specification ("all" or frame index)
|
|
mask_black_bg: Mask out black background pixels
|
|
mask_white_bg: Mask out white background pixels
|
|
show_cam: Include camera visualization
|
|
mask_sky: Apply sky segmentation mask
|
|
target_dir: Output directory for intermediate files
|
|
prediction_mode: "Predicted Pointmap" or "Predicted Depthmap"
|
|
|
|
Returns:
|
|
trimesh.Scene: Processed 3D scene containing point cloud and cameras
|
|
|
|
Raises:
|
|
ValueError: If input predictions structure is invalid
|
|
ImportError: If trimesh is not available
|
|
"""
|
|
if trimesh is None:
|
|
raise ImportError("trimesh is required for GLB export. Install with: pip install trimesh")
|
|
|
|
if not isinstance(predictions, dict):
|
|
raise ValueError("predictions must be a dictionary")
|
|
|
|
if conf_thres is None:
|
|
conf_thres = 10.0
|
|
|
|
print("Building GLB scene")
|
|
|
|
# Parse frame filter
|
|
selected_frame_idx = None
|
|
if filter_by_frames != "all" and filter_by_frames != "All":
|
|
try:
|
|
selected_frame_idx = int(filter_by_frames.split(":")[0])
|
|
except (ValueError, IndexError):
|
|
pass
|
|
|
|
# Select prediction source
|
|
if "Pointmap" in prediction_mode:
|
|
print("Using Pointmap Branch")
|
|
if "world_points" in predictions:
|
|
pred_world_points = predictions["world_points"]
|
|
pred_world_points_conf = predictions.get(
|
|
"world_points_conf", np.ones_like(pred_world_points[..., 0])
|
|
)
|
|
else:
|
|
print("Warning: world_points not found, falling back to depth-based points")
|
|
pred_world_points = predictions["world_points_from_depth"]
|
|
pred_world_points_conf = predictions.get(
|
|
"depth_conf", np.ones_like(pred_world_points[..., 0])
|
|
)
|
|
else:
|
|
print("Using Depthmap and Camera Branch")
|
|
pred_world_points = predictions["world_points_from_depth"]
|
|
pred_world_points_conf = predictions.get(
|
|
"depth_conf", np.ones_like(pred_world_points[..., 0])
|
|
)
|
|
|
|
images = predictions["images"]
|
|
camera_matrices = predictions["extrinsic"]
|
|
|
|
# Apply sky segmentation if enabled
|
|
if mask_sky and target_dir is not None:
|
|
pred_world_points_conf = _apply_sky_mask(
|
|
pred_world_points_conf, target_dir, images
|
|
)
|
|
|
|
# Apply frame filter
|
|
if selected_frame_idx is not None:
|
|
pred_world_points = pred_world_points[selected_frame_idx][None]
|
|
pred_world_points_conf = pred_world_points_conf[selected_frame_idx][None]
|
|
images = images[selected_frame_idx][None]
|
|
camera_matrices = camera_matrices[selected_frame_idx][None]
|
|
|
|
# Prepare vertices and colors
|
|
vertices_3d = pred_world_points.reshape(-1, 3)
|
|
|
|
# Handle different image formats
|
|
if images.ndim == 4 and images.shape[1] == 3: # NCHW format
|
|
colors_rgb = np.transpose(images, (0, 2, 3, 1))
|
|
else:
|
|
colors_rgb = images
|
|
colors_rgb = (colors_rgb.reshape(-1, 3) * 255).astype(np.uint8)
|
|
|
|
# Apply confidence filtering
|
|
conf = pred_world_points_conf.reshape(-1)
|
|
conf_threshold = np.percentile(conf, conf_thres) if conf_thres > 0 else 0.0
|
|
conf_mask = (conf >= conf_threshold) & (conf > 1e-5)
|
|
|
|
# Apply background masking
|
|
if mask_black_bg:
|
|
black_bg_mask = colors_rgb.sum(axis=1) >= 16
|
|
conf_mask = conf_mask & black_bg_mask
|
|
|
|
if mask_white_bg:
|
|
white_bg_mask = ~(
|
|
(colors_rgb[:, 0] > 240) &
|
|
(colors_rgb[:, 1] > 240) &
|
|
(colors_rgb[:, 2] > 240)
|
|
)
|
|
conf_mask = conf_mask & white_bg_mask
|
|
|
|
vertices_3d = vertices_3d[conf_mask]
|
|
colors_rgb = colors_rgb[conf_mask]
|
|
|
|
# Handle empty point cloud
|
|
if vertices_3d is None or np.asarray(vertices_3d).size == 0:
|
|
vertices_3d = np.array([[1, 0, 0]])
|
|
colors_rgb = np.array([[255, 255, 255]])
|
|
scene_scale = 1
|
|
else:
|
|
lower_percentile = np.percentile(vertices_3d, 5, axis=0)
|
|
upper_percentile = np.percentile(vertices_3d, 95, axis=0)
|
|
scene_scale = np.linalg.norm(upper_percentile - lower_percentile)
|
|
|
|
colormap = matplotlib.colormaps.get_cmap("gist_rainbow")
|
|
|
|
# Build scene
|
|
scene_3d = trimesh.Scene()
|
|
point_cloud_data = trimesh.PointCloud(vertices=vertices_3d, colors=colors_rgb)
|
|
scene_3d.add_geometry(point_cloud_data)
|
|
|
|
# Prepare camera matrices
|
|
num_cameras = len(camera_matrices)
|
|
extrinsics_matrices = np.zeros((num_cameras, 4, 4))
|
|
extrinsics_matrices[:, :3, :4] = camera_matrices
|
|
extrinsics_matrices[:, 3, 3] = 1
|
|
|
|
# Add cameras
|
|
if show_cam:
|
|
for i in range(num_cameras):
|
|
world_to_camera = extrinsics_matrices[i]
|
|
camera_to_world = np.linalg.inv(world_to_camera)
|
|
rgba_color = colormap(i / num_cameras)
|
|
current_color = tuple(int(255 * x) for x in rgba_color[:3])
|
|
integrate_camera_into_scene(scene_3d, camera_to_world, current_color, scene_scale)
|
|
|
|
# Align scene
|
|
scene_3d = apply_scene_alignment(scene_3d, extrinsics_matrices)
|
|
|
|
print("GLB Scene built")
|
|
return scene_3d
|
|
|
|
|
|
def _apply_sky_mask(
|
|
conf: np.ndarray,
|
|
target_dir: str,
|
|
images: np.ndarray
|
|
) -> np.ndarray:
|
|
"""Apply sky segmentation mask to confidence scores."""
|
|
try:
|
|
import onnxruntime
|
|
except ImportError:
|
|
print("Warning: onnxruntime not available, skipping sky masking")
|
|
return conf
|
|
|
|
target_dir_images = os.path.join(target_dir, "images")
|
|
if not os.path.exists(target_dir_images):
|
|
print(f"Warning: Images directory not found at {target_dir_images}")
|
|
return conf
|
|
|
|
image_list = sorted(os.listdir(target_dir_images))
|
|
S, H, W = conf.shape if hasattr(conf, "shape") else (len(images), images.shape[1], images.shape[2])
|
|
|
|
skyseg_model_path = "skyseg.onnx"
|
|
if not os.path.exists(skyseg_model_path):
|
|
print("Downloading skyseg.onnx...")
|
|
download_file_from_url(
|
|
"https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx",
|
|
skyseg_model_path
|
|
)
|
|
|
|
skyseg_session = onnxruntime.InferenceSession(skyseg_model_path)
|
|
sky_mask_list = []
|
|
|
|
for i, image_name in enumerate(image_list[:S]):
|
|
image_filepath = os.path.join(target_dir_images, image_name)
|
|
mask_filepath = os.path.join(target_dir, "sky_masks", image_name)
|
|
|
|
if os.path.exists(mask_filepath):
|
|
sky_mask = cv2.imread(mask_filepath, cv2.IMREAD_GRAYSCALE)
|
|
else:
|
|
sky_mask = segment_sky(image_filepath, skyseg_session, mask_filepath)
|
|
|
|
if sky_mask.shape[0] != H or sky_mask.shape[1] != W:
|
|
sky_mask = cv2.resize(sky_mask, (W, H), interpolation=cv2.INTER_LINEAR)
|
|
|
|
sky_mask_list.append(_mask_to_float(sky_mask))
|
|
|
|
sky_mask_array = np.array(sky_mask_list)
|
|
sky_mask_binary = (sky_mask_array > _SKYSEG_SOFT_THRESHOLD).astype(np.float32)
|
|
return conf * sky_mask_binary
|
|
|
|
|
|
def integrate_camera_into_scene(
|
|
scene: "trimesh.Scene",
|
|
transform: np.ndarray,
|
|
face_colors: Tuple[int, int, int],
|
|
scene_scale: float,
|
|
frustum_thickness: float = 1.0,
|
|
):
|
|
"""
|
|
Integrates a camera mesh into the 3D scene.
|
|
|
|
Args:
|
|
scene: The 3D scene to add the camera model
|
|
transform: Transformation matrix for camera positioning
|
|
face_colors: RGB color tuple for the camera
|
|
scene_scale: Scale of the scene
|
|
frustum_thickness: Multiplier for frustum edge thickness (>1 = thicker)
|
|
"""
|
|
cam_width = scene_scale * 0.05
|
|
cam_height = scene_scale * 0.1
|
|
|
|
rot_45_degree = np.eye(4)
|
|
rot_45_degree[:3, :3] = Rotation.from_euler("z", 45, degrees=True).as_matrix()
|
|
rot_45_degree[2, 3] = -cam_height
|
|
|
|
opengl_transform = get_opengl_conversion_matrix()
|
|
complete_transform = transform @ opengl_transform @ rot_45_degree
|
|
camera_cone_shape = trimesh.creation.cone(cam_width, cam_height, sections=4)
|
|
|
|
# Build thicker frustum by stacking rotated copies
|
|
slight_rotation = np.eye(4)
|
|
slight_rotation[:3, :3] = Rotation.from_euler("z", 2, degrees=True).as_matrix()
|
|
|
|
shell_scales = [1.0, 0.95]
|
|
shell_transforms = [np.eye(4), slight_rotation]
|
|
# Add extra shells for thickness
|
|
if frustum_thickness > 1.0:
|
|
n_extra = max(1, int(frustum_thickness - 1))
|
|
for k in range(1, n_extra + 1):
|
|
# Progressively rotated and scaled copies
|
|
angle = 2.0 + k * 2.0
|
|
scale = 1.0 + k * 0.02
|
|
rot = np.eye(4)
|
|
rot[:3, :3] = Rotation.from_euler("z", angle, degrees=True).as_matrix()
|
|
shell_scales.append(scale)
|
|
shell_transforms.append(rot)
|
|
rot_neg = np.eye(4)
|
|
rot_neg[:3, :3] = Rotation.from_euler("z", -angle, degrees=True).as_matrix()
|
|
shell_scales.append(scale)
|
|
shell_transforms.append(rot_neg)
|
|
|
|
vertices_parts = []
|
|
for s, t_mat in zip(shell_scales, shell_transforms):
|
|
vertices_parts.append(
|
|
transform_points(t_mat, s * camera_cone_shape.vertices)
|
|
)
|
|
vertices_combined = np.concatenate(vertices_parts)
|
|
vertices_transformed = transform_points(complete_transform, vertices_combined)
|
|
|
|
mesh_faces = compute_camera_faces_multi(camera_cone_shape, len(shell_scales))
|
|
camera_mesh = trimesh.Trimesh(vertices=vertices_transformed, faces=mesh_faces)
|
|
camera_mesh.visual.face_colors[:, :3] = face_colors
|
|
scene.add_geometry(camera_mesh)
|
|
|
|
|
|
def apply_scene_alignment(
|
|
scene_3d: "trimesh.Scene",
|
|
extrinsics_matrices: np.ndarray
|
|
) -> "trimesh.Scene":
|
|
"""
|
|
Aligns the 3D scene based on the extrinsics of the first camera.
|
|
|
|
Args:
|
|
scene_3d: The 3D scene to be aligned
|
|
extrinsics_matrices: Camera extrinsic matrices
|
|
|
|
Returns:
|
|
Aligned 3D scene
|
|
"""
|
|
opengl_conversion_matrix = get_opengl_conversion_matrix()
|
|
|
|
align_rotation = np.eye(4)
|
|
align_rotation[:3, :3] = Rotation.from_euler("y", 180, degrees=True).as_matrix()
|
|
|
|
initial_transformation = (
|
|
np.linalg.inv(extrinsics_matrices[0]) @ opengl_conversion_matrix @ align_rotation
|
|
)
|
|
scene_3d.apply_transform(initial_transformation)
|
|
return scene_3d
|
|
|
|
|
|
def get_opengl_conversion_matrix() -> np.ndarray:
|
|
"""Returns the OpenGL conversion matrix (flips Y and Z axes)."""
|
|
matrix = np.identity(4)
|
|
matrix[1, 1] = -1
|
|
matrix[2, 2] = -1
|
|
return matrix
|
|
|
|
|
|
def transform_points(
|
|
transformation: np.ndarray,
|
|
points: np.ndarray,
|
|
dim: Optional[int] = None
|
|
) -> np.ndarray:
|
|
"""
|
|
Applies a 4x4 transformation to a set of points.
|
|
|
|
Args:
|
|
transformation: Transformation matrix
|
|
points: Points to be transformed
|
|
dim: Dimension for reshaping the result
|
|
|
|
Returns:
|
|
Transformed points
|
|
"""
|
|
points = np.asarray(points)
|
|
initial_shape = points.shape[:-1]
|
|
dim = dim or points.shape[-1]
|
|
|
|
transformation = transformation.swapaxes(-1, -2)
|
|
points = points @ transformation[..., :-1, :] + transformation[..., -1:, :]
|
|
|
|
return points[..., :dim].reshape(*initial_shape, dim)
|
|
|
|
|
|
def compute_camera_faces(cone_shape: "trimesh.Trimesh") -> np.ndarray:
|
|
"""Computes the faces for the camera mesh."""
|
|
faces_list = []
|
|
num_vertices_cone = len(cone_shape.vertices)
|
|
|
|
for face in cone_shape.faces:
|
|
if 0 in face:
|
|
continue
|
|
v1, v2, v3 = face
|
|
v1_offset, v2_offset, v3_offset = face + num_vertices_cone
|
|
v1_offset_2, v2_offset_2, v3_offset_2 = face + 2 * num_vertices_cone
|
|
|
|
faces_list.extend([
|
|
(v1, v2, v2_offset),
|
|
(v1, v1_offset, v3),
|
|
(v3_offset, v2, v3),
|
|
(v1, v2, v2_offset_2),
|
|
(v1, v1_offset_2, v3),
|
|
(v3_offset_2, v2, v3),
|
|
])
|
|
|
|
faces_list += [(v3, v2, v1) for v1, v2, v3 in faces_list]
|
|
return np.array(faces_list)
|
|
|
|
|
|
def compute_camera_faces_multi(cone_shape: "trimesh.Trimesh", num_shells: int) -> np.ndarray:
|
|
"""Computes faces for a camera mesh with multiple shells (for thicker frustums).
|
|
|
|
Connects each consecutive pair of vertex shells to form the frustum edges.
|
|
"""
|
|
faces_list = []
|
|
nv = len(cone_shape.vertices)
|
|
|
|
for s in range(num_shells - 1):
|
|
off_a = s * nv
|
|
off_b = (s + 1) * nv
|
|
for face in cone_shape.faces:
|
|
if 0 in face:
|
|
continue
|
|
v1, v2, v3 = face
|
|
faces_list.extend([
|
|
(v1 + off_a, v2 + off_a, v2 + off_b),
|
|
(v1 + off_a, v1 + off_b, v3 + off_a),
|
|
(v3 + off_b, v2 + off_a, v3 + off_a),
|
|
])
|
|
|
|
faces_list += [(v3, v2, v1) for v1, v2, v3 in faces_list]
|
|
return np.array(faces_list)
|
|
|
|
|
|
def segment_sky(
|
|
image_path: str,
|
|
onnx_session,
|
|
mask_filename: str
|
|
) -> np.ndarray:
|
|
"""
|
|
Segments sky from an image using an ONNX model.
|
|
|
|
Args:
|
|
image_path: Path to input image
|
|
onnx_session: ONNX runtime session with loaded model
|
|
mask_filename: Path to save the output mask
|
|
|
|
Returns:
|
|
Continuous non-sky confidence map in [0, 1]
|
|
"""
|
|
image = cv2.imread(image_path)
|
|
result_map = run_skyseg(onnx_session, _SKYSEG_INPUT_SIZE, image)
|
|
result_map_original = cv2.resize(
|
|
result_map, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_LINEAR
|
|
)
|
|
output_mask = _result_map_to_non_sky_conf(result_map_original)
|
|
|
|
os.makedirs(os.path.dirname(mask_filename), exist_ok=True)
|
|
cv2.imwrite(mask_filename, _mask_to_uint8(output_mask))
|
|
return output_mask
|
|
|
|
|
|
def run_skyseg(
|
|
onnx_session,
|
|
input_size: Tuple[int, int],
|
|
image: np.ndarray
|
|
) -> np.ndarray:
|
|
"""
|
|
Runs sky segmentation inference using ONNX model.
|
|
|
|
Args:
|
|
onnx_session: ONNX runtime session
|
|
input_size: Target size for model input (width, height)
|
|
image: Input image in BGR format
|
|
|
|
Returns:
|
|
Segmentation mask
|
|
"""
|
|
temp_image = copy.deepcopy(image)
|
|
resize_image = cv2.resize(temp_image, dsize=(input_size[0], input_size[1]))
|
|
x = cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB)
|
|
x = np.array(x, dtype=np.float32)
|
|
mean = [0.485, 0.456, 0.406]
|
|
std = [0.229, 0.224, 0.225]
|
|
x = (x / 255 - mean) / std
|
|
x = x.transpose(2, 0, 1)
|
|
x = x.reshape(-1, 3, input_size[0], input_size[1]).astype("float32")
|
|
|
|
input_name = onnx_session.get_inputs()[0].name
|
|
output_name = onnx_session.get_outputs()[0].name
|
|
onnx_result = onnx_session.run([output_name], {input_name: x})
|
|
|
|
onnx_result = np.array(onnx_result).squeeze()
|
|
min_value = np.min(onnx_result)
|
|
max_value = np.max(onnx_result)
|
|
onnx_result = (onnx_result - min_value) / (max_value - min_value)
|
|
onnx_result *= 255
|
|
return onnx_result.astype("uint8")
|
|
|
|
|
|
def download_file_from_url(url: str, filename: str):
|
|
"""Downloads a file from a URL, handling redirects."""
|
|
import requests
|
|
|
|
try:
|
|
response = requests.get(url, allow_redirects=False)
|
|
response.raise_for_status()
|
|
|
|
if response.status_code == 302:
|
|
redirect_url = response.headers["Location"]
|
|
response = requests.get(redirect_url, stream=True)
|
|
response.raise_for_status()
|
|
else:
|
|
print(f"Unexpected status code: {response.status_code}")
|
|
return
|
|
|
|
with open(filename, "wb") as f:
|
|
for chunk in response.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
print(f"Downloaded {filename} successfully.")
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Error downloading file: {e}")
|