diff --git a/cli_demo.py b/cli_demo.py deleted file mode 100644 index 33aaf8f..0000000 --- a/cli_demo.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/colmap_writer.py b/colmap_writer.py new file mode 100644 index 0000000..66636ac --- /dev/null +++ b/colmap_writer.py @@ -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() \ No newline at end of file diff --git a/export_cam_pose.py b/export_cam_pose.py deleted file mode 100644 index 37b5a95..0000000 --- a/export_cam_pose.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/npz_reader.py b/npz_reader.py deleted file mode 100644 index f79ad7c..0000000 --- a/npz_reader.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/reconstruct.py b/reconstruct.py deleted file mode 100644 index 9c73c16..0000000 --- a/reconstruct.py +++ /dev/null @@ -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() \ No newline at end of file