415 lines
14 KiB
Python
415 lines
14 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.
|
||
|
||
import os
|
||
import glob
|
||
import time
|
||
import threading
|
||
import argparse
|
||
from typing import List, Optional
|
||
|
||
import numpy as np
|
||
import torch
|
||
from tqdm.auto import tqdm
|
||
import viser
|
||
import viser.transforms as viser_tf
|
||
import cv2
|
||
|
||
|
||
try:
|
||
import onnxruntime
|
||
except ImportError:
|
||
print("onnxruntime not found. Sky segmentation may not work.")
|
||
|
||
from visual_util import segment_sky, download_file_from_url
|
||
from vggt.models.vggt import VGGT
|
||
from vggt.utils.load_fn import load_and_preprocess_images
|
||
from vggt.utils.geometry import closed_form_inverse_se3, unproject_depth_map_to_point_map
|
||
from vggt.utils.pose_enc import pose_encoding_to_extri_intri
|
||
|
||
|
||
def viser_wrapper(
|
||
pred_dict: dict,
|
||
port: int = 8080,
|
||
init_conf_threshold: float = 50.0, # represents percentage (e.g., 50 means filter lowest 50%)
|
||
use_point_map: bool = False,
|
||
background_mode: bool = False,
|
||
mask_sky: bool = False,
|
||
image_folder: str = None,
|
||
):
|
||
"""
|
||
Visualize predicted 3D points and camera poses with viser.
|
||
|
||
Args:
|
||
pred_dict (dict):
|
||
{
|
||
"images": (S, 3, H, W) - Input images,
|
||
"world_points": (S, H, W, 3),
|
||
"world_points_conf": (S, H, W),
|
||
"depth": (S, H, W, 1),
|
||
"depth_conf": (S, H, W),
|
||
"extrinsic": (S, 3, 4),
|
||
"intrinsic": (S, 3, 3),
|
||
}
|
||
port (int): Port number for the viser server.
|
||
init_conf_threshold (float): Initial percentage of low-confidence points to filter out.
|
||
use_point_map (bool): Whether to visualize world_points or use depth-based points.
|
||
background_mode (bool): Whether to run the server in background thread.
|
||
mask_sky (bool): Whether to apply sky segmentation to filter out sky points.
|
||
image_folder (str): Path to the folder containing input images.
|
||
"""
|
||
print(f"Starting viser server on port {port}")
|
||
|
||
server = viser.ViserServer(host="0.0.0.0", port=port)
|
||
server.gui.configure_theme(titlebar_content=None, control_layout="collapsible")
|
||
|
||
# Unpack prediction dict
|
||
images = pred_dict["images"] # (S, 3, H, W)
|
||
world_points_map = pred_dict["world_points"] # (S, H, W, 3)
|
||
conf_map = pred_dict["world_points_conf"] # (S, H, W)
|
||
|
||
depth_map = pred_dict["depth"] # (S, H, W, 1)
|
||
depth_conf = pred_dict["depth_conf"] # (S, H, W)
|
||
|
||
extrinsics_cam = pred_dict["extrinsic"] # (S, 3, 4)
|
||
intrinsics_cam = pred_dict["intrinsic"] # (S, 3, 3)
|
||
|
||
# Compute world points from depth if not using the precomputed point map
|
||
if not use_point_map:
|
||
world_points = unproject_depth_map_to_point_map(depth_map, extrinsics_cam, intrinsics_cam)
|
||
conf = depth_conf
|
||
else:
|
||
world_points = world_points_map
|
||
conf = conf_map
|
||
|
||
# Apply sky segmentation if enabled
|
||
if mask_sky and image_folder is not None:
|
||
conf = apply_sky_segmentation(conf, image_folder)
|
||
|
||
# Convert images from (S, 3, H, W) to (S, H, W, 3)
|
||
# Then flatten everything for the point cloud
|
||
colors = images.transpose(0, 2, 3, 1) # now (S, H, W, 3)
|
||
S, H, W, _ = world_points.shape
|
||
|
||
# Flatten
|
||
points = world_points.reshape(-1, 3)
|
||
colors_flat = (colors.reshape(-1, 3) * 255).astype(np.uint8)
|
||
conf_flat = conf.reshape(-1)
|
||
|
||
cam_to_world_mat = closed_form_inverse_se3(extrinsics_cam) # shape (S, 4, 4) typically
|
||
# For convenience, we store only (3,4) portion
|
||
cam_to_world = cam_to_world_mat[:, :3, :]
|
||
|
||
# Compute scene center and recenter
|
||
scene_center = np.mean(points, axis=0)
|
||
points_centered = points - scene_center
|
||
cam_to_world[..., -1] -= scene_center
|
||
|
||
# Store frame indices so we can filter by frame
|
||
frame_indices = np.repeat(np.arange(S), H * W)
|
||
|
||
# Build the viser GUI
|
||
gui_show_frames = server.gui.add_checkbox(
|
||
"Show Cameras",
|
||
initial_value=True,
|
||
)
|
||
|
||
# Now the slider represents percentage of points to filter out
|
||
gui_points_conf = server.gui.add_slider(
|
||
"Confidence Percent",
|
||
min=0,
|
||
max=100,
|
||
step=0.1,
|
||
initial_value=init_conf_threshold,
|
||
)
|
||
|
||
gui_frame_selector = server.gui.add_dropdown(
|
||
"Show Points from Frames",
|
||
options=["All"] + [str(i) for i in range(S)],
|
||
initial_value="All",
|
||
)
|
||
|
||
# Create the main point cloud handle
|
||
# Compute the threshold value as the given percentile
|
||
init_threshold_val = np.percentile(conf_flat, init_conf_threshold)
|
||
init_conf_mask = (conf_flat >= init_threshold_val) & (conf_flat > 0.1)
|
||
point_cloud = server.scene.add_point_cloud(
|
||
name="viser_pcd",
|
||
points=points_centered[init_conf_mask],
|
||
colors=colors_flat[init_conf_mask],
|
||
point_size=0.001,
|
||
point_shape="circle",
|
||
)
|
||
|
||
# We will store references to frames & frustums so we can toggle visibility
|
||
frames: List[viser.FrameHandle] = []
|
||
frustums: List[viser.CameraFrustumHandle] = []
|
||
|
||
def visualize_frames(extrinsics: np.ndarray, images_: np.ndarray) -> None:
|
||
"""
|
||
Add camera frames and frustums to the scene.
|
||
extrinsics: (S, 3, 4)
|
||
images_: (S, 3, H, W)
|
||
"""
|
||
# Clear any existing frames or frustums
|
||
for f in frames:
|
||
f.remove()
|
||
frames.clear()
|
||
for fr in frustums:
|
||
fr.remove()
|
||
frustums.clear()
|
||
|
||
# Optionally attach a callback that sets the viewpoint to the chosen camera
|
||
def attach_callback(frustum: viser.CameraFrustumHandle, frame: viser.FrameHandle) -> None:
|
||
@frustum.on_click
|
||
def _(_) -> None:
|
||
for client in server.get_clients().values():
|
||
client.camera.wxyz = frame.wxyz
|
||
client.camera.position = frame.position
|
||
|
||
img_ids = range(S)
|
||
for img_id in tqdm(img_ids):
|
||
cam2world_3x4 = extrinsics[img_id]
|
||
T_world_camera = viser_tf.SE3.from_matrix(cam2world_3x4)
|
||
|
||
# Add a small frame axis
|
||
frame_axis = server.scene.add_frame(
|
||
f"frame_{img_id}",
|
||
wxyz=T_world_camera.rotation().wxyz,
|
||
position=T_world_camera.translation(),
|
||
axes_length=0.05,
|
||
axes_radius=0.002,
|
||
origin_radius=0.002,
|
||
)
|
||
frames.append(frame_axis)
|
||
|
||
# Convert the image for the frustum
|
||
img = images_[img_id] # shape (3, H, W)
|
||
img = (img.transpose(1, 2, 0) * 255).astype(np.uint8)
|
||
h, w = img.shape[:2]
|
||
|
||
# If you want correct FOV from intrinsics, do something like:
|
||
# fx = intrinsics_cam[img_id, 0, 0]
|
||
# fov = 2 * np.arctan2(h/2, fx)
|
||
# For demonstration, we pick a simple approximate FOV:
|
||
fy = 1.1 * h
|
||
fov = 2 * np.arctan2(h / 2, fy)
|
||
|
||
# Add the frustum
|
||
frustum_cam = server.scene.add_camera_frustum(
|
||
f"frame_{img_id}/frustum",
|
||
fov=fov,
|
||
aspect=w / h,
|
||
scale=0.05,
|
||
image=img,
|
||
line_width=1.0,
|
||
)
|
||
frustums.append(frustum_cam)
|
||
attach_callback(frustum_cam, frame_axis)
|
||
|
||
def update_point_cloud() -> None:
|
||
"""Update the point cloud based on current GUI selections."""
|
||
# Here we compute the threshold value based on the current percentage
|
||
current_percentage = gui_points_conf.value
|
||
threshold_val = np.percentile(conf_flat, current_percentage)
|
||
|
||
print(f"Threshold absolute value: {threshold_val}, percentage: {current_percentage}%")
|
||
|
||
conf_mask = (conf_flat >= threshold_val) & (conf_flat > 1e-5)
|
||
|
||
if gui_frame_selector.value == "All":
|
||
frame_mask = np.ones_like(conf_mask, dtype=bool)
|
||
else:
|
||
selected_idx = int(gui_frame_selector.value)
|
||
frame_mask = frame_indices == selected_idx
|
||
|
||
combined_mask = conf_mask & frame_mask
|
||
point_cloud.points = points_centered[combined_mask]
|
||
point_cloud.colors = colors_flat[combined_mask]
|
||
|
||
@gui_points_conf.on_update
|
||
def _(_) -> None:
|
||
update_point_cloud()
|
||
|
||
@gui_frame_selector.on_update
|
||
def _(_) -> None:
|
||
update_point_cloud()
|
||
|
||
@gui_show_frames.on_update
|
||
def _(_) -> None:
|
||
"""Toggle visibility of camera frames and frustums."""
|
||
for f in frames:
|
||
f.visible = gui_show_frames.value
|
||
for fr in frustums:
|
||
fr.visible = gui_show_frames.value
|
||
|
||
# Add the camera frames to the scene
|
||
visualize_frames(cam_to_world, images)
|
||
|
||
print("Starting viser server...")
|
||
# If background_mode is True, spawn a daemon thread so the main thread can continue.
|
||
if background_mode:
|
||
|
||
def server_loop():
|
||
while True:
|
||
time.sleep(0.001)
|
||
|
||
thread = threading.Thread(target=server_loop, daemon=True)
|
||
thread.start()
|
||
else:
|
||
while True:
|
||
time.sleep(0.01)
|
||
|
||
return server
|
||
|
||
|
||
# Helper functions for sky segmentation
|
||
|
||
|
||
def apply_sky_segmentation(conf: np.ndarray, image_folder: str) -> np.ndarray:
|
||
"""
|
||
Apply sky segmentation to confidence scores.
|
||
|
||
Args:
|
||
conf (np.ndarray): Confidence scores with shape (S, H, W)
|
||
image_folder (str): Path to the folder containing input images
|
||
|
||
Returns:
|
||
np.ndarray: Updated confidence scores with sky regions masked out
|
||
"""
|
||
S, H, W = conf.shape
|
||
sky_masks_dir = image_folder.rstrip("/") + "_sky_masks"
|
||
os.makedirs(sky_masks_dir, exist_ok=True)
|
||
|
||
# Download skyseg.onnx if it doesn't exist
|
||
if not os.path.exists("skyseg.onnx"):
|
||
print("Downloading skyseg.onnx...")
|
||
download_file_from_url("https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx", "skyseg.onnx")
|
||
|
||
skyseg_session = onnxruntime.InferenceSession("skyseg.onnx")
|
||
image_files = sorted(glob.glob(os.path.join(image_folder, "*")))
|
||
sky_mask_list = []
|
||
|
||
print("Generating sky masks...")
|
||
for i, image_path in enumerate(tqdm(image_files[:S])): # Limit to the number of images in the batch
|
||
image_name = os.path.basename(image_path)
|
||
mask_filepath = os.path.join(sky_masks_dir, image_name)
|
||
|
||
if os.path.exists(mask_filepath):
|
||
sky_mask = cv2.imread(mask_filepath, cv2.IMREAD_GRAYSCALE)
|
||
else:
|
||
sky_mask = segment_sky(image_path, skyseg_session, mask_filepath)
|
||
|
||
# Resize mask to match H×W if needed
|
||
if sky_mask.shape[0] != H or sky_mask.shape[1] != W:
|
||
sky_mask = cv2.resize(sky_mask, (W, H))
|
||
|
||
sky_mask_list.append(sky_mask)
|
||
|
||
# Convert list to numpy array with shape S×H×W
|
||
sky_mask_array = np.array(sky_mask_list)
|
||
# Apply sky mask to confidence scores
|
||
sky_mask_binary = (sky_mask_array > 0.1).astype(np.float32)
|
||
conf = conf * sky_mask_binary
|
||
|
||
print("Sky segmentation applied successfully")
|
||
return conf
|
||
|
||
|
||
parser = argparse.ArgumentParser(description="VGGT demo with viser for 3D visualization")
|
||
parser.add_argument(
|
||
"--image_folder", type=str, default="examples/kitchen/images/", help="Path to folder containing images"
|
||
)
|
||
parser.add_argument("--use_point_map", action="store_true", help="Use point map instead of depth-based points")
|
||
parser.add_argument("--background_mode", action="store_true", help="Run the viser server in background mode")
|
||
parser.add_argument("--port", type=int, default=8080, help="Port number for the viser server")
|
||
parser.add_argument(
|
||
"--conf_threshold", type=float, default=25.0, help="Initial percentage of low-confidence points to filter out"
|
||
)
|
||
parser.add_argument("--mask_sky", action="store_true", help="Apply sky segmentation to filter out sky points")
|
||
|
||
|
||
def main():
|
||
"""
|
||
Main function for the VGGT demo with viser for 3D visualization.
|
||
|
||
This function:
|
||
1. Loads the VGGT model
|
||
2. Processes input images from the specified folder
|
||
3. Runs inference to generate 3D points and camera poses
|
||
4. Optionally applies sky segmentation to filter out sky points
|
||
5. Visualizes the results using viser
|
||
|
||
Command-line arguments:
|
||
--image_folder: Path to folder containing input images
|
||
--use_point_map: Use point map instead of depth-based points
|
||
--background_mode: Run the viser server in background mode
|
||
--port: Port number for the viser server
|
||
--conf_threshold: Initial percentage of low-confidence points to filter out
|
||
--mask_sky: Apply sky segmentation to filter out sky points
|
||
"""
|
||
args = parser.parse_args()
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
print(f"Using device: {device}")
|
||
|
||
print("Initializing and loading VGGT model...")
|
||
# model = VGGT.from_pretrained("facebook/VGGT-1B")
|
||
|
||
model = VGGT()
|
||
_URL = "https://huggingface.co/facebook/VGGT-1B/resolve/main/model.pt"
|
||
model.load_state_dict(torch.hub.load_state_dict_from_url(_URL))
|
||
|
||
model.eval()
|
||
model = model.to(device)
|
||
|
||
# Use the provided image folder path
|
||
print(f"Loading images from {args.image_folder}...")
|
||
image_names = glob.glob(os.path.join(args.image_folder, "*"))
|
||
print(f"Found {len(image_names)} images")
|
||
|
||
images = load_and_preprocess_images(image_names).to(device)
|
||
print(f"Preprocessed images shape: {images.shape}")
|
||
|
||
print("Running inference...")
|
||
with torch.no_grad():
|
||
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
|
||
predictions = model(images)
|
||
|
||
print("Converting pose encoding to extrinsic and intrinsic matrices...")
|
||
extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images.shape[-2:])
|
||
predictions["extrinsic"] = extrinsic
|
||
predictions["intrinsic"] = intrinsic
|
||
|
||
print("Processing model outputs...")
|
||
for key in predictions.keys():
|
||
if isinstance(predictions[key], torch.Tensor):
|
||
predictions[key] = predictions[key].cpu().numpy().squeeze(0) # remove batch dimension and convert to numpy
|
||
|
||
if args.use_point_map:
|
||
print("Visualizing 3D points from point map")
|
||
else:
|
||
print("Visualizing 3D points by unprojecting depth map by cameras")
|
||
|
||
if args.mask_sky:
|
||
print("Sky segmentation enabled - will filter out sky points")
|
||
|
||
print("Starting viser visualization...")
|
||
|
||
viser_server = viser_wrapper(
|
||
predictions,
|
||
port=args.port,
|
||
init_conf_threshold=args.conf_threshold,
|
||
use_point_map=args.use_point_map,
|
||
background_mode=args.background_mode,
|
||
mask_sky=args.mask_sky,
|
||
image_folder=args.image_folder,
|
||
)
|
||
print("Visualization complete")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|