nodes/wTechArtist-ComfyUI_VVL_VideoCamera_Advanced
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

npz reader works, reconstruction fails, alignment wo overlapping camera poses

Browse Source
This commit is contained in:
Atakan Topaloglu
2025-03-18 15:06:47 +03:00
parent f1bc88b9f0
commit 7efebaba5a
7 changed files with 2212 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -7,7 +7,7 @@ __pycache__/
**/__pycache__/
*.py[cod]
*$py.class
input_images*
# C extensions
*.so
cli_demo.py
+336
View File
@@ -0,0 +1,336 @@
#!/usr/bin/env 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 cv2
import torch
import numpy as np
import sys
import shutil
from datetime import datetime
import glob
import gc
import time
import argparse
sys.path.append("vggt/")
from visual_util import predictions_to_glb
from vggt.models.vggt import VGGT
from vggt.utils.load_fn import load_and_preprocess_images
from vggt.utils.pose_enc import pose_encoding_to_extri_intri
from vggt.utils.geometry import unproject_depth_map_to_point_map
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Initializing and loading VGGT model...")
model = VGGT.from_pretrained("facebook/VGGT-1B") # alternative loading method
# 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)
# -------------------------------------------------------------------------
# 1) Core model inference
# -------------------------------------------------------------------------
def run_model(target_dir, model) -> dict:
"""
Run the VGGT model on images in the 'target_dir/images' folder and return predictions.
"""
print(f"Processing images from {target_dir}")
# Device check
device = "cuda" if torch.cuda.is_available() else "cpu"
if not torch.cuda.is_available():
raise ValueError("CUDA is not available. Check your environment.")
# Move model to device
model = model.to(device)
model.eval()
# Load and preprocess images
image_names = glob.glob(os.path.join(target_dir, "images", "*"))
image_names = sorted(image_names)
print(f"Found {len(image_names)} images")
if len(image_names) == 0:
raise ValueError("No images found. Check your input.")
images = load_and_preprocess_images(image_names).to(device)
print(f"Preprocessed images shape: {images.shape}")
# Run inference
print("Running inference...")
with torch.no_grad():
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
predictions = model(images)
# Convert pose encoding to extrinsic and intrinsic matrices
print("Converting pose encoding to extrinsic and intrinsic matrices...")
extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images.shape[-2:])
print(f"extrinsic.shape {extrinsic.shape} \n")
print(f"extrinsic {extrinsic}")
print(f"extrinsic.shape {intrinsic.shape} \n")
print(f"intrinsic {intrinsic}")
predictions["extrinsic"] = extrinsic
predictions["intrinsic"] = intrinsic
# Convert tensors to numpy arrays
for key in predictions.keys():
if isinstance(predictions[key], torch.Tensor):
predictions[key] = predictions[key].cpu().numpy().squeeze(0) # remove batch dimension
# Generate world points from depth map
print("Computing world points from depth map...")
depth_map = predictions["depth"] # (S, H, W, 1)
world_points = unproject_depth_map_to_point_map(depth_map, predictions["extrinsic"], predictions["intrinsic"])
print(world_points.shape)
print(world_points)
predictions["world_points_from_depth"] = world_points
# Clean up
torch.cuda.empty_cache()
print(f"predictions are: {predictions}")
return predictions
# -------------------------------------------------------------------------
# 2) Handle uploaded video/images --> produce target_dir + images folder
# -------------------------------------------------------------------------
def handle_uploads(input_video, input_images):
"""
Create a new folder (target_dir/images) and copy user-provided images or extract frames from video.
Returns (target_dir, list of image paths).
"""
start_time = time.time()
gc.collect()
torch.cuda.empty_cache()
# Create a unique folder name
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
target_dir = f"input_images_{timestamp}"
target_dir_images = os.path.join(target_dir, "images")
# Clean up if folder already exists (unlikely)
if os.path.exists(target_dir):
shutil.rmtree(target_dir)
os.makedirs(target_dir)
os.makedirs(target_dir_images)
image_paths = []
# --- Handle images ---
if input_images is not None:
for file_path in input_images:
dst_path = os.path.join(target_dir_images, os.path.basename(file_path))
shutil.copy(file_path, dst_path)
image_paths.append(dst_path)
# --- Handle video ---
if input_video is not None:
video_path = input_video
vs = cv2.VideoCapture(video_path)
fps = vs.get(cv2.CAP_PROP_FPS)
frame_interval = int(fps * 1) # extract 1 frame per second
count = 0
video_frame_num = 0
while True:
gotit, frame = vs.read()
if not gotit:
break
count += 1
if count % frame_interval == 0:
image_path = os.path.join(target_dir_images, f"{video_frame_num:06}.png")
cv2.imwrite(image_path, frame)
image_paths.append(image_path)
video_frame_num += 1
# Sort images
image_paths = sorted(image_paths)
end_time = time.time()
print(f"Files copied to {target_dir_images}; took {end_time - start_time:.3f} seconds")
return target_dir, image_paths
# -------------------------------------------------------------------------
# 3) Reconstruction: performs the 3D reconstruction and saves the GLB file
# -------------------------------------------------------------------------
def gradio_demo(
target_dir,
conf_thres=3.0,
frame_filter="All",
mask_black_bg=False,
mask_white_bg=False,
show_cam=True,
mask_sky=False,
prediction_mode="Pointmap Regression",
):
"""
Perform reconstruction using the images in target_dir/images.
Returns the path to the GLB file and a log message.
"""
if not os.path.isdir(target_dir) or target_dir == "None":
return None, "No valid target directory found. Please provide valid input."
start_time = time.time()
gc.collect()
torch.cuda.empty_cache()
# Get list of image files in target directory
target_dir_images = os.path.join(target_dir, "images")
all_files = sorted(os.listdir(target_dir_images)) if os.path.isdir(target_dir_images) else []
frame_filter_choices = ["All"] + [f"{i}: {filename}" for i, filename in enumerate(all_files)]
print("Running run_model...")
with torch.no_grad():
predictions = run_model(target_dir, model)
# Save predictions to disk
prediction_save_path = os.path.join(target_dir, "predictions.npz")
np.savez(prediction_save_path, **predictions)
if frame_filter is None:
frame_filter = "All"
# Build a GLB file name
glbfile = os.path.join(
target_dir,
f"glbscene_{conf_thres}{frame_filter.replace('.', '').replace(':', '').replace(' ', '')}"
f"maskb{mask_black_bg}maskw{mask_white_bg}_cam{show_cam}_sky{mask_sky}_pred{prediction_mode.replace(' ', '')}.glb",
)
# Convert predictions to GLB
glbscene = predictions_to_glb(
predictions,
conf_thres=conf_thres,
filter_by_frames=frame_filter,
mask_black_bg=mask_black_bg,
mask_white_bg=mask_white_bg,
show_cam=show_cam,
mask_sky=mask_sky,
target_dir=target_dir,
prediction_mode=prediction_mode,
)
glbscene.export(file_obj=glbfile)
# Cleanup
del predictions
gc.collect()
torch.cuda.empty_cache()
end_time = time.time()
print(f"Total time: {end_time - start_time:.2f} seconds")
log_msg = f"Reconstruction successful. Processed {len(all_files)} frames."
return glbfile, log_msg
# -------------------------------------------------------------------------
# 4) Main function to run from command line
# -------------------------------------------------------------------------
def main():
parser = argparse.ArgumentParser(
description="VGGT 3D Reconstruction from images/video (command-line mode)"
)
parser.add_argument("--video", type=str, default=None, help="Path to input video file.")
parser.add_argument(
"--images",
type=str,
nargs="*",
default=None,
help="Paths to one or more input image files.",
)
parser.add_argument(
"--image_dir",
type=str,
default=None,
help="Path to a directory containing input images.",
)
parser.add_argument(
"--conf_thres",
type=float,
default=50.0,
help="Confidence threshold percentage (default: 50.0).",
)
parser.add_argument(
"--mask_black_bg",
action="store_true",
help="Filter black background points (flag; default: False).",
)
parser.add_argument(
"--mask_white_bg",
action="store_true",
help="Filter white background points (flag; default: False).",
)
parser.add_argument(
"--show_cam",
action="store_true",
help="Include camera visualization (flag; default: False).",
)
parser.add_argument(
"--mask_sky",
action="store_true",
help="Filter sky points (flag; default: False).",
)
parser.add_argument(
"--prediction_mode",
type=str,
choices=["Depthmap and Camera Branch", "Pointmap Branch"],
default="Depthmap and Camera Branch",
help="Prediction mode to use.",
)
args = parser.parse_args()
# Ensure that at least one type of input is provided.
if args.video is None and args.image_dir is None and (args.images is None or len(args.images) == 0):
parser.error("You must provide at least a video file, an image directory, or one or more image files.")
# If an image directory is provided, use it to build a list of image files.
image_list = None
if args.image_dir is not None:
image_list = []
# Look for common image file types.
for ext in (".png", ".jpg", ".jpeg", ".bmp"):
image_list.extend(glob.glob(os.path.join(args.image_dir, f"*{ext}")))
image_list = sorted(image_list)
if not image_list:
parser.error(f"No image files found in directory: {args.image_dir}")
elif args.images is not None and len(args.images) > 0:
image_list = args.images
# Process inputs: copy images / extract frames from video
target_dir, image_paths = handle_uploads(args.video, image_list)
print("Images for reconstruction:")
for ipath in image_paths:
print(" ", ipath)
# For command-line, we simply use "All" for frame filtering.
frame_filter = "All"
# Run reconstruction
glbfile, log_msg = gradio_demo(
target_dir,
conf_thres=args.conf_thres,
frame_filter=frame_filter,
mask_black_bg=args.mask_black_bg,
mask_white_bg=args.mask_white_bg,
show_cam=args.show_cam,
mask_sky=args.mask_sky,
prediction_mode=args.prediction_mode,
)
print(log_msg)
print("Output GLB file saved at:", glbfile)
if __name__ == "__main__":
main()
deneme.txt
+259
View File
@@ -0,0 +1,259 @@
Initializing and loading VGGT model...
Initializing and loading VGGT model...
Files copied to input_images_20250318_113826_158596/images; took 0.213 seconds
Images for reconstruction:
input_images_20250318_113826_158596/images/sandalye1.jpg
input_images_20250318_113826_158596/images/sandalye2.jpg
input_images_20250318_113826_158596/images/sandalye3.jpg
input_images_20250318_113826_158596/images/sandalye4.jpg
Running run_model...
Processing images from input_images_20250318_113826_158596
Found 4 images
Preprocessed images shape: torch.Size([4, 3, 518, 518])
Running inference...
Converting pose encoding to extrinsic and intrinsic matrices...
extrinsic.shape torch.Size([1, 4, 3, 4])
extrinsic tensor([[[[ 1.0000e+00, 3.9872e-05, 2.8053e-04, -3.3790e-05],
[-3.9849e-05, 1.0000e+00, -8.0769e-05, 1.1223e-05],
[-2.8053e-04, 8.0758e-05, 1.0000e+00, -1.1633e-04]],
[[ 8.7897e-01, -1.3935e-01, 4.5607e-01, -4.3532e-01],
[ 1.3684e-01, 9.8984e-01, 3.8704e-02, -2.3849e-02],
[-4.5683e-01, 2.8389e-02, 8.8910e-01, 1.1196e-01]],
[[ 5.4024e-01, 2.4198e-01, -8.0597e-01, 7.7377e-01],
[-2.4684e-01, 9.6120e-01, 1.2313e-01, -1.2373e-01],
[ 8.0450e-01, 1.3242e-01, 5.7901e-01, 4.3251e-01]],
[[-8.7065e-01, 1.3212e-01, -4.7383e-01, 4.1820e-01],
[-1.6372e-01, 8.3051e-01, 5.3240e-01, -5.3343e-01],
[ 4.6386e-01, 5.4111e-01, -7.0145e-01, 1.6244e+00]]]],
device='cuda:0')
extrinsic.shape torch.Size([1, 4, 3, 3])
intrinsic tensor([[[[635.5547, 0.0000, 259.0000],
[ 0.0000, 632.4077, 259.0000],
[ 0.0000, 0.0000, 1.0000]],
[[673.3128, 0.0000, 259.0000],
[ 0.0000, 668.9824, 259.0000],
[ 0.0000, 0.0000, 1.0000]],
[[647.9485, 0.0000, 259.0000],
[ 0.0000, 645.5082, 259.0000],
[ 0.0000, 0.0000, 1.0000]],
[[657.2827, 0.0000, 259.0000],
[ 0.0000, 654.9764, 259.0000],
[ 0.0000, 0.0000, 1.0000]]]], device='cuda:0')
Computing world points from depth map...
(4, 518, 518, 3)
[[[[-0.32635777 -0.3277628 0.80045778]
[-0.32256861 -0.32521258 0.79423084]
[-0.32393473 -0.32785957 0.80069504]
...
[ 0.3017545 -0.30691559 0.74978018]
[ 0.29972646 -0.30366574 0.74184234]
[ 0.29839372 -0.30114284 0.73568019]]
[[-0.31818604 -0.31832302 0.78041975]
[-0.3173857 -0.31875226 0.78147245]
[-0.31967265 -0.32229703 0.79016251]
...
[ 0.30374933 -0.30775166 0.75473664]
[ 0.30392252 -0.3067282 0.75222743]
[ 0.29965597 -0.30124899 0.73879209]]
[[-0.31972387 -0.3186213 0.78419041]
[-0.32490816 -0.32504109 0.79998947]
[-0.32480285 -0.32619915 0.80283983]
...
[ 0.30304874 -0.30585151 0.75299595]
[ 0.30336437 -0.30497799 0.75084603]
[ 0.2993832 -0.29980799 0.73811962]]
...
[[-0.33901007 0.33658178 0.83136156]
[-0.32582639 0.32474639 0.80213335]
[-0.31179238 0.31196827 0.77057688]
...
[ 0.30756927 0.30935648 0.7642391 ]
[ 0.29614456 0.29670384 0.73298793]
[ 0.27577959 0.27522561 0.67993796]]
[[-0.34688456 0.34574441 0.85066747]
[-0.35423694 0.35443977 0.87205807]
[-0.33679196 0.33829626 0.83234564]
...
[ 0.2973025 0.30019644 0.73872963]
[ 0.2851795 0.28683233 0.70584956]
[ 0.26956483 0.27007303 0.66461615]]
[[-0.34381086 0.34401416 0.84313129]
[-0.34717483 0.34872546 0.85467628]
[-0.34365569 0.34653311 0.84930432]
...
[ 0.28845039 0.29238992 0.71673515]
[ 0.29376423 0.29661765 0.72709672]
[ 0.29365914 0.29536141 0.72401809]]]
[[[-0.27946509 -0.29477718 0.69031692]
[-0.26531197 -0.29030559 0.68010908]
[-0.26449681 -0.29057286 0.68089283]
...
[ 0.30653217 -0.34329744 0.87838063]
[ 0.30861049 -0.34085725 0.87222004]
[ 0.31028803 -0.33932096 0.86838007]]
[[-0.2688719 -0.28984225 0.68177667]
[-0.25869788 -0.28679426 0.67483183]
[-0.24405615 -0.28214502 0.6641481 ]
...
[ 0.30616927 -0.34339063 0.88149869]
[ 0.30734013 -0.34308526 0.8808186 ]
[ 0.31049251 -0.33810836 0.86809097]]
[[-0.25766528 -0.28472603 0.67272506]
[-0.2610214 -0.28645689 0.67693901]
[-0.25939894 -0.28643477 0.67704571]
...
[ 0.30660648 -0.34162906 0.87984014]
[ 0.30754778 -0.34185431 0.88052861]
[ 0.31080704 -0.33663796 0.86713213]]
...
[[-0.26050822 0.39260375 0.78422529]
[-0.23273498 0.37591514 0.75874836]
[-0.19342003 0.35200265 0.72174699]
...
[ 0.3749378 0.27278311 1.05199034]
[ 0.37321451 0.28736655 1.09755724]
[ 0.37564739 0.28077125 1.07865886]]
[[-0.26893818 0.39936239 0.79274108]
[-0.27326109 0.40263491 0.79886063]
[-0.24693766 0.38678373 0.7747721 ]
...
[ 0.3755698 0.27180026 1.0451034 ]
[ 0.37722163 0.26900011 1.03775596]
[ 0.38097745 0.25502119 0.99644519]]
[[-0.23415218 0.37915543 0.75883711]
[-0.25543611 0.39302342 0.78156286]
[-0.244081 0.38648061 0.77219268]
...
[ 0.38006765 0.25102234 0.97855636]
[ 0.38073792 0.25290644 0.9853092 ]
[ 0.38054548 0.25944408 1.00614963]]]
[[[-0.23509424 -0.41133813 1.08391587]
[-0.2495936 -0.40332443 1.06414579]
[-0.24803409 -0.4034823 1.06426301]
...
[ 0.01272387 -0.23979517 0.54143505]
[ 0.01129287 -0.23923831 0.5401517 ]
[ 0.02229242 -0.24042661 0.54119231]]
[[-0.25275668 -0.40132411 1.06256198]
[-0.26313683 -0.39544312 1.04792776]
[-0.25849008 -0.39717296 1.05189651]
...
[ 0.01571067 -0.23917281 0.54219104]
[ 0.01542839 -0.23877725 0.54111911]
[-0.00361067 -0.23577585 0.53654601]]
[[-0.26319948 -0.39502293 1.05013188]
[-0.2504369 -0.40083964 1.0641572 ]
[-0.25789636 -0.39645873 1.05314673]
...
[ 0.00777232 -0.23702825 0.54087802]
[ 0.00771979 -0.23666892 0.53985876]
[-0.01090445 -0.23375314 0.53536792]]
...
[[-0.17039119 0.40640899 1.60073241]
[-0.19096508 0.3885153 1.55737619]
[-0.19465709 0.38497204 1.54676109]
...
[-0.02004279 0.38856792 0.6495685 ]
[ 0.0340838 0.42424778 0.66639606]
[ 0.03972201 0.4278182 0.66690115]]
[[-0.15641601 0.42065152 1.62900647]
[-0.18368865 0.39695608 1.57258144]
[-0.19292603 0.38866549 1.55122962]
...
[-0.03154959 0.38239768 0.64596448]
[-0.0069916 0.39855119 0.65292407]
[-0.0308108 0.38258874 0.64367914]]
[[-0.19888478 0.38662463 1.54783547]
[-0.19650977 0.38824308 1.5489444 ]
[-0.22517428 0.3632996 1.49018289]
...
[-0.0785455 0.35255071 0.63039331]
[-0.05018382 0.371291 0.63871399]
[-0.04201678 0.37657947 0.64020454]]]
[[[ 0.24408485 -0.35839285 1.02185841]
[ 0.22540079 -0.36147691 1.03593353]
[ 0.22383672 -0.36140655 1.03576627]
...
[-0.33928249 -0.29961887 0.82623133]
[-0.33998192 -0.29910583 0.82418417]
[-0.33834201 -0.29527485 0.80836729]]
[[ 0.22273443 -0.36127259 1.04010822]
[ 0.21128833 -0.36304563 1.04818302]
[ 0.21446391 -0.36209773 1.04408448]
...
[-0.33926179 -0.29841579 0.8256513 ]
[-0.33948677 -0.29722773 0.82082631]
[-0.34343427 -0.30132617 0.83782139]]
[[ 0.2181925 -0.36106792 1.04436677]
[ 0.22282258 -0.35983951 1.03905615]
[ 0.22128324 -0.35976713 1.03887214]
...
[-0.34051235 -0.29900312 0.83243455]
[-0.34124158 -0.29853749 0.83059836]
[-0.34494142 -0.30232457 0.8462184 ]]
...
[[ 0.71527314 0.81422378 1.12961781]
[ 0.63372771 0.72739113 1.16137676]
[ 0.63649916 0.73292072 1.15831854]
...
[-0.36567945 1.18487171 0.38276409]
[-0.36379322 1.25170052 0.33228194]
[-0.3619157 1.32145128 0.27955365]]
[[ 0.75121778 0.85613816 1.11595538]
[ 0.69141656 0.79312866 1.13858521]
[ 0.63250065 0.73083639 1.161045 ]
...
[-0.36201864 1.24967404 0.33818433]
[-0.36173734 1.29338213 0.30484941]
[-0.36453465 1.29021636 0.30607512]]
[[ 0.64923229 0.74657437 1.15902805]
[ 0.68824982 0.79207206 1.14100994]
[ 0.63004851 0.73044378 1.163124 ]
...
[-0.35799937 1.32063673 0.28931232]
[-0.36057686 1.32153301 0.28752395]
[-0.37615691 1.11483374 0.43864225]]]]
Building GLB scene
Using Depthmap and Camera Branch
GLB Scene built
Total time: 2.53 seconds
Reconstruction successful. Processed 4 frames.
Output GLB file saved at: input_images_20250318_113826_158596/glbscene_50.0AllmaskbFalsemaskwFalse_camFalse_skyFalse_predDepthmapandCameraBranch.glb
deneme2.txt
+1118
View File
File diff suppressed because it is too large Load Diff
merged_scene.glb
BIN
View File
Binary file not shown.
npz_reader.py
+229
View File
@@ -0,0 +1,229 @@
import numpy as np
import os
import sys
import argparse
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
def load_predictions(npz_file):
"""Load and return the predictions from an NPZ file."""
if not os.path.exists(npz_file):
print(f"Error: File {npz_file} does not exist.")
sys.exit(1)
try:
# Load the NPZ file
data = np.load(npz_file, allow_pickle=True)
# Convert to dictionary for easier handling
predictions = {key: data[key] for key in data.files}
return predictions
except Exception as e:
print(f"Error loading predictions file: {e}")
sys.exit(1)
def print_predictions_summary(predictions):
"""Print a summary of the predictions."""
print("\n======== PREDICTIONS SUMMARY ========")
print(f"Number of entries: {len(predictions)}")
for key, value in predictions.items():
if isinstance(value, np.ndarray):
print(f"\n{key}:")
print(f" Shape: {value.shape}")
print(f" Type: {value.dtype}")
# Print some statistics for numerical arrays
if np.issubdtype(value.dtype, np.number):
try:
print(f" Min: {np.min(value)}")
print(f" Max: {np.max(value)}")
print(f" Mean: {np.mean(value)}")
except:
print(" (Could not compute statistics)")
# Print a sample of the data
print(" Sample data:")
if value.size > 0:
# For multi-dimensional arrays, flatten for sample display
flat_value = value.reshape(-1)
sample_size = min(5, flat_value.size)
print(f" {flat_value[:sample_size]} ...")
else:
print(f"\n{key}: {value}")
def print_full_array(key, array):
"""Print the full contents of an array."""
print(f"\n======== FULL CONTENTS OF '{key}' ========")
print(f"Shape: {array.shape}")
print(f"Type: {array.dtype}")
# For large arrays, ask for confirmation
if array.size > 1000:
confirm = input(f"Warning: This array has {array.size} elements. Do you want to print it all? (y/n): ")
if confirm.lower() != 'y':
print("Printing canceled.")
return
# Print the array
print("\nArray contents:")
print(array)
def visualize_depth_maps(predictions):
"""Visualize the depth maps from the predictions."""
if 'depth' not in predictions:
print("No depth maps found in predictions.")
return
depth_maps = predictions['depth']
# Determine how many depth maps to display
num_frames = depth_maps.shape[0] if depth_maps.ndim > 3 else 1
# For multi-dimensional arrays, adapt accordingly
if depth_maps.ndim == 4: # (frames, height, width, 1)
num_frames = min(4, depth_maps.shape[0]) # Show up to 4 frames
fig, axes = plt.subplots(1, num_frames, figsize=(15, 5))
if num_frames == 1:
axes = [axes] # Make it iterable for consistency
for i, ax in enumerate(axes):
if i < depth_maps.shape[0]:
# Remove the channel dimension if it exists
if depth_maps.shape[-1] == 1:
depth_map = depth_maps[i, :, :, 0]
else:
depth_map = depth_maps[i]
im = ax.imshow(depth_map, cmap='viridis')
ax.set_title(f"Frame {i}")
plt.colorbar(im, ax=ax)
plt.tight_layout()
plt.show()
def visualize_point_cloud(predictions):
"""Visualize the 3D point cloud."""
point_cloud_key = None
if 'world_points_from_depth' in predictions:
point_cloud_key = 'world_points_from_depth'
if point_cloud_key is None:
print("No point cloud data found in predictions.")
return
points = predictions[point_cloud_key]
# Choose a specific frame if multiple are available
if points.ndim == 4: # (frames, height, width, 3)
frame_idx = 0
print(f"Showing point cloud for frame {frame_idx} (shape: {points.shape})")
points = points[frame_idx]
# Flatten the spatial dimensions to get a list of 3D points
points_flat = points.reshape(-1, 3)
# Create a scatter plot in 3D
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')
# Downsample if there are too many points
max_points = 10000
if len(points_flat) > max_points:
# Random sampling
idx = np.random.choice(len(points_flat), max_points, replace=False)
points_flat = points_flat[idx]
# Plot the points
ax.scatter(
points_flat[:, 0],
points_flat[:, 1],
points_flat[:, 2],
c='b',
marker='.',
alpha=0.5,
s=1
)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_title('3D Point Cloud')
# Equal aspect ratio
ax.set_box_aspect([1, 1, 1])
plt.tight_layout()
plt.show()
def main():
parser = argparse.ArgumentParser(description='Read and analyze predictions.npz file from VGGT model.')
parser.add_argument('npz_file', help='Path to the predictions.npz file')
parser.add_argument('--key', help='Specific key to print from the npz file')
parser.add_argument('--visualize', choices=['depth', 'points', 'all'],
help='Visualize specific data: depth maps or 3D points')
args = parser.parse_args()
# Load predictions
predictions = load_predictions(args.npz_file)
# If a specific key is provided, print it and exit
if args.key:
if args.key in predictions:
print_full_array(args.key, predictions[args.key])
return
else:
print(f"Error: Key '{args.key}' not found.")
print("Available keys:")
for key in predictions.keys():
print(f"- {key}")
return
# Print summary
print_predictions_summary(predictions)
# Handle visualization request
if args.visualize:
if args.visualize == 'depth' or args.visualize == 'all':
visualize_depth_maps(predictions)
if args.visualize == 'points' or args.visualize == 'all':
visualize_point_cloud(predictions)
else:
# Interactive mode
while True:
print("\n======== OPTIONS ========")
print("1: List all available keys")
print("2: Print contents of a specific key")
print("3: Visualize depth maps")
print("4: Visualize 3D point cloud")
print("5: Exit")
choice = input("\nEnter your choice (1-5): ")
if choice == '1':
print("\nAvailable keys:")
for key in sorted(predictions.keys()):
print(f"- {key}")
elif choice == '2':
key = input("Enter the key name: ")
if key in predictions:
print_full_array(key, predictions[key])
else:
print(f"Error: Key '{key}' not found.")
print("Available keys:")
for k in sorted(predictions.keys()):
print(f"- {k}")
elif choice == '3':
visualize_depth_maps(predictions)
elif choice == '4':
visualize_point_cloud(predictions)
elif choice == '5':
break
else:
print("Invalid choice. Please try again.")
if __name__ == "__main__":
main()
reconstruct.py
+269
View File
@@ -0,0 +1,269 @@
import numpy as np
import open3d as o3d
import copy
import trimesh
import os
import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation
def merge_point_clouds(prediction_sets, voxel_size=0.02):
"""
Register and merge multiple point clouds from overlapping image sets
Args:
prediction_sets: List of prediction dictionaries containing point clouds and camera parameters
voxel_size: Resolution for downsampling and registration
Returns:
merged_scene: Final trimesh.Scene with merged point cloud
"""
# 1. Extract point clouds and camera positions from each prediction set
point_clouds = []
camera_positions = []
for i, predictions in enumerate(prediction_sets):
# Extract point cloud data
if "world_points" in predictions:
points = predictions["world_points"]
else:
points = predictions["world_points_from_depth"]
# Get confidence values if available
if "world_points_conf" in predictions:
conf = predictions["world_points_conf"]
else:
conf = np.ones_like(points[..., 0])
# Get RGB images
images = predictions["images"]
# Extract camera matrices
extrinsics = predictions["extrinsic"]
# Create colored point cloud
vertices, colors = extract_colored_point_cloud(points, images, conf, threshold=50.0)
# Create Open3D point cloud
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(vertices)
pcd.colors = o3d.utility.Vector3dVector(colors / 255.0) # Normalize colors to [0,1]
# Downsample to speed up registration
pcd_down = pcd.voxel_down_sample(voxel_size)
# Extract camera centers from extrinsic matrices
cam_centers = extract_camera_centers(extrinsics)
point_clouds.append(pcd_down)
camera_positions.append(cam_centers)
print(f"Point cloud {i}: {len(pcd_down.points)} points, {len(cam_centers)} cameras")
# 2. Register point clouds using both camera positions and point features
transformations = register_point_clouds(point_clouds, camera_positions)
# 3. Transform and merge point clouds
merged_cloud = o3d.geometry.PointCloud()
for i, pcd in enumerate(point_clouds):
# Apply transformation
if i > 0: # First cloud is reference
pcd_transformed = copy.deepcopy(pcd)
pcd_transformed.transform(transformations[i-1])
merged_cloud += pcd_transformed
else:
merged_cloud += pcd
# 4. Clean and optimize final point cloud
print("Cleaning merged point cloud...")
merged_cloud = merged_cloud.voxel_down_sample(voxel_size)
merged_cloud, _ = merged_cloud.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
# 5. Create trimesh scene for output
merged_scene = trimesh.Scene()
# Convert merged point cloud to trimesh
vertices = np.asarray(merged_cloud.points)
colors = (np.asarray(merged_cloud.colors) * 255).astype(np.uint8)
merged_trimesh = trimesh.PointCloud(vertices=vertices, colors=colors)
merged_scene.add_geometry(merged_trimesh)
return merged_scene
def extract_colored_point_cloud(points, images, confidence, threshold=50.0):
"""Extract points and colors based on confidence threshold"""
S, H, W, _ = points.shape
vertices_3d = 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: # NHWC format
colors_rgb = images
colors_rgb = (colors_rgb.reshape(-1, 3) * 255).astype(np.uint8)
conf = confidence.reshape(-1)
# Apply confidence threshold
if threshold > 0:
conf_threshold = np.percentile(conf, threshold)
mask = (conf >= conf_threshold) & (conf > 1e-5)
vertices_3d = vertices_3d[mask]
colors_rgb = colors_rgb[mask]
return vertices_3d, colors_rgb
def extract_camera_centers(extrinsics):
"""Extract camera centers from extrinsic matrices"""
centers = []
for ext in extrinsics:
# Create 4x4 transformation matrix
T = np.eye(4)
T[:3, :4] = ext
# Get camera center in world coordinates (-R^T * t)
R = T[:3, :3]
t = T[:3, 3]
center = -np.dot(R.T, t)
centers.append(center)
return np.array(centers)
def register_point_clouds(point_clouds, camera_positions):
"""
Register point clouds using both camera positions and ICP
Returns:
List of transformation matrices to align each point cloud to the first one
"""
if len(point_clouds) <= 1:
return []
transformations = []
target_cloud = point_clouds[0]
target_cameras = camera_positions[0]
# Process each subsequent point cloud
for i in range(1, len(point_clouds)):
source_cloud = point_clouds[i]
source_cameras = camera_positions[i]
# 1. Initial alignment using camera positions (if sequential)
initial_transformation = estimate_transformation_from_cameras(source_cameras, target_cameras)
# Apply initial transformation
source_aligned = copy.deepcopy(source_cloud)
source_aligned.transform(initial_transformation)
# 2. Fine registration using Colored ICP
# Prepare for colored ICP
source_aligned.estimate_normals()
target_cloud.estimate_normals()
# Run Colored ICP (slower but more accurate for colored point clouds)
print(f"Running colored ICP for point cloud {i}...")
reg_p2p = o3d.pipelines.registration.registration_colored_icp(
source_aligned, target_cloud, 0.05,
initial_transformation,
o3d.pipelines.registration.TransformationEstimationForColoredICP(),
o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200)
)
# Combine transformations
final_transformation = reg_p2p.transformation @ initial_transformation
transformations.append(final_transformation)
print(f"Point cloud {i} registered with fitness: {reg_p2p.fitness}, RMSE: {reg_p2p.inlier_rmse}")
return transformations
def estimate_transformation_from_cameras(source_cameras, target_cameras):
"""
Estimate transformation between point clouds using camera positions
using a simplified version of Umeyama algorithm
"""
if len(source_cameras) < 3 or len(target_cameras) < 3:
# Not enough camera positions for reliable estimation
return np.eye(4)
# Find corresponding cameras by proximity
# This works if the camera positions are in sequential order
s_cam = source_cameras
t_cam = target_cameras
# If the number of cameras differs, use the minimum length
min_len = min(len(s_cam), len(t_cam))
# Try to find best alignment using first and last few cameras
source_pts = np.vstack([s_cam[:min(3, min_len)], s_cam[-min(3, min_len):]])
target_pts = np.vstack([t_cam[:min(3, min_len)], t_cam[-min(3, min_len):]])
# Calculate centroids
source_centroid = np.mean(source_pts, axis=0)
target_centroid = np.mean(target_pts, axis=0)
# Center the points
source_centered = source_pts - source_centroid
target_centered = target_pts - target_centroid
# Compute rotation using singular value decomposition
H = np.dot(source_centered.T, target_centered)
U, _, Vt = np.linalg.svd(H)
R = np.dot(Vt.T, U.T)
# Ensure proper rotation matrix (det=1)
if np.linalg.det(R) < 0:
Vt[-1, :] *= -1
R = np.dot(Vt.T, U.T)
# Translation
t = target_centroid - np.dot(R, source_centroid)
# Create transformation matrix
transformation = np.eye(4)
transformation[:3, :3] = R
transformation[:3, 3] = t
return transformation
def main():
"""Example usage"""
import argparse
parser = argparse.ArgumentParser(description='Merge multiple point clouds from VGGT predictions')
parser.add_argument('--prediction_dirs', type=str, nargs='+', required=True,
help='Directories containing prediction.npz files')
parser.add_argument('--output', type=str, default='merged_scene.glb',
help='Output GLB file path')
parser.add_argument('--voxel_size', type=float, default=0.02,
help='Voxel size for downsampling')
args = parser.parse_args()
# Load predictions
prediction_sets = []
for dir_path in args.prediction_dirs:
npz_path = os.path.join(dir_path, 'predictions.npz')
if not os.path.exists(npz_path):
print(f"Warning: No predictions.npz found in {dir_path}")
continue
predictions = dict(np.load(npz_path, allow_pickle=True))
prediction_sets.append(predictions)
# Merge point clouds
merged_scene = merge_point_clouds(prediction_sets, args.voxel_size)
# Export merged scene
merged_scene.export(file_obj=args.output)
print(f"Merged scene saved to {args.output}")
if __name__ == "__main__":
main()