Point cloud and camera estimation works well
This commit is contained in:
Atakan Topaloglu
@@ -11,6 +11,7 @@ reconstructions/
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
output_colmap*
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
|
||||
@@ -0,0 +1,277 @@
|
||||
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)):
|
||||
# Extract image dimensions
|
||||
height, width = shape
|
||||
|
||||
# Extract focal length and principal point
|
||||
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, sample_step=20):
|
||||
"""Extract a sparse point cloud from point maps with colors from images"""
|
||||
points = []
|
||||
point_colors = [] # Will store actual colors from images
|
||||
point_origins = [] # Track which image/pixel each point came from
|
||||
|
||||
for img_idx, (point_map, img) in enumerate(zip(point_maps, images)):
|
||||
h, w, _ = point_map.shape
|
||||
for y in range(0, h, sample_step):
|
||||
for x in range(0, w, sample_step):
|
||||
# Get 3D point coordinates
|
||||
point3d = point_map[y, x]
|
||||
|
||||
# Skip invalid points
|
||||
if np.linalg.norm(point3d) < 0.001 or np.linalg.norm(point3d) > 1000:
|
||||
continue
|
||||
|
||||
# Extract the color from the image
|
||||
color = img[y, x]
|
||||
|
||||
# Add point, color, and origin info
|
||||
points.append(point3d)
|
||||
point_colors.append(color)
|
||||
point_origins.append((img_idx, y, x))
|
||||
|
||||
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"""
|
||||
# if len(points) == 0:
|
||||
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)):
|
||||
# Get image dimensions
|
||||
height, width = shape
|
||||
|
||||
# Transform point to camera coordinates
|
||||
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)):
|
||||
# Extract rotation matrix and translation vector
|
||||
rot_matrix = extrinsic[:3, :3]
|
||||
trans_vector = extrinsic[:3, 3]
|
||||
|
||||
# Convert to camera-to-world transformation
|
||||
rot_c2w = rot_matrix.T # Transpose for inverse rotation
|
||||
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 # Assuming one camera per image
|
||||
image_name = os.path.basename(image_names[img_idx])
|
||||
|
||||
# Write image pose
|
||||
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")
|
||||
|
||||
# Write 2D point observations
|
||||
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
|
||||
|
||||
# Convert color to integers
|
||||
r, g, b = [int(c * 255) for c in color]
|
||||
error = 1.0 # Placeholder for reprojection error
|
||||
|
||||
# Format track information
|
||||
track_str = " ".join([f"{img_id} {point2d_idx}" for img_id, point2d_idx in track])
|
||||
|
||||
# Write point information
|
||||
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, image_paths, output_dir, image_shapes, images):
|
||||
"""Convert model estimates to COLMAP format files"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Extract sparse point cloud with colors from images
|
||||
print("Extracting colored sparse point cloud...")
|
||||
points3d, point_colors, point_origins = extract_sparse_point_cloud(point_maps, images, sample_step=20)
|
||||
print(f"Extracted {len(points3d)} colored points")
|
||||
|
||||
# Compute point visibility in each image
|
||||
print("Computing point visibility...")
|
||||
image_points, point_tracks = compute_point_visibility(
|
||||
points3d, extrinsics, intrinsics, image_shapes
|
||||
)
|
||||
|
||||
# Write COLMAP files
|
||||
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)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get list of images
|
||||
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")
|
||||
|
||||
# Process images with VGGT model
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
|
||||
|
||||
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)
|
||||
|
||||
# Also load images in RGB format for coloring
|
||||
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("Estimating depth maps...")
|
||||
depth_map, _ = 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)
|
||||
)
|
||||
|
||||
# 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
|
||||
|
||||
# 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:
|
||||
# Use simple resizing - could be improved with better interpolation
|
||||
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
|
||||
convert_to_colmap_format(
|
||||
extrinsic_np, intrinsic_np, point_map_np, image_paths,
|
||||
args.output_dir, image_shapes, resized_images
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,127 @@
|
||||
#!/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()
|
||||
Binary file not shown.
+296
-79
@@ -29,6 +29,247 @@ 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
|
||||
@@ -142,7 +383,7 @@ def process_image_batch(image_batch, model, batch_id, output_dir):
|
||||
glb_path = os.path.join(batch_dir, "scene.glb")
|
||||
glbscene = predictions_to_glb(
|
||||
predictions,
|
||||
conf_thres=90.0,
|
||||
conf_thres=50.0,
|
||||
filter_by_frames="All",
|
||||
mask_black_bg=True,
|
||||
mask_white_bg=True,
|
||||
@@ -268,92 +509,68 @@ def find_overlapping_cameras(batch1_cameras, batch2_cameras, batch1_images, batc
|
||||
|
||||
|
||||
def register_consecutive_point_clouds(source_pcd, target_pcd, source_cameras, target_cameras,
|
||||
source_images, target_images):
|
||||
source_images, target_images, output_dir=None, batch_id=None):
|
||||
"""
|
||||
Register two point clouds from consecutive batches using both overlapping cameras
|
||||
and point cloud features
|
||||
|
||||
Returns:
|
||||
Transformation matrix to align source to target
|
||||
Register two point clouds from consecutive batches with improved initialization
|
||||
"""
|
||||
print("Registering consecutive point clouds...")
|
||||
|
||||
# 1. Find overlapping cameras
|
||||
source_overlap_idx, target_overlap_idx = find_overlapping_cameras(
|
||||
# 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
|
||||
)
|
||||
|
||||
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]
|
||||
# 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
|
||||
)
|
||||
|
||||
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
|
||||
# 3. Apply initial transformation and proceed with ICP refinement
|
||||
source_aligned = copy.deepcopy(source_pcd)
|
||||
source_aligned.transform(initial_transformation)
|
||||
|
||||
# 4. Refine alignment with Colored ICP
|
||||
# Prepare for colored ICP
|
||||
# Prepare for 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)
|
||||
)
|
||||
# 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
|
||||
|
||||
# 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}")
|
||||
# 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, voxel_size=0.02):
|
||||
def register_all_point_clouds(batch_results, output_dir, voxel_size=0.02):
|
||||
"""
|
||||
Register all point clouds from overlapping batches
|
||||
|
||||
Returns:
|
||||
List of transformation matrices and list of point clouds
|
||||
Register all point clouds with better initialization and visualization
|
||||
"""
|
||||
point_clouds = []
|
||||
camera_positions = []
|
||||
@@ -363,7 +580,7 @@ def register_all_point_clouds(batch_results, voxel_size=0.02):
|
||||
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
|
||||
predictions, conf_threshold=30.0, voxel_size=voxel_size, mask_bg=True
|
||||
)
|
||||
point_clouds.append(pcd)
|
||||
camera_positions.append(cameras)
|
||||
@@ -377,25 +594,24 @@ def register_all_point_clouds(batch_results, voxel_size=0.02):
|
||||
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
|
||||
source_images, target_images, output_dir, i
|
||||
)
|
||||
|
||||
transformations.append(transformation)
|
||||
|
||||
# Apply transformations to point clouds
|
||||
# Apply transformations to align all 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
|
||||
# Calculate cumulative 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
|
||||
|
||||
# Apply transformation
|
||||
transformed_pcd = copy.deepcopy(point_clouds[i])
|
||||
transformed_pcd.transform(cumulative_transform)
|
||||
aligned_point_clouds.append(transformed_pcd)
|
||||
@@ -487,8 +703,9 @@ def reconstruct_from_images_with_overlap(image_dir, output_path, batch_size=8, o
|
||||
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)
|
||||
# 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)
|
||||
@@ -540,15 +757,15 @@ def main():
|
||||
help="Output GLB file path for the merged scene"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch_size", type=int, default=30,
|
||||
"--batch_size", type=int, default=28,
|
||||
help="Number of images per batch"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overlap_size", type=int, default=20,
|
||||
"--overlap_size", type=int, default=16,
|
||||
help="Number of overlapping images between consecutive batches"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--conf_threshold", type=float, default=90.0,
|
||||
"--conf_threshold", type=float, default=20.0,
|
||||
help="Confidence threshold percentage for filtering points (0-100)"
|
||||
)
|
||||
parser.add_argument(
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
import pandas as pd
|
||||
|
||||
def subtract_consecutive_rows(csv_file, output_file):
|
||||
# Read the csv file into a DataFrame
|
||||
df = pd.read_csv(csv_file)
|
||||
|
||||
# Check if the DataFrame has at least two rows
|
||||
if len(df) < 2:
|
||||
print("The CSV file must have at least two rows.")
|
||||
return
|
||||
|
||||
# Subtract consecutive rows
|
||||
result_df = df.diff().iloc[1:]
|
||||
|
||||
# Save the result to a new CSV file
|
||||
result_df.to_csv(output_file, index=False)
|
||||
print(f"Result saved to {output_file}")
|
||||
|
||||
# Example usage
|
||||
csv_file_path = 'camera_poses2.csv'
|
||||
output_file_path = 'camera_poses_sub2.csv'
|
||||
subtract_consecutive_rows(csv_file_path, output_file_path)
|
||||
Reference in New Issue
Block a user