diff --git a/vggt_to_colmap.py b/vggt_to_colmap.py index 55cf255..67bbdb0 100644 --- a/vggt_to_colmap.py +++ b/vggt_to_colmap.py @@ -7,6 +7,8 @@ import struct from scipy.spatial.transform import Rotation import sys from PIL import Image +import cv2 +import requests # Add VGGT to path sys.path.append("vggt/") @@ -22,13 +24,15 @@ def load_model(device=None): device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") - model = VGGT.from_pretrained("facebook/VGGT-1B") + 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) return model, device -def process_images(image_dir, model, device, conf_threshold=50.0): +def process_images(image_dir, model, device): """Process images with VGGT and return predictions.""" image_names = glob.glob(os.path.join(image_dir, "*")) image_names = sorted([f for f in image_names if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) @@ -74,6 +78,7 @@ def process_images(image_dir, model, device, conf_threshold=50.0): # Store original images for color extraction predictions["original_images"] = original_images + predictions["images"] = np.array(original_images) / 255.0 # Add normalized images as in original code return predictions, image_names @@ -99,108 +104,266 @@ def extrinsic_to_colmap_format(extrinsics): return np.array(quaternions), np.array(translations) -def is_sky_point(point, rgb): - """Check if a point is likely to be sky based on position and color.""" - # Check if point is above horizon (y is up in most 3D coordinates) - # and has blue-ish color (more blue than red or green) - return (point[1] > 1.0 and # High up - rgb[2] > 1.2 * rgb[0] and # More blue than red - rgb[2] > 1.2 * rgb[1]) # More blue than green +def download_file_from_url(url, filename): + """Downloads a file from a URL, handling redirects.""" + try: + # Get the redirect URL + response = requests.get(url, allow_redirects=False) + response.raise_for_status() # Raise HTTPError for bad requests (4xx or 5xx) -def is_black_bg_point(rgb, threshold=30): - """Check if a point is likely a black background pixel.""" - # Check if all RGB values are very low - return np.mean(rgb) < threshold + if response.status_code == 302: # Expecting a redirect + redirect_url = response.headers["Location"] + response = requests.get(redirect_url, stream=True) + response.raise_for_status() + else: + response = requests.get(url, stream=True) + response.raise_for_status() -def is_white_bg_point(rgb, threshold=230): - """Check if a point is likely a white background pixel.""" - # Check if all RGB values are very high - return np.mean(rgb) > threshold and np.std(rgb) < 15 + with open(filename, "wb") as f: + for chunk in response.iter_content(chunk_size=8192): + f.write(chunk) + print(f"Downloaded {filename} successfully.") + return True + + except requests.exceptions.RequestException as e: + print(f"Error downloading file: {e}") + return False + +def segment_sky(image_path, onnx_session, mask_filename=None): + """ + Segments sky from an image using an ONNX model. + """ + assert mask_filename is not None + image = cv2.imread(image_path) + + result_map = run_skyseg(onnx_session, [320, 320], image) + # resize the result_map to the original image size + result_map_original = cv2.resize(result_map, (image.shape[1], image.shape[0])) + + # Fix: Invert the mask so that 255 = non-sky, 0 = sky + # The model outputs low values for sky, high values for non-sky + output_mask = np.zeros_like(result_map_original) + output_mask[result_map_original < 32] = 255 # Use threshold of 32 + + os.makedirs(os.path.dirname(mask_filename), exist_ok=True) + cv2.imwrite(mask_filename, output_mask) + return output_mask + +def run_skyseg(onnx_session, input_size, image): + """ + Runs sky segmentation inference using ONNX model. + """ + import copy + + # Pre process:Resize, BGR->RGB, Transpose, PyTorch standardization, float32 cast + temp_image = copy.deepcopy(image) + resize_image = cv2.resize(temp_image, dsize=(input_size[0], input_size[1])) + x = cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB) + x = np.array(x, dtype=np.float32) + mean = [0.485, 0.456, 0.406] + std = [0.229, 0.224, 0.225] + x = (x / 255 - mean) / std + x = x.transpose(2, 0, 1) + x = x.reshape(-1, 3, input_size[0], input_size[1]).astype("float32") + + # Inference + input_name = onnx_session.get_inputs()[0].name + output_name = onnx_session.get_outputs()[0].name + onnx_result = onnx_session.run([output_name], {input_name: x}) + + # Post process + onnx_result = np.array(onnx_result).squeeze() + min_value = np.min(onnx_result) + max_value = np.max(onnx_result) + onnx_result = (onnx_result - min_value) / (max_value - min_value) + onnx_result *= 255 + onnx_result = onnx_result.astype("uint8") + + return onnx_result def filter_and_prepare_points(predictions, conf_threshold, mask_sky=False, mask_black_bg=False, - mask_white_bg=False, stride=4): + mask_white_bg=False, stride=4, target_dir=None, prediction_mode="Depthmap and Camera Branch"): """ Filter points based on confidence and prepare for COLMAP format. - Uses stride to sample fewer points for efficiency. - Extracts RGB color from original images. - Applies optional filtering for sky and backgrounds. + Implementation matches the conventions in the original VGGT code. """ - # Scale confidence threshold from percentage to actual value range - max_conf = np.max(predictions["depth_conf"]) - conf_threshold = conf_threshold / 100.0 * max_conf + # Use the correct point cloud based on prediction mode + if "Pointmap" in prediction_mode: + print("Using Pointmap Branch") + if "world_points" in predictions: + pred_world_points = predictions["world_points"] + pred_world_points_conf = predictions.get("world_points_conf", np.ones_like(pred_world_points[..., 0])) + else: + print("Warning: world_points not found in predictions, falling back to depth-based points") + pred_world_points = predictions["world_points_from_depth"] + pred_world_points_conf = predictions.get("depth_conf", np.ones_like(pred_world_points[..., 0])) + else: + print("Using Depthmap and Camera Branch") + pred_world_points = predictions["world_points_from_depth"] + pred_world_points_conf = predictions.get("depth_conf", np.ones_like(pred_world_points[..., 0])) - print(f"Confidence threshold: {conf_threshold:.4f} (scaled from {conf_threshold*100/max_conf:.2f}%)") - print(f"Filtering options - Sky: {mask_sky}, Black BG: {mask_black_bg}, White BG: {mask_white_bg}") + # Get images for color extraction + colors_rgb = predictions["images"] # Already in normalized format (0-1) + colors_rgb = (colors_rgb * 255).astype(np.uint8) - # Initialize containers - points3D = [] - point_indices = {} # Maps point hash to its index - image_points2D = [[] for _ in range(len(predictions["depth"]))] - - # Get world points and confidences - world_points = predictions["world_points_from_depth"] # (num_images, H, W, 3) - confidences = predictions["depth_conf"] # (num_images, H, W) - original_images = predictions["original_images"] # List of numpy arrays (H, W, 3) - - # Process each image - for img_idx in range(len(world_points)): - h, w = world_points[img_idx].shape[:2] - orig_img = original_images[img_idx] - - # Handle possible resizing between original image and depth map - orig_h, orig_w = orig_img.shape[:2] - scale_y, scale_x = orig_h / h, orig_w / w - - # Sample points with stride - for y in range(0, h, stride): - for x in range(0, w, stride): - if confidences[img_idx, y, x] > conf_threshold: - # Get 3D point - point3D = world_points[img_idx, y, x] + # Handle sky segmentation if requested + if mask_sky and target_dir is not None: + print("Applying sky segmentation mask") + try: + import onnxruntime + + skyseg_session = None + target_dir_images = os.path.join(target_dir, "images") + os.makedirs(target_dir_images, exist_ok=True) + + # Save the RGB images temporarily if they don't exist + image_list = [] + for i, img in enumerate(colors_rgb): + img_path = os.path.join(target_dir_images, f"image_{i:04d}.png") + image_list.append(os.path.basename(img_path)) + if not os.path.exists(img_path): + cv2.imwrite(img_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) + + # Ensure the output directory for masks exists + sky_masks_dir = os.path.join(target_dir, "sky_masks") + os.makedirs(sky_masks_dir, exist_ok=True) + + # Download skyseg.onnx if it doesn't exist + if not os.path.exists("skyseg.onnx"): + print("Downloading skyseg.onnx...") + download_success = download_file_from_url( + "https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx", + "skyseg.onnx" + ) + if not download_success: + print("Failed to download skyseg model, skipping sky filtering") + mask_sky = False + + if mask_sky: # Continue only if download succeeded + S, H, W = pred_world_points_conf.shape + sky_mask_list = [] + + for i, image_name in enumerate(image_list): + image_filepath = os.path.join(target_dir_images, image_name) + mask_filepath = os.path.join(sky_masks_dir, image_name) - # Skip invalid points - if not np.all(np.isfinite(point3D)): - continue - - # Get RGB color from original image - # Map coordinates from depth map to original image - orig_y = min(int(y * scale_y), orig_h - 1) - orig_x = min(int(x * scale_x), orig_w - 1) - rgb = orig_img[orig_y, orig_x] - - # Apply filters - if mask_sky and is_sky_point(point3D, rgb): - continue - if mask_black_bg and is_black_bg_point(rgb): - continue - if mask_white_bg and is_white_bg_point(rgb): - continue - - # Generate a hash for this point - point_hash = hash_point(point3D, scale=100) - - if point_hash not in point_indices: - # Create new point - point_idx = len(points3D) - point_indices[point_hash] = point_idx - - point_entry = { - "id": point_idx, - "xyz": point3D, - "rgb": rgb, - "error": 1.0, - "track": [(img_idx, len(image_points2D[img_idx]))] - } - points3D.append(point_entry) + # Check if mask already exists + if os.path.exists(mask_filepath): + # Load existing mask + sky_mask = cv2.imread(mask_filepath, cv2.IMREAD_GRAYSCALE) else: - # Update existing point's track - point_idx = point_indices[point_hash] - points3D[point_idx]["track"].append((img_idx, len(image_points2D[img_idx]))) + # Generate new mask + if skyseg_session is None: + skyseg_session = onnxruntime.InferenceSession("skyseg.onnx") + sky_mask = segment_sky(image_filepath, skyseg_session, mask_filepath) - # Add 2D point to image - image_points2D[img_idx].append((x, y, point_indices[point_hash])) + # Resize mask to match H×W if needed + if sky_mask.shape[0] != H or sky_mask.shape[1] != W: + sky_mask = cv2.resize(sky_mask, (W, H)) + + sky_mask_list.append(sky_mask) + + # Convert list to numpy array with shape S×H×W + sky_mask_array = np.array(sky_mask_list) + + # Apply sky mask to confidence scores + sky_mask_binary = (sky_mask_array > 0.1).astype(np.float32) + pred_world_points_conf = pred_world_points_conf * sky_mask_binary + + except (ImportError, Exception) as e: + print(f"Error in sky segmentation: {e}") + print("Falling back to basic sky filtering") + mask_sky = False - print(f"Generated {len(points3D)} 3D points with {sum(len(pts) for pts in image_points2D)} observations") + # Reshape for processing + vertices_3d = pred_world_points.reshape(-1, 3) + conf = pred_world_points_conf.reshape(-1) + colors_rgb_flat = colors_rgb.reshape(-1, 3) + + # Apply confidence threshold exactly as in original code + if conf_threshold == 0.0: + conf_thres_value = 0.0 + else: + conf_thres_value = np.percentile(conf, conf_threshold) + + print(f"Using confidence threshold: {conf_threshold}% (value: {conf_thres_value:.4f})") + conf_mask = (conf >= conf_thres_value) & (conf > 1e-5) + + # Apply black background filter if requested + if mask_black_bg: + print("Filtering black background") + black_bg_mask = colors_rgb_flat.sum(axis=1) >= 16 + conf_mask = conf_mask & black_bg_mask + + # Apply white background filter if requested + if mask_white_bg: + print("Filtering white background") + white_bg_mask = ~((colors_rgb_flat[:, 0] > 240) & (colors_rgb_flat[:, 1] > 240) & (colors_rgb_flat[:, 2] > 240)) + conf_mask = conf_mask & white_bg_mask + + # Filter points + filtered_vertices = vertices_3d[conf_mask] + filtered_colors = colors_rgb_flat[conf_mask] + + # Check if we have any points left + if len(filtered_vertices) == 0: + print("Warning: No points remaining after filtering. Using default point.") + filtered_vertices = np.array([[0, 0, 0]]) + filtered_colors = np.array([[200, 200, 200]]) + + print(f"Filtered to {len(filtered_vertices)} points") + + # Now prepare the filtered points for COLMAP format + points3D = [] + point_indices = {} + image_points2D = [[] for _ in range(len(pred_world_points))] + + # Get the original 3D indices from the flat indices + S, H, W = pred_world_points.shape[0:3] + flat_indices = np.where(conf_mask)[0] + + for flat_idx in flat_indices[::stride]: # Use stride to reduce point count + # Convert flat index back to 3D coordinates + img_idx = flat_idx // (H * W) + remainder = flat_idx % (H * W) + y = remainder // W + x = remainder % W + + if img_idx >= S: # Handle edge cases + continue + + point3D = vertices_3d[flat_idx] + rgb = filtered_colors[flat_idx - flat_indices[0]] if flat_idx >= flat_indices[0] else filtered_colors[0] + + # Skip invalid points + if not np.all(np.isfinite(point3D)): + continue + + # Generate a hash for this point to identify duplicates + point_hash = hash_point(point3D, scale=100) + + if point_hash not in point_indices: + # Create new point + point_idx = len(points3D) + point_indices[point_hash] = point_idx + + point_entry = { + "id": point_idx, + "xyz": point3D, + "rgb": rgb, + "error": 1.0, + "track": [(img_idx, len(image_points2D[img_idx]))] + } + points3D.append(point_entry) + else: + # Update existing point's track + point_idx = point_indices[point_hash] + points3D[point_idx]["track"].append((img_idx, len(image_points2D[img_idx]))) + + # Add 2D point to image + image_points2D[img_idx].append((x, y, point_indices[point_hash])) + + print(f"Prepared {len(points3D)} 3D points with {sum(len(pts) for pts in image_points2D)} observations for COLMAP") return points3D, image_points2D def hash_point(point, scale=100): @@ -378,7 +541,7 @@ def main(): parser.add_argument("--output_dir", type=str, default="colmap_output", help="Directory to save COLMAP files") parser.add_argument("--conf_threshold", type=float, default=50.0, - help="Confidence threshold (0-100) for including points") + help="Confidence threshold (0-100%) for including points") parser.add_argument("--mask_sky", action="store_true", help="Filter out points likely to be sky") parser.add_argument("--mask_black_bg", action="store_true", @@ -389,31 +552,41 @@ def main(): help="Output binary COLMAP files instead of text") parser.add_argument("--stride", type=int, default=1, help="Stride for point sampling (higher = fewer points)") + parser.add_argument("--prediction_mode", type=str, default="Depthmap and Camera Branch", + choices=["Depthmap and Camera Branch", "Pointmap Branch"], + help="Which prediction branch to use") args = parser.parse_args() # Create output directory os.makedirs(args.output_dir, exist_ok=True) + # Create temporary directory for processing + temp_dir = os.path.join(args.output_dir, "temp") + os.makedirs(os.path.join(temp_dir, "images"), exist_ok=True) + # Load model model, device = load_model() # Process images - predictions, image_names = process_images( - args.image_dir, model, device, args.conf_threshold) + predictions, image_names = process_images(args.image_dir, model, device) # Extract quaternions and translations print("Converting camera parameters to COLMAP format...") quaternions, translations = extrinsic_to_colmap_format(predictions["extrinsic"]) - # Filter and prepare points + # Filter and prepare points using the original VGGT conventions print(f"Filtering points with confidence threshold {args.conf_threshold}% and stride {args.stride}...") points3D, image_points2D = filter_and_prepare_points( - predictions, args.conf_threshold, + predictions, + args.conf_threshold, mask_sky=args.mask_sky, mask_black_bg=args.mask_black_bg, mask_white_bg=args.mask_white_bg, - stride=args.stride) + stride=args.stride, + target_dir=temp_dir, + prediction_mode=args.prediction_mode + ) # Get image dimensions height, width = predictions["depth"].shape[1:3] @@ -444,4 +617,4 @@ def main(): print(f"COLMAP files successfully written to {args.output_dir}") if __name__ == "__main__": - main() \ No newline at end of file + main() \ No newline at end of file