Added colmap writer for point cloud and camera initialization
This commit is contained in:
Atakan Topaloglu
-336
@@ -1,336 +0,0 @@
|
||||
#!/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()
|
||||
@@ -0,0 +1,291 @@
|
||||
import numpy as np
|
||||
import os
|
||||
import torch
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
import glob
|
||||
import argparse
|
||||
from PIL import Image
|
||||
|
||||
def rotation_matrix_to_quaternion(rot_matrix):
|
||||
"""Convert rotation matrix to quaternion (w, x, y, z)"""
|
||||
rotation = R.from_matrix(rot_matrix)
|
||||
return rotation.as_quat()[[3, 0, 1, 2]] # Reorder from (x, y, z, w) to (w, x, y, z)
|
||||
|
||||
def write_cameras_txt(intrinsics, image_shapes, output_dir):
|
||||
"""Write cameras.txt file"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
with open(os.path.join(output_dir, "cameras.txt"), "w") as f:
|
||||
f.write("# Camera list with one line of data per camera:\n")
|
||||
f.write("# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n")
|
||||
f.write(f"# Number of cameras: {len(intrinsics)}\n")
|
||||
|
||||
for i, (intrin, shape) in enumerate(zip(intrinsics, image_shapes)):
|
||||
height, width = shape
|
||||
|
||||
fx, fy = intrin[0, 0], intrin[1, 1]
|
||||
cx, cy = intrin[0, 2], intrin[1, 2]
|
||||
|
||||
# Using PINHOLE model (fx, fy, cx, cy)
|
||||
camera_id = i + 1
|
||||
f.write(f"{camera_id} PINHOLE {width} {height} {fx} {fy} {cx} {cy}\n")
|
||||
|
||||
def extract_sparse_point_cloud(point_maps, images, confidences, sample_step=20, confidence_threshold=0.0):
|
||||
"""Extract a sparse point cloud with filtering by confidence threshold"""
|
||||
points = []
|
||||
point_colors = []
|
||||
point_origins = []
|
||||
|
||||
for img_idx, (point_map, img, conf_map) in enumerate(zip(point_maps, images, confidences)):
|
||||
h, w, _ = point_map.shape
|
||||
for y in range(0, h, sample_step):
|
||||
for x in range(0, w, sample_step):
|
||||
conf = conf_map[y, x]
|
||||
|
||||
if conf < confidence_threshold:
|
||||
continue
|
||||
|
||||
point3d = point_map[y, x]
|
||||
|
||||
if np.linalg.norm(point3d) < 0.001 or np.linalg.norm(point3d) > 1000:
|
||||
continue
|
||||
|
||||
color = img[y, x]
|
||||
|
||||
points.append(point3d)
|
||||
point_colors.append(color)
|
||||
point_origins.append((img_idx, y, x))
|
||||
|
||||
if len(points) == 0:
|
||||
print("Warning: No points passed the confidence threshold. Try lowering the threshold.")
|
||||
return np.array([[0, 0, 0]]), np.array([[0.5, 0.5, 0.5]]), [(0, 0, 0)]
|
||||
|
||||
return np.array(points), np.array(point_colors), point_origins
|
||||
|
||||
def filter_points_by_distance(points, colors, origins, min_distance=0.05):
|
||||
"""Filter points that are too close to each other"""
|
||||
return points, colors, origins
|
||||
|
||||
def compute_point_visibility(points3d, extrinsics, intrinsics, image_shapes):
|
||||
"""Compute which images can see each 3D point"""
|
||||
num_points = len(points3d)
|
||||
num_images = len(extrinsics)
|
||||
|
||||
# For each image, store visible 2D projections
|
||||
image_points = [[] for _ in range(num_images)]
|
||||
|
||||
# For each 3D point, store which images see it
|
||||
point_tracks = [[] for _ in range(num_points)]
|
||||
|
||||
for point_idx, point3d in enumerate(points3d):
|
||||
for img_idx, (extrinsic, intrinsic, shape) in enumerate(zip(extrinsics, intrinsics, image_shapes)):
|
||||
height, width = shape
|
||||
|
||||
point_cam = np.dot(extrinsic[:3, :3], point3d) + extrinsic[:3, 3]
|
||||
|
||||
# Check if point is in front of camera
|
||||
if point_cam[2] <= 0:
|
||||
continue
|
||||
|
||||
# Project to image plane
|
||||
point_homo = np.dot(intrinsic, point_cam / point_cam[2])
|
||||
x, y = point_homo[0], point_homo[1]
|
||||
|
||||
# Check if within image boundaries
|
||||
if 0 <= x < width and 0 <= y < height:
|
||||
image_points[img_idx].append((x, y, point_idx))
|
||||
point_tracks[point_idx].append((img_idx + 1, len(image_points[img_idx]) - 1))
|
||||
|
||||
return image_points, point_tracks
|
||||
|
||||
def write_images_txt(extrinsics, image_points, image_names, output_dir):
|
||||
"""Write images.txt file"""
|
||||
with open(os.path.join(output_dir, "images.txt"), "w") as f:
|
||||
f.write("# Image list with two lines of data per image:\n")
|
||||
f.write("# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n")
|
||||
f.write("# POINTS2D[] as (X, Y, POINT3D_ID)\n")
|
||||
f.write(f"# Number of images: {len(extrinsics)}\n")
|
||||
|
||||
for img_idx, (extrinsic, points) in enumerate(zip(extrinsics, image_points)):
|
||||
|
||||
rot_matrix = extrinsic[:3, :3]
|
||||
trans_vector = extrinsic[:3, 3]
|
||||
|
||||
rot_c2w = rot_matrix.T
|
||||
trans_c2w = -np.dot(rot_c2w, trans_vector) # -R^T * t
|
||||
|
||||
quat_c2w = rotation_matrix_to_quaternion(rot_c2w)
|
||||
|
||||
image_id = img_idx + 1
|
||||
camera_id = img_idx + 1
|
||||
image_name = os.path.basename(image_names[img_idx])
|
||||
|
||||
|
||||
f.write(f"{image_id} {quat_c2w[0]:.9f} {quat_c2w[1]:.9f} {quat_c2w[2]:.9f} {quat_c2w[3]:.9f} {trans_c2w[0]:.9f} {trans_c2w[1]:.9f} {trans_c2w[2]:.9f} {camera_id} {image_name}\n")
|
||||
|
||||
|
||||
points2d_line = " ".join([f"{x:.6f} {y:.6f} {point3d_id}" for x, y, point3d_id in points])
|
||||
f.write(f"{points2d_line}\n")
|
||||
|
||||
def write_points3d_txt(points3d, point_colors, point_tracks, output_dir):
|
||||
"""Write points3D.txt file"""
|
||||
with open(os.path.join(output_dir, "points3D.txt"), "w") as f:
|
||||
f.write("# 3D point list with one line of data per point:\n")
|
||||
f.write("# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n")
|
||||
f.write(f"# Number of points: {len(points3d)}\n")
|
||||
|
||||
valid_points = 0
|
||||
for point_idx, (point3d, color, track) in enumerate(zip(points3d, point_colors, point_tracks)):
|
||||
# Skip points that aren't visible in any image
|
||||
if not track:
|
||||
continue
|
||||
|
||||
r, g, b = [int(c * 255) for c in color]
|
||||
error = 1.0 # Placeholder for reprojection error
|
||||
|
||||
track_str = " ".join([f"{img_id} {point2d_idx}" for img_id, point2d_idx in track])
|
||||
|
||||
f.write(f"{point_idx} {point3d[0]:.9f} {point3d[1]:.9f} {point3d[2]:.9f} {r} {g} {b} {error:.9f} {track_str}\n")
|
||||
valid_points += 1
|
||||
|
||||
print(f"Written {valid_points} valid points (visible in at least one image)")
|
||||
|
||||
def convert_to_colmap_format(extrinsics, intrinsics, point_maps, point_confidences, image_paths, output_dir, image_shapes, images, sample_step=20, confidence_threshold=0.0):
|
||||
"""Convert model estimates to COLMAP format files with confidence filtering"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
print(f"Extracting colored sparse point cloud (confidence threshold: {confidence_threshold})...")
|
||||
points3d, point_colors, point_origins = extract_sparse_point_cloud(
|
||||
point_maps, images, point_confidences, sample_step, confidence_threshold
|
||||
)
|
||||
print(f"Extracted {len(points3d)} colored points that meet the confidence threshold")
|
||||
|
||||
print("Computing point visibility...")
|
||||
image_points, point_tracks = compute_point_visibility(
|
||||
points3d, extrinsics, intrinsics, image_shapes
|
||||
)
|
||||
|
||||
print("Writing COLMAP format files...")
|
||||
write_cameras_txt(intrinsics, image_shapes, output_dir)
|
||||
write_images_txt(extrinsics, image_points, image_paths, output_dir)
|
||||
write_points3d_txt(points3d, point_colors, point_tracks, output_dir)
|
||||
|
||||
print(f"COLMAP format files written to {output_dir}")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Convert VGGT estimates to COLMAP format")
|
||||
parser.add_argument("image_dir", help="Directory containing input images")
|
||||
parser.add_argument("output_dir", help="Directory to write COLMAP files")
|
||||
parser.add_argument("--sample-step", type=int, default=10, help="Point sampling step (lower = denser point cloud)")
|
||||
parser.add_argument("--confidence-threshold", type=float, default=0.0,
|
||||
help="Minimum confidence value for points (0.0-1.0, higher = stricter filtering)")
|
||||
parser.add_argument("--use-point-branch", action="store_true",
|
||||
help="Use point map branch instead of depth map for 3D points")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
image_paths = sorted(glob.glob(os.path.join(args.image_dir, "*.jpg"))) + \
|
||||
sorted(glob.glob(os.path.join(args.image_dir, "*.png")))
|
||||
|
||||
if not image_paths:
|
||||
raise ValueError(f"No images found in {args.image_dir}")
|
||||
|
||||
print(f"Found {len(image_paths)} images")
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float16
|
||||
|
||||
print(f"Using device: {device}")
|
||||
print("Loading VGGT model...")
|
||||
|
||||
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
|
||||
|
||||
model = VGGT.from_pretrained("facebook/VGGT-1B").to(device)
|
||||
|
||||
print("Processing images...")
|
||||
images_tensor = load_and_preprocess_images(image_paths).to(device)
|
||||
|
||||
original_images = []
|
||||
image_shapes = []
|
||||
for img_path in image_paths:
|
||||
with Image.open(img_path) as img:
|
||||
# Convert to RGB and numpy array for coloring purposes
|
||||
np_img = np.array(img.convert('RGB')) / 255.0 # Normalize to 0-1
|
||||
original_images.append(np_img)
|
||||
# Store image shape
|
||||
image_shapes.append((img.height, img.width))
|
||||
|
||||
with torch.no_grad():
|
||||
with torch.cuda.amp.autocast(dtype=dtype):
|
||||
images_batch = images_tensor[None] # add batch dimension
|
||||
aggregated_tokens_list, ps_idx = model.aggregator(images_batch)
|
||||
|
||||
print("Estimating camera parameters...")
|
||||
pose_enc = model.camera_head(aggregated_tokens_list)[-1]
|
||||
extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc, images_batch.shape[-2:])
|
||||
|
||||
print(f"Branch selection: {'Point Branch' if args.use_point_branch else 'Depth Branch'}")
|
||||
|
||||
if args.use_point_branch:
|
||||
print("Using point map branch...")
|
||||
point_map, point_conf = model.point_head(aggregated_tokens_list, images_batch, ps_idx)
|
||||
else:
|
||||
print("Using depth map branch (default, more accurate)...")
|
||||
depth_map, depth_conf = model.depth_head(aggregated_tokens_list, images_batch, ps_idx)
|
||||
|
||||
print("Unprojecting to 3D points...")
|
||||
point_map = unproject_depth_map_to_point_map(
|
||||
depth_map.squeeze(0), extrinsic.squeeze(0), intrinsic.squeeze(0)
|
||||
)
|
||||
point_conf = depth_conf
|
||||
|
||||
# Convert to numpy arrays if they're not already
|
||||
if isinstance(extrinsic, torch.Tensor):
|
||||
extrinsic_np = extrinsic.squeeze(0).cpu().numpy()
|
||||
else:
|
||||
extrinsic_np = extrinsic
|
||||
|
||||
if isinstance(intrinsic, torch.Tensor):
|
||||
intrinsic_np = intrinsic.squeeze(0).cpu().numpy()
|
||||
else:
|
||||
intrinsic_np = intrinsic
|
||||
|
||||
if isinstance(point_map, torch.Tensor):
|
||||
point_map_np = point_map.cpu().numpy()
|
||||
else:
|
||||
point_map_np = point_map
|
||||
|
||||
if isinstance(point_conf, torch.Tensor):
|
||||
point_conf_np = point_conf.squeeze(0).cpu().numpy()
|
||||
else:
|
||||
point_conf_np = point_conf
|
||||
|
||||
# # Normalize confidence values to 0-1 range
|
||||
# if np.max(point_conf_np) > 1.0:
|
||||
# point_conf_np = point_conf_np / np.max(point_conf_np)
|
||||
|
||||
print(f"Confidence range: {np.min(point_conf_np):.4f} - {np.max(point_conf_np):.4f}")
|
||||
|
||||
# Resize original images to match point map if needed
|
||||
resized_images = []
|
||||
for i, img in enumerate(original_images):
|
||||
h, w, _ = point_map_np[i].shape
|
||||
if img.shape[0] != h or img.shape[1] != w:
|
||||
|
||||
from skimage.transform import resize
|
||||
resized_img = resize(img, (h, w, 3), anti_aliasing=True)
|
||||
resized_images.append(resized_img)
|
||||
else:
|
||||
resized_images.append(img)
|
||||
|
||||
convert_to_colmap_format(
|
||||
extrinsic_np, intrinsic_np, point_map_np, point_conf_np, image_paths,
|
||||
args.output_dir, image_shapes, resized_images,
|
||||
args.sample_step, args.confidence_threshold
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,127 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
import os
|
||||
import numpy as np
|
||||
import json
|
||||
import csv
|
||||
import glob
|
||||
import argparse
|
||||
|
||||
def extract_camera_transforms(extrinsics):
|
||||
"""Extract camera positions and orientations from extrinsic matrices"""
|
||||
cameras = []
|
||||
for ext in extrinsics:
|
||||
# Create 4x4 transformation matrix
|
||||
T = np.eye(4)
|
||||
T[:3, :4] = ext
|
||||
|
||||
# Extract rotation matrix and position
|
||||
R = T[:3, :3]
|
||||
t = T[:3, 3]
|
||||
|
||||
# Get camera center in world coordinates (-R^T * t)
|
||||
center = -np.dot(R.T, t)
|
||||
|
||||
# Store camera data
|
||||
camera_data = {
|
||||
'position': center.tolist(),
|
||||
'rotation_matrix': R.tolist(),
|
||||
'transform_matrix': T.tolist()
|
||||
}
|
||||
cameras.append(camera_data)
|
||||
|
||||
return cameras
|
||||
|
||||
def export_cameras_csv(cameras, output_path):
|
||||
"""Export camera data to CSV format"""
|
||||
with open(output_path, 'w', newline='') as f:
|
||||
writer = csv.writer(f)
|
||||
# Write header
|
||||
writer.writerow(['frame', 'pos_x', 'pos_y', 'pos_z',
|
||||
'r11', 'r12', 'r13',
|
||||
'r21', 'r22', 'r23',
|
||||
'r31', 'r32', 'r33'])
|
||||
|
||||
# Write data for each camera
|
||||
for i, cam in enumerate(cameras):
|
||||
pos = cam['position']
|
||||
rot = cam['rotation_matrix']
|
||||
row = [i,
|
||||
pos[0], pos[1], pos[2],
|
||||
rot[0][0], rot[0][1], rot[0][2],
|
||||
rot[1][0], rot[1][1], rot[1][2],
|
||||
rot[2][0], rot[2][1], rot[2][2]]
|
||||
writer.writerow(row)
|
||||
|
||||
def export_cameras_json(cameras, output_path):
|
||||
"""Export camera data to JSON format"""
|
||||
camera_data = {
|
||||
'cameras': [
|
||||
{
|
||||
'frame': i,
|
||||
'position': cam['position'],
|
||||
'rotation_matrix': cam['rotation_matrix'],
|
||||
'transform_matrix': cam['transform_matrix']
|
||||
}
|
||||
for i, cam in enumerate(cameras)
|
||||
]
|
||||
}
|
||||
|
||||
with open(output_path, 'w') as f:
|
||||
json.dump(camera_data, f, indent=2)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Export camera poses from inference results to CSV and JSON"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--npz_files", type=str, required=True,
|
||||
help="Path to directory containing NPZ files or glob pattern to match them"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_prefix", type=str, default="camera_poses0",
|
||||
help="Prefix for output files (will create .csv and .json files)"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get all NPZ files
|
||||
if os.path.isdir(args.npz_files):
|
||||
npz_paths = glob.glob(os.path.join(args.npz_files, "**", "*.npz"), recursive=True)
|
||||
else:
|
||||
npz_paths = glob.glob(args.npz_files)
|
||||
|
||||
if not npz_paths:
|
||||
raise ValueError(f"No NPZ files found at {args.npz_files}")
|
||||
|
||||
print(f"Found {len(npz_paths)} NPZ files")
|
||||
|
||||
# Collect all camera transforms
|
||||
all_cameras = []
|
||||
|
||||
for npz_path in npz_paths:
|
||||
print(f"Processing {npz_path}")
|
||||
predictions = dict(np.load(npz_path, allow_pickle=True))
|
||||
|
||||
if "extrinsic" in predictions:
|
||||
extrinsics = predictions["extrinsic"]
|
||||
cameras = extract_camera_transforms(extrinsics)
|
||||
all_cameras.extend(cameras)
|
||||
|
||||
if not all_cameras:
|
||||
raise ValueError("No camera transforms found in the provided NPZ files")
|
||||
|
||||
print(f"Extracted {len(all_cameras)} camera transforms")
|
||||
|
||||
# Export to CSV and JSON
|
||||
csv_path = f"{args.output_prefix}.csv"
|
||||
json_path = f"{args.output_prefix}.json"
|
||||
|
||||
export_cameras_csv(all_cameras, csv_path)
|
||||
export_cameras_json(all_cameras, json_path)
|
||||
|
||||
print(f"Exports complete!")
|
||||
print(f"CSV file saved to: {csv_path}")
|
||||
print(f"JSON file saved to: {json_path}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
-229
@@ -1,229 +0,0 @@
|
||||
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()
|
||||
-792
@@ -1,792 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# Integrated VGGT + Point Cloud Registration with Overlapping Image Batches
|
||||
|
||||
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
|
||||
import copy
|
||||
import trimesh
|
||||
import open3d as o3d
|
||||
from tqdm import tqdm
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
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
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Overlapping Batch Processing Functions
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
|
||||
def improve_camera_based_initialization(source_cameras, target_cameras, source_images, target_images):
|
||||
"""
|
||||
Create a better initial transformation using camera poses to prevent bottom-to-top matching issues
|
||||
"""
|
||||
print("Finding optimal camera alignment...")
|
||||
|
||||
# Find overlapping camera frames if any
|
||||
source_filenames = [os.path.splitext(os.path.basename(path))[0] for path in source_images]
|
||||
target_filenames = [os.path.splitext(os.path.basename(path))[0] for path in target_images]
|
||||
common_files = set(source_filenames) & set(target_filenames)
|
||||
|
||||
# Get indices in each batch
|
||||
source_overlap_idx = [source_filenames.index(f) for f in common_files if f in source_filenames]
|
||||
target_overlap_idx = [target_filenames.index(f) for f in common_files if f in target_filenames]
|
||||
|
||||
# Try multiple alignment strategies and choose the best one
|
||||
alignment_options = []
|
||||
|
||||
# Option 1: Use overlapping cameras if available
|
||||
if len(source_overlap_idx) >= 3:
|
||||
print(f"Strategy 1: Using {len(source_overlap_idx)} explicitly overlapping cameras")
|
||||
source_pts = source_cameras[source_overlap_idx]
|
||||
target_pts = target_cameras[target_overlap_idx]
|
||||
|
||||
# Calculate transformation
|
||||
transform, error = compute_rigid_transform(source_pts, target_pts)
|
||||
alignment_options.append(("Overlapping Cameras", transform, error, source_overlap_idx, target_overlap_idx))
|
||||
|
||||
# Option 2: Try start-to-end alignment (last cameras of target to first cameras of source)
|
||||
num_match = min(8, min(len(source_cameras), len(target_cameras)))
|
||||
source_start_idx = list(range(num_match))
|
||||
target_end_idx = list(range(len(target_cameras) - num_match, len(target_cameras)))
|
||||
|
||||
print(f"Strategy 2: Matching last {num_match} cameras of target to first {num_match} cameras of source")
|
||||
source_pts = source_cameras[:num_match]
|
||||
target_pts = target_cameras[-num_match:]
|
||||
transform, error = compute_rigid_transform(source_pts, target_pts)
|
||||
alignment_options.append(("Start-to-End", transform, error, source_start_idx, target_end_idx))
|
||||
|
||||
# Option 3: Try end-to-start alignment (first cameras of target to last cameras of source)
|
||||
source_end_idx = list(range(len(source_cameras) - num_match, len(source_cameras)))
|
||||
target_start_idx = list(range(num_match))
|
||||
|
||||
print(f"Strategy 3: Matching first {num_match} cameras of target to last {num_match} cameras of source")
|
||||
source_pts = source_cameras[-num_match:]
|
||||
target_pts = target_cameras[:num_match]
|
||||
transform, error = compute_rigid_transform(source_pts, target_pts)
|
||||
alignment_options.append(("End-to-Start", transform, error, source_end_idx, target_start_idx))
|
||||
|
||||
# Choose the best alignment based on error
|
||||
alignment_options.sort(key=lambda x: x[2])
|
||||
best_alignment = alignment_options[0]
|
||||
|
||||
print(f"Selected alignment strategy: {best_alignment[0]} with error: {best_alignment[2]:.6f}")
|
||||
return best_alignment[1], best_alignment[3], best_alignment[4]
|
||||
|
||||
|
||||
def compute_rigid_transform(source_pts, target_pts):
|
||||
"""Calculate optimal rigid transform between point sets and return error metric"""
|
||||
# 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 SVD
|
||||
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
|
||||
|
||||
# Calculate error metric (average distance after transform)
|
||||
transformed_source = np.dot(source_centered, R.T) + target_centroid
|
||||
error = np.mean(np.linalg.norm(transformed_source - target_pts, axis=1))
|
||||
|
||||
return transformation, error
|
||||
|
||||
|
||||
def visualize_registration_overlap(source_pcd, target_pcd, source_cameras, target_cameras,
|
||||
source_overlap_idx, target_overlap_idx, transformation, output_path):
|
||||
"""
|
||||
Create a visualization GLB showing both point clouds and their camera positions
|
||||
"""
|
||||
# Transform source to align with target
|
||||
source_transformed = copy.deepcopy(source_pcd)
|
||||
source_transformed.transform(transformation)
|
||||
|
||||
# Color the point clouds differently
|
||||
source_points = np.asarray(source_transformed.points)
|
||||
source_colors = np.ones((len(source_points), 3)) * [1.0, 0.7, 0.0] # Yellow for source
|
||||
|
||||
target_points = np.asarray(target_pcd.points)
|
||||
target_colors = np.ones((len(target_points), 3)) * [0.0, 0.65, 0.93] # Blue for target
|
||||
|
||||
# Combine the point clouds
|
||||
all_points = np.vstack((source_points, target_points))
|
||||
all_colors = np.vstack((source_colors, target_colors))
|
||||
all_colors = (all_colors * 255).astype(np.uint8)
|
||||
|
||||
# Create trimesh scene
|
||||
scene = trimesh.Scene()
|
||||
|
||||
# Add the combined point cloud
|
||||
cloud = trimesh.PointCloud(vertices=all_points, colors=all_colors)
|
||||
scene.add_geometry(cloud)
|
||||
|
||||
# Transform source cameras
|
||||
transformed_source_cameras = []
|
||||
for cam in source_cameras:
|
||||
# Convert to homogeneous coordinates
|
||||
cam_h = np.ones(4)
|
||||
cam_h[:3] = cam
|
||||
# Transform
|
||||
cam_transformed = transformation @ cam_h
|
||||
transformed_source_cameras.append(cam_transformed[:3])
|
||||
|
||||
# Add camera positions as spheres
|
||||
for i, cam_pos in enumerate(transformed_source_cameras):
|
||||
# Red for overlapping source cameras, orange for others
|
||||
color = [255, 0, 0, 255] if i in source_overlap_idx else [255, 165, 0, 255]
|
||||
sphere = trimesh.primitives.Sphere(radius=0.05, center=cam_pos)
|
||||
sphere.visual.face_colors = color
|
||||
scene.add_geometry(sphere)
|
||||
|
||||
for i, cam_pos in enumerate(target_cameras):
|
||||
# Green for overlapping target cameras, cyan for others
|
||||
color = [0, 255, 0, 255] if i in target_overlap_idx else [0, 255, 255, 255]
|
||||
sphere = trimesh.primitives.Sphere(radius=0.05, center=cam_pos)
|
||||
sphere.visual.face_colors = color
|
||||
scene.add_geometry(sphere)
|
||||
|
||||
# Add camera trajectory lines using cylinders instead of paths
|
||||
# Source trajectory (red)
|
||||
if len(transformed_source_cameras) >= 2:
|
||||
for i in range(len(transformed_source_cameras) - 1):
|
||||
p1 = transformed_source_cameras[i]
|
||||
p2 = transformed_source_cameras[i+1]
|
||||
|
||||
# Create a cylinder between consecutive camera positions
|
||||
direction = p2 - p1
|
||||
length = np.linalg.norm(direction)
|
||||
if length > 0: # Avoid zero-length cylinders
|
||||
# Create a cylinder
|
||||
radius = 0.01 # Cylinder radius
|
||||
center = (p1 + p2) / 2
|
||||
|
||||
# Calculate rotation to align cylinder with direction
|
||||
direction_norm = direction / length
|
||||
z_axis = np.array([0, 0, 1])
|
||||
|
||||
if np.allclose(direction_norm, z_axis) or np.allclose(direction_norm, -z_axis):
|
||||
rotation = np.eye(3)
|
||||
else:
|
||||
# Find rotation from z-axis to direction
|
||||
cross = np.cross(z_axis, direction_norm)
|
||||
sin_angle = np.linalg.norm(cross)
|
||||
cos_angle = np.dot(z_axis, direction_norm)
|
||||
|
||||
# Skew-symmetric cross-product matrix
|
||||
K = np.array([[0, -cross[2], cross[1]],
|
||||
[cross[2], 0, -cross[0]],
|
||||
[-cross[1], cross[0], 0]])
|
||||
|
||||
# Rodrigues formula
|
||||
if sin_angle > 0:
|
||||
rotation = np.eye(3) + K + K @ K * (1 - cos_angle) / (sin_angle ** 2)
|
||||
else:
|
||||
rotation = np.eye(3)
|
||||
|
||||
# Create cylinder
|
||||
transform = np.eye(4)
|
||||
transform[:3, :3] = rotation
|
||||
transform[:3, 3] = center
|
||||
|
||||
cylinder = trimesh.creation.cylinder(radius=radius, height=length)
|
||||
cylinder.apply_transform(transform)
|
||||
cylinder.visual.face_colors = [255, 0, 0, 100] # Red, semi-transparent
|
||||
scene.add_geometry(cylinder)
|
||||
|
||||
# Target trajectory (green)
|
||||
if len(target_cameras) >= 2:
|
||||
for i in range(len(target_cameras) - 1):
|
||||
p1 = target_cameras[i]
|
||||
p2 = target_cameras[i+1]
|
||||
|
||||
direction = p2 - p1
|
||||
length = np.linalg.norm(direction)
|
||||
if length > 0:
|
||||
radius = 0.01
|
||||
center = (p1 + p2) / 2
|
||||
|
||||
direction_norm = direction / length
|
||||
z_axis = np.array([0, 0, 1])
|
||||
|
||||
if np.allclose(direction_norm, z_axis) or np.allclose(direction_norm, -z_axis):
|
||||
rotation = np.eye(3)
|
||||
else:
|
||||
cross = np.cross(z_axis, direction_norm)
|
||||
sin_angle = np.linalg.norm(cross)
|
||||
cos_angle = np.dot(z_axis, direction_norm)
|
||||
|
||||
K = np.array([[0, -cross[2], cross[1]],
|
||||
[cross[2], 0, -cross[0]],
|
||||
[-cross[1], cross[0], 0]])
|
||||
|
||||
if sin_angle > 0:
|
||||
rotation = np.eye(3) + K + K @ K * (1 - cos_angle) / (sin_angle ** 2)
|
||||
else:
|
||||
rotation = np.eye(3)
|
||||
|
||||
transform = np.eye(4)
|
||||
transform[:3, :3] = rotation
|
||||
transform[:3, 3] = center
|
||||
|
||||
cylinder = trimesh.creation.cylinder(radius=radius, height=length)
|
||||
cylinder.apply_transform(transform)
|
||||
cylinder.visual.face_colors = [0, 255, 0, 100] # Green, semi-transparent
|
||||
scene.add_geometry(cylinder)
|
||||
|
||||
# Export the scene
|
||||
scene.export(output_path)
|
||||
print(f"Registration visualization saved to {output_path}")
|
||||
|
||||
return scene
|
||||
|
||||
def create_overlapping_batches(image_paths, batch_size=16, overlap_size=8):
|
||||
"""
|
||||
Create overlapping batches of sequential images
|
||||
|
||||
Args:
|
||||
image_paths: List of paths to input images (already sorted)
|
||||
batch_size: Number of images per batch
|
||||
overlap_size: Number of overlapping images between consecutive batches
|
||||
|
||||
Returns:
|
||||
List of image path batches
|
||||
"""
|
||||
if batch_size <= overlap_size:
|
||||
raise ValueError(f"Batch size ({batch_size}) must be larger than overlap size ({overlap_size})")
|
||||
|
||||
if len(image_paths) <= batch_size:
|
||||
return [image_paths]
|
||||
|
||||
batches = []
|
||||
step_size = batch_size - overlap_size
|
||||
|
||||
for i in range(0, len(image_paths), step_size):
|
||||
batch_end = min(i + batch_size, len(image_paths))
|
||||
batch = image_paths[i:batch_end]
|
||||
batches.append(batch)
|
||||
|
||||
# If we can't form a full batch anymore, stop
|
||||
if batch_end == len(image_paths):
|
||||
break
|
||||
|
||||
# Ensure the last batch is at least half the batch_size to avoid tiny batches
|
||||
if len(batches) > 1 and len(batches[-1]) < batch_size // 2:
|
||||
# Extend the last batch by including more images from the previous batch
|
||||
batches[-2].extend(batches[-1])
|
||||
batches.pop()
|
||||
|
||||
return batches
|
||||
|
||||
|
||||
def setup_model():
|
||||
"""Initialize and return the VGGT model"""
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
print(f"Initializing and loading VGGT model on {device}...")
|
||||
|
||||
model = VGGT.from_pretrained("facebook/VGGT-1B")
|
||||
model.eval()
|
||||
model = model.to(device)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def process_image_batch(image_batch, model, batch_id, output_dir):
|
||||
"""
|
||||
Process a batch of images and save results
|
||||
|
||||
Args:
|
||||
image_batch: List of paths to images in this batch
|
||||
model: VGGT model instance
|
||||
batch_id: ID for this batch
|
||||
output_dir: Directory to save batch results
|
||||
|
||||
Returns:
|
||||
Dictionary with batch predictions and paths
|
||||
"""
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
# Create batch directory
|
||||
batch_dir = os.path.join(output_dir, f"batch_{batch_id:03d}")
|
||||
batch_images_dir = os.path.join(batch_dir, "images")
|
||||
os.makedirs(batch_images_dir, exist_ok=True)
|
||||
|
||||
# Copy images to batch directory for reference
|
||||
batch_image_paths = []
|
||||
for i, path in enumerate(image_batch):
|
||||
dst_path = os.path.join(batch_images_dir, f"{i:04d}{os.path.splitext(path)[1]}")
|
||||
shutil.copy(path, dst_path)
|
||||
batch_image_paths.append(dst_path)
|
||||
|
||||
# Process images with VGGT
|
||||
print(f"Processing batch {batch_id} with {len(batch_image_paths)} images...")
|
||||
|
||||
# Load and preprocess images
|
||||
images = load_and_preprocess_images(batch_image_paths).to(device)
|
||||
print(f"Preprocessed images shape: {images.shape}")
|
||||
|
||||
# Run inference
|
||||
with torch.no_grad():
|
||||
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
|
||||
predictions = model(images)
|
||||
|
||||
# Convert 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
|
||||
|
||||
# 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
|
||||
depth_map = predictions["depth"]
|
||||
world_points = unproject_depth_map_to_point_map(depth_map, predictions["extrinsic"], predictions["intrinsic"])
|
||||
predictions["world_points_from_depth"] = world_points
|
||||
|
||||
# Save predictions to NPZ file
|
||||
npz_path = os.path.join(batch_dir, "predictions.npz")
|
||||
np.savez(npz_path, **predictions)
|
||||
|
||||
# Generate individual batch GLB
|
||||
glb_path = os.path.join(batch_dir, "scene.glb")
|
||||
glbscene = predictions_to_glb(
|
||||
predictions,
|
||||
conf_thres=50.0,
|
||||
filter_by_frames="All",
|
||||
mask_black_bg=True,
|
||||
mask_white_bg=True,
|
||||
show_cam=True,
|
||||
mask_sky=False,
|
||||
target_dir=batch_dir,
|
||||
prediction_mode="Depthmap and Camera Branch"
|
||||
)
|
||||
glbscene.export(file_obj=glb_path)
|
||||
|
||||
result = {
|
||||
"predictions": predictions,
|
||||
"batch_dir": batch_dir,
|
||||
"glb_path": glb_path,
|
||||
"npz_path": npz_path,
|
||||
"image_paths": batch_image_paths,
|
||||
"batch_id": batch_id
|
||||
}
|
||||
|
||||
# Clean CUDA memory
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Point Cloud Extraction and Registration
|
||||
# -------------------------------------------------------------------------
|
||||
def extract_point_cloud_from_predictions(predictions, conf_threshold=50.0, voxel_size=0.02, mask_bg=True):
|
||||
"""Extract Open3D point cloud from predictions with filtering"""
|
||||
# Get points and confidence
|
||||
if "world_points" in predictions:
|
||||
points = predictions["world_points"]
|
||||
conf = predictions.get("world_points_conf", np.ones_like(points[..., 0]))
|
||||
else:
|
||||
points = predictions["world_points_from_depth"]
|
||||
conf = predictions.get("depth_conf", np.ones_like(points[..., 0]))
|
||||
|
||||
# Get images
|
||||
images = predictions["images"]
|
||||
|
||||
# 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
|
||||
|
||||
# Reshape data
|
||||
vertices_3d = points.reshape(-1, 3)
|
||||
colors_rgb = (colors_rgb.reshape(-1, 3) * 255).astype(np.uint8)
|
||||
conf = conf.reshape(-1)
|
||||
|
||||
# Apply confidence threshold
|
||||
if conf_threshold > 0:
|
||||
conf_threshold_value = np.percentile(conf, conf_threshold)
|
||||
conf_mask = (conf >= conf_threshold_value) & (conf > 1e-5)
|
||||
else:
|
||||
conf_mask = np.ones_like(conf, dtype=bool)
|
||||
|
||||
# Filter background
|
||||
if mask_bg:
|
||||
# Filter black background
|
||||
black_bg_mask = colors_rgb.sum(axis=1) >= 16
|
||||
# Filter white background
|
||||
white_bg_mask = ~((colors_rgb[:, 0] > 240) & (colors_rgb[:, 1] > 240) & (colors_rgb[:, 2] > 240))
|
||||
conf_mask = conf_mask & black_bg_mask & white_bg_mask
|
||||
|
||||
# Apply masks
|
||||
vertices_3d = vertices_3d[conf_mask]
|
||||
colors_rgb = colors_rgb[conf_mask]
|
||||
|
||||
# Create Open3D point cloud
|
||||
pcd = o3d.geometry.PointCloud()
|
||||
pcd.points = o3d.utility.Vector3dVector(vertices_3d)
|
||||
pcd.colors = o3d.utility.Vector3dVector(colors_rgb / 255.0)
|
||||
|
||||
# Downsample to make processing more efficient
|
||||
pcd = pcd.voxel_down_sample(voxel_size)
|
||||
|
||||
# Get camera positions
|
||||
extrinsics = predictions["extrinsic"]
|
||||
camera_centers = extract_camera_centers(extrinsics)
|
||||
|
||||
return pcd, camera_centers
|
||||
|
||||
|
||||
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 find_overlapping_cameras(batch1_cameras, batch2_cameras, batch1_images, batch2_images):
|
||||
"""
|
||||
Identify overlapping cameras between batches based on image filenames
|
||||
|
||||
Returns:
|
||||
Indices of overlapping cameras in each batch
|
||||
"""
|
||||
# Extract base filenames without directories and extensions
|
||||
batch1_filenames = [os.path.splitext(os.path.basename(path))[0] for path in batch1_images]
|
||||
batch2_filenames = [os.path.splitext(os.path.basename(path))[0] for path in batch2_images]
|
||||
|
||||
# Find common filenames
|
||||
common_files = set(batch1_filenames) & set(batch2_filenames)
|
||||
|
||||
# Get indices in each batch
|
||||
batch1_indices = [batch1_filenames.index(f) for f in common_files if f in batch1_filenames]
|
||||
batch2_indices = [batch2_filenames.index(f) for f in common_files if f in batch2_filenames]
|
||||
|
||||
return batch1_indices, batch2_indices
|
||||
|
||||
|
||||
def register_consecutive_point_clouds(source_pcd, target_pcd, source_cameras, target_cameras,
|
||||
source_images, target_images, output_dir=None, batch_id=None):
|
||||
"""
|
||||
Register two point clouds from consecutive batches with improved initialization
|
||||
"""
|
||||
print("Registering consecutive point clouds...")
|
||||
|
||||
# 1. Better initialization using camera poses
|
||||
initial_transformation, source_overlap_idx, target_overlap_idx = improve_camera_based_initialization(
|
||||
source_cameras, target_cameras, source_images, target_images
|
||||
)
|
||||
|
||||
# 2. Visualize pre-registration alignment
|
||||
if output_dir is not None and batch_id is not None:
|
||||
pre_reg_path = os.path.join(output_dir, f"pre_registration_batch_{batch_id}.glb")
|
||||
visualize_registration_overlap(
|
||||
source_pcd, target_pcd, source_cameras, target_cameras,
|
||||
source_overlap_idx, target_overlap_idx, initial_transformation, pre_reg_path
|
||||
)
|
||||
|
||||
# 3. Apply initial transformation and proceed with ICP refinement
|
||||
source_aligned = copy.deepcopy(source_pcd)
|
||||
source_aligned.transform(initial_transformation)
|
||||
|
||||
# Prepare for ICP
|
||||
source_aligned.estimate_normals()
|
||||
target_pcd.estimate_normals()
|
||||
|
||||
# Try different refinement methods
|
||||
try:
|
||||
print("Running colored ICP for fine alignment...")
|
||||
reg_p2p = o3d.pipelines.registration.registration_colored_icp(
|
||||
source_aligned, target_pcd, 0.05,
|
||||
np.eye(4), # Start from current alignment
|
||||
o3d.pipelines.registration.TransformationEstimationForColoredICP(),
|
||||
o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=50)
|
||||
)
|
||||
print(f"ICP fitness: {reg_p2p.fitness}, RMSE: {reg_p2p.inlier_rmse}")
|
||||
|
||||
# Only use ICP refinement if it significantly improves the alignment
|
||||
if reg_p2p.fitness > 0.3: # Adjust threshold as needed
|
||||
final_transformation = reg_p2p.transformation @ initial_transformation
|
||||
else:
|
||||
print("ICP did not improve alignment significantly, using initial transformation")
|
||||
final_transformation = initial_transformation
|
||||
except Exception as e:
|
||||
print(f"ICP refinement failed: {str(e)}")
|
||||
final_transformation = initial_transformation
|
||||
|
||||
# 4. Visualize post-registration alignment
|
||||
if output_dir is not None and batch_id is not None:
|
||||
post_reg_path = os.path.join(output_dir, f"post_registration_batch_{batch_id}.glb")
|
||||
visualize_registration_overlap(
|
||||
source_pcd, target_pcd, source_cameras, target_cameras,
|
||||
source_overlap_idx, target_overlap_idx, final_transformation, post_reg_path
|
||||
)
|
||||
|
||||
return final_transformation
|
||||
|
||||
|
||||
def register_all_point_clouds(batch_results, output_dir, voxel_size=0.02):
|
||||
"""
|
||||
Register all point clouds with better initialization and visualization
|
||||
"""
|
||||
point_clouds = []
|
||||
camera_positions = []
|
||||
transformations = []
|
||||
|
||||
# Extract point cloud from each batch
|
||||
for batch_result in batch_results:
|
||||
predictions = batch_result["predictions"]
|
||||
pcd, cameras = extract_point_cloud_from_predictions(
|
||||
predictions, conf_threshold=30.0, voxel_size=voxel_size, mask_bg=True
|
||||
)
|
||||
point_clouds.append(pcd)
|
||||
camera_positions.append(cameras)
|
||||
|
||||
# Register consecutive point clouds
|
||||
for i in range(1, len(point_clouds)):
|
||||
source_pcd = point_clouds[i]
|
||||
target_pcd = point_clouds[i-1]
|
||||
source_cameras = camera_positions[i]
|
||||
target_cameras = camera_positions[i-1]
|
||||
source_images = batch_results[i]["image_paths"]
|
||||
target_images = batch_results[i-1]["image_paths"]
|
||||
|
||||
# Pass both output_dir and batch_id
|
||||
transformation = register_consecutive_point_clouds(
|
||||
source_pcd, target_pcd, source_cameras, target_cameras,
|
||||
source_images, target_images, output_dir, i
|
||||
)
|
||||
|
||||
transformations.append(transformation)
|
||||
|
||||
# Apply transformations to align all point clouds
|
||||
aligned_point_clouds = [point_clouds[0]] # First cloud is reference
|
||||
|
||||
for i in range(1, len(point_clouds)):
|
||||
# Calculate cumulative transformation
|
||||
cumulative_transform = np.eye(4)
|
||||
for j in range(i-1, -1, -1):
|
||||
cumulative_transform = transformations[j] @ cumulative_transform
|
||||
|
||||
# Apply transformation
|
||||
transformed_pcd = copy.deepcopy(point_clouds[i])
|
||||
transformed_pcd.transform(cumulative_transform)
|
||||
aligned_point_clouds.append(transformed_pcd)
|
||||
|
||||
return aligned_point_clouds, transformations
|
||||
|
||||
|
||||
def merge_aligned_point_clouds(aligned_point_clouds, voxel_size=0.02):
|
||||
"""
|
||||
Merge aligned point clouds into a single consistent model
|
||||
"""
|
||||
if not aligned_point_clouds:
|
||||
return None
|
||||
|
||||
# Merge all aligned clouds
|
||||
merged_cloud = o3d.geometry.PointCloud()
|
||||
for pcd in aligned_point_clouds:
|
||||
merged_cloud += pcd
|
||||
|
||||
# Clean and optimize the merged point cloud
|
||||
print("Cleaning and optimizing 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)
|
||||
|
||||
return merged_cloud
|
||||
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Main Reconstruction Pipeline
|
||||
# -------------------------------------------------------------------------
|
||||
def reconstruct_from_images_with_overlap(image_dir, output_path, batch_size=8, overlap_size=2,
|
||||
conf_threshold=50.0, voxel_size=0.02):
|
||||
"""
|
||||
Complete end-to-end reconstruction pipeline with overlapping image batches
|
||||
|
||||
Args:
|
||||
image_dir: Directory containing input images
|
||||
output_path: Path to save the output merged GLB file
|
||||
batch_size: Number of images per batch
|
||||
overlap_size: Number of overlapping images between batches
|
||||
conf_threshold: Confidence threshold for filtering points (0-100)
|
||||
voxel_size: Voxel size for downsampling point clouds
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Create output directory for all results
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
output_dir = f"reconstructions/reconstruction_{timestamp}"
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Get all images in directory
|
||||
image_paths = []
|
||||
for ext in [".png", ".jpg", ".jpeg", ".bmp"]:
|
||||
image_paths.extend(glob.glob(os.path.join(image_dir, f"*{ext}")))
|
||||
image_paths.extend(glob.glob(os.path.join(image_dir, f"*{ext.upper()}")))
|
||||
|
||||
# Sort images to ensure sequential processing
|
||||
image_paths = sorted(image_paths)
|
||||
|
||||
if not image_paths:
|
||||
raise ValueError(f"No images found in {image_dir}")
|
||||
|
||||
print(f"Found {len(image_paths)} images in {image_dir}")
|
||||
|
||||
# Create overlapping batches
|
||||
batches = create_overlapping_batches(image_paths, batch_size, overlap_size)
|
||||
print(f"Created {len(batches)} overlapping batches")
|
||||
|
||||
# Initialize model
|
||||
model = setup_model()
|
||||
|
||||
# Process each batch
|
||||
batch_results = []
|
||||
for i, batch in enumerate(batches):
|
||||
print(f"\nProcessing batch {i+1}/{len(batches)} with {len(batch)} images")
|
||||
batch_result = process_image_batch(batch, model, i, output_dir)
|
||||
batch_results.append(batch_result)
|
||||
|
||||
# Save memory by removing large prediction data after extraction
|
||||
if "predictions" in batch_result:
|
||||
del batch_result["predictions"]
|
||||
gc.collect()
|
||||
|
||||
# Register and merge point clouds
|
||||
print("\nRegistering and merging point clouds...")
|
||||
|
||||
# Load predictions again from NPZ files to save memory during processing
|
||||
for i, result in enumerate(batch_results):
|
||||
predictions = dict(np.load(result["npz_path"], allow_pickle=True))
|
||||
batch_results[i]["predictions"] = predictions
|
||||
|
||||
# Register point clouds - FIXED LINE BELOW
|
||||
aligned_point_clouds, transformations = register_all_point_clouds(batch_results, output_dir, voxel_size)
|
||||
# Register point cloud
|
||||
|
||||
# Merge point clouds
|
||||
merged_cloud = merge_aligned_point_clouds(aligned_point_clouds, voxel_size)
|
||||
|
||||
# Convert to trimesh scene
|
||||
print("Creating final trimesh scene...")
|
||||
vertices = np.asarray(merged_cloud.points)
|
||||
colors = (np.asarray(merged_cloud.colors) * 255).astype(np.uint8)
|
||||
|
||||
scene = trimesh.Scene()
|
||||
point_cloud = trimesh.PointCloud(vertices=vertices, colors=colors)
|
||||
scene.add_geometry(point_cloud)
|
||||
|
||||
# Export merged result to GLB
|
||||
if not output_path.endswith('.glb'):
|
||||
output_path += '.glb'
|
||||
|
||||
merged_output_path = os.path.join(output_dir, os.path.basename(output_path))
|
||||
print(f"Exporting merged scene to {merged_output_path}...")
|
||||
scene.export(file_obj=merged_output_path)
|
||||
|
||||
# Also copy to requested output location
|
||||
shutil.copy(merged_output_path, output_path)
|
||||
|
||||
end_time = time.time()
|
||||
print(f"\nReconstruction completed in {end_time - start_time:.2f} seconds")
|
||||
print(f"Final point cloud has {len(vertices)} points")
|
||||
print(f"\nResults:")
|
||||
print(f"- Individual batch GLBs in subdirectories of {output_dir}")
|
||||
print(f"- Final merged GLB at {output_path}")
|
||||
print(f"- Working files in {output_dir}")
|
||||
|
||||
return output_path, output_dir
|
||||
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Main Function and CLI Interface
|
||||
# -------------------------------------------------------------------------
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="3D reconstruction from sequential images with overlapping batches"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image_dir", type=str, required=True,
|
||||
help="Directory containing input images in sequential order"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", type=str, default="merged_scene.glb",
|
||||
help="Output GLB file path for the merged scene"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch_size", type=int, default=28,
|
||||
help="Number of images per batch"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overlap_size", type=int, default=16,
|
||||
help="Number of overlapping images between consecutive batches"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--conf_threshold", type=float, default=20.0,
|
||||
help="Confidence threshold percentage for filtering points (0-100)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--voxel_size", type=float, default=0.02,
|
||||
help="Voxel size for point cloud downsampling"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
output_path, output_dir = reconstruct_from_images_with_overlap(
|
||||
image_dir=args.image_dir,
|
||||
output_path=args.output,
|
||||
batch_size=args.batch_size,
|
||||
overlap_size=args.overlap_size,
|
||||
conf_threshold=args.conf_threshold,
|
||||
voxel_size=args.voxel_size
|
||||
)
|
||||
|
||||
print(f"Reconstruction complete! Output saved to {output_path}")
|
||||
print(f"Working directory with individual GLBs: {output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user