575 lines
20 KiB
Python
575 lines
20 KiB
Python
#!/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 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=90.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):
|
|
"""
|
|
Register two point clouds from consecutive batches using both overlapping cameras
|
|
and point cloud features
|
|
|
|
Returns:
|
|
Transformation matrix to align source to target
|
|
"""
|
|
print("Registering consecutive point clouds...")
|
|
|
|
# 1. Find overlapping cameras
|
|
source_overlap_idx, target_overlap_idx = find_overlapping_cameras(
|
|
source_cameras, target_cameras, source_images, target_images
|
|
)
|
|
|
|
if len(source_overlap_idx) < 2:
|
|
print("Warning: Less than 2 overlapping cameras found, alignment may be inaccurate")
|
|
# Fall back to using all cameras if not enough overlap
|
|
source_pts = source_cameras
|
|
target_pts = target_cameras[:len(source_pts)] if len(target_cameras) > len(source_pts) else target_cameras
|
|
else:
|
|
# Use overlapping cameras for initial alignment
|
|
source_pts = source_cameras[source_overlap_idx]
|
|
target_pts = target_cameras[target_overlap_idx]
|
|
|
|
print(f"Found {len(source_overlap_idx)} overlapping cameras for alignment")
|
|
|
|
# 2. Initial alignment using camera positions
|
|
# 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)
|
|
|
|
# Initial transformation matrix
|
|
initial_transformation = np.eye(4)
|
|
initial_transformation[:3, :3] = R
|
|
initial_transformation[:3, 3] = t
|
|
|
|
# 3. Apply initial transformation to source point cloud
|
|
source_aligned = copy.deepcopy(source_pcd)
|
|
source_aligned.transform(initial_transformation)
|
|
|
|
# 4. Refine alignment with Colored ICP
|
|
# Prepare for colored ICP
|
|
source_aligned.estimate_normals()
|
|
target_pcd.estimate_normals()
|
|
|
|
# Run Colored ICP
|
|
print("Running colored ICP for fine alignment...")
|
|
reg_p2p = o3d.pipelines.registration.registration_colored_icp(
|
|
source_aligned, target_pcd, 0.05,
|
|
initial_transformation,
|
|
o3d.pipelines.registration.TransformationEstimationForColoredICP(),
|
|
o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=50)
|
|
)
|
|
|
|
# Final transformation matrix (combines initial and refinement)
|
|
final_transformation = reg_p2p.transformation @ initial_transformation
|
|
|
|
print(f"Registration complete with fitness: {reg_p2p.fitness}, RMSE: {reg_p2p.inlier_rmse}")
|
|
|
|
return final_transformation
|
|
|
|
|
|
def register_all_point_clouds(batch_results, voxel_size=0.02):
|
|
"""
|
|
Register all point clouds from overlapping batches
|
|
|
|
Returns:
|
|
List of transformation matrices and list of point clouds
|
|
"""
|
|
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=50.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"]
|
|
|
|
transformation = register_consecutive_point_clouds(
|
|
source_pcd, target_pcd, source_cameras, target_cameras,
|
|
source_images, target_images
|
|
)
|
|
|
|
transformations.append(transformation)
|
|
|
|
# Apply transformations to point clouds
|
|
aligned_point_clouds = [point_clouds[0]] # First cloud is reference
|
|
|
|
# Forward transformations - transform each cloud to first cloud's reference frame
|
|
for i in range(1, len(point_clouds)):
|
|
# Initialize with identity transformation
|
|
cumulative_transform = np.eye(4)
|
|
|
|
# Compose transformations from current cloud back to the first cloud
|
|
for j in range(i-1, -1, -1):
|
|
cumulative_transform = transformations[j] @ cumulative_transform
|
|
|
|
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
|
|
aligned_point_clouds, transformations = register_all_point_clouds(batch_results, voxel_size)
|
|
|
|
# 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=30,
|
|
help="Number of images per batch"
|
|
)
|
|
parser.add_argument(
|
|
"--overlap_size", type=int, default=20,
|
|
help="Number of overlapping images between consecutive batches"
|
|
)
|
|
parser.add_argument(
|
|
"--conf_threshold", type=float, default=90.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() |