Alignment is upside down, likely due to poor initialization
This commit is contained in:
Atakan Topaloglu
+1
-1
@@ -7,7 +7,7 @@ __pycache__/
|
||||
**/__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
input_images*
|
||||
reconstructions/
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
|
||||
Binary file not shown.
+480
-174
@@ -1,101 +1,188 @@
|
||||
#!/usr/bin/env python
|
||||
# Integrated VGGT + Point Cloud Registration with Overlapping Image Batches
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
import open3d as o3d
|
||||
import sys
|
||||
import shutil
|
||||
from datetime import datetime
|
||||
import glob
|
||||
import gc
|
||||
import time
|
||||
import argparse
|
||||
import copy
|
||||
import trimesh
|
||||
import os
|
||||
import open3d as o3d
|
||||
from tqdm import tqdm
|
||||
import matplotlib.pyplot as plt
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
sys.path.append("vggt/")
|
||||
|
||||
def merge_point_clouds(prediction_sets, voxel_size=0.02):
|
||||
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):
|
||||
"""
|
||||
Register and merge multiple point clouds from overlapping image sets
|
||||
Create overlapping batches of sequential images
|
||||
|
||||
Args:
|
||||
prediction_sets: List of prediction dictionaries containing point clouds and camera parameters
|
||||
voxel_size: Resolution for downsampling and registration
|
||||
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:
|
||||
merged_scene: Final trimesh.Scene with merged point cloud
|
||||
List of image path batches
|
||||
"""
|
||||
# 1. Extract point clouds and camera positions from each prediction set
|
||||
point_clouds = []
|
||||
camera_positions = []
|
||||
if batch_size <= overlap_size:
|
||||
raise ValueError(f"Batch size ({batch_size}) must be larger than overlap size ({overlap_size})")
|
||||
|
||||
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")
|
||||
if len(image_paths) <= batch_size:
|
||||
return [image_paths]
|
||||
|
||||
# 2. Register point clouds using both camera positions and point features
|
||||
transformations = register_point_clouds(point_clouds, camera_positions)
|
||||
batches = []
|
||||
step_size = batch_size - overlap_size
|
||||
|
||||
# 3. Transform and merge point clouds
|
||||
merged_cloud = o3d.geometry.PointCloud()
|
||||
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
|
||||
|
||||
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
|
||||
# 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()
|
||||
|
||||
# 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
|
||||
return batches
|
||||
|
||||
|
||||
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)
|
||||
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
|
||||
@@ -103,17 +190,43 @@ def extract_colored_point_cloud(points, images, confidence, threshold=50.0):
|
||||
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 = confidence.reshape(-1)
|
||||
conf = conf.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]
|
||||
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)
|
||||
|
||||
return vertices_3d, colors_rgb
|
||||
# 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):
|
||||
@@ -133,76 +246,56 @@ def extract_camera_centers(extrinsics):
|
||||
return np.array(centers)
|
||||
|
||||
|
||||
def register_point_clouds(point_clouds, camera_positions):
|
||||
def find_overlapping_cameras(batch1_cameras, batch2_cameras, batch1_images, batch2_images):
|
||||
"""
|
||||
Register point clouds using both camera positions and ICP
|
||||
Identify overlapping cameras between batches based on image filenames
|
||||
|
||||
Returns:
|
||||
List of transformation matrices to align each point cloud to the first one
|
||||
Indices of overlapping cameras in each batch
|
||||
"""
|
||||
if len(point_clouds) <= 1:
|
||||
return []
|
||||
# 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]
|
||||
|
||||
transformations = []
|
||||
target_cloud = point_clouds[0]
|
||||
target_cameras = camera_positions[0]
|
||||
# Find common filenames
|
||||
common_files = set(batch1_filenames) & set(batch2_filenames)
|
||||
|
||||
# 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
|
||||
# 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 estimate_transformation_from_cameras(source_cameras, target_cameras):
|
||||
def register_consecutive_point_clouds(source_pcd, target_pcd, source_cameras, target_cameras,
|
||||
source_images, target_images):
|
||||
"""
|
||||
Estimate transformation between point clouds using camera positions
|
||||
using a simplified version of Umeyama algorithm
|
||||
Register two point clouds from consecutive batches using both overlapping cameras
|
||||
and point cloud features
|
||||
|
||||
Returns:
|
||||
Transformation matrix to align source to target
|
||||
"""
|
||||
if len(source_cameras) < 3 or len(target_cameras) < 3:
|
||||
# Not enough camera positions for reliable estimation
|
||||
return np.eye(4)
|
||||
print("Registering consecutive point clouds...")
|
||||
|
||||
# Find corresponding cameras by proximity
|
||||
# This works if the camera positions are in sequential order
|
||||
s_cam = source_cameras
|
||||
t_cam = target_cameras
|
||||
# 1. Find overlapping cameras
|
||||
source_overlap_idx, target_overlap_idx = find_overlapping_cameras(
|
||||
source_cameras, target_cameras, source_images, target_images
|
||||
)
|
||||
|
||||
# If the number of cameras differs, use the minimum length
|
||||
min_len = min(len(s_cam), len(t_cam))
|
||||
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]
|
||||
|
||||
# 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):]])
|
||||
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)
|
||||
@@ -211,7 +304,7 @@ def estimate_transformation_from_cameras(source_cameras, target_cameras):
|
||||
source_centered = source_pts - source_centroid
|
||||
target_centered = target_pts - target_centroid
|
||||
|
||||
# Compute rotation using singular value decomposition
|
||||
# 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)
|
||||
@@ -224,45 +317,258 @@ def estimate_transformation_from_cameras(source_cameras, target_cameras):
|
||||
# Translation
|
||||
t = target_centroid - np.dot(R, source_centroid)
|
||||
|
||||
# Create transformation matrix
|
||||
transformation = np.eye(4)
|
||||
transformation[:3, :3] = R
|
||||
transformation[:3, 3] = t
|
||||
# Initial transformation matrix
|
||||
initial_transformation = np.eye(4)
|
||||
initial_transformation[:3, :3] = R
|
||||
initial_transformation[:3, 3] = t
|
||||
|
||||
return transformation
|
||||
# 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():
|
||||
"""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')
|
||||
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()
|
||||
|
||||
# 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)
|
||||
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
|
||||
)
|
||||
|
||||
# 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}")
|
||||
print(f"Reconstruction complete! Output saved to {output_path}")
|
||||
print(f"Working directory with individual GLBs: {output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user