File size: 6,003 Bytes
4eeefd1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 |
# 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
from pathlib import Path
from argparse import ArgumentParser
from tqdm import tqdm
from tqdm.contrib.concurrent import process_map
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
from rec_utils.datasets import ARKitDataset
# -------------------------------------------------------------------------
# 1) Core model inference
# -------------------------------------------------------------------------
def run_model(model, target_dir, device, max_images) -> dict:
"""
Run the VGGT model on images in the 'target_dir/images' folder and return predictions.
"""
print(f"Processing images from {target_dir}")
if not torch.cuda.is_available():
raise ValueError("CUDA is not available. Check your environment.")
# Load and preprocess images
image_names = [*target_dir.glob("*.jpg")]
image_names = sorted(image_names)
print(f"Found {len(image_names)} images")
if len(image_names) == 0:
raise ValueError(f"No images found at {target_dir}. Check your upload.")
if len(image_names) > max_images:
print(f"Downsampling {len(image_names)} images to {max_images} images")
image_names = [image_names[i] for i in np.linspace(0, len(image_names) - 1, max_images).round().astype(int)]
images = load_and_preprocess_images(image_names).to(device)
print(f"Preprocessed images shape: {images.shape}")
# Run inference
print("Running inference...")
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
with torch.no_grad():
with torch.cuda.amp.autocast(dtype=dtype):
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:])
predictions["poses"] = extrinsic
predictions["Ks"] = intrinsic
# Convert tensors to numpy
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["poses"], predictions["Ks"])
# predictions["world_points_from_depth"] = world_points
# Clean up
torch.cuda.empty_cache()
predictions["image_names"] = [str(image_name) for image_name in image_names]
return predictions
def process_scene(
model,
scene_name,
input_dir,
output_dir,
device,
max_images=10000,
force=False
):
"""
Perform reconstruction using the already-created target_dir/images.
"""
if not force and (output_dir / "predictions.npz").exists():
print(f"Skipping scene {scene_name} because it already exists")
return
start_time = time.time()
gc.collect()
torch.cuda.empty_cache()
print("Running run_model...")
with torch.no_grad():
predictions = run_model(model, input_dir, device, max_images)
# Save predictions
del predictions["images"]
np.savez(output_dir / "predictions.npz", **predictions)
del predictions
gc.collect()
torch.cuda.empty_cache()
end_time = time.time()
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--scene_names", nargs="+", default=os.listdir("/workspace-SR006.nfs2/datasets/arkitscenes/offline_prepared_data/posed_images/"))
parser.add_argument("--input_dir", type=str, default='/workspace-SR006.nfs2/datasets/arkitscenes/offline_prepared_data/posed_images/')
parser.add_argument("--output_dir", type=str, default='output/arkit_250')
parser.add_argument("--max_images", type=int, default=250)
parser.add_argument("--conf_thres", type=float, default=3.0)
parser.add_argument("--job_num", "-n", type=int, default=1)
parser.add_argument("--job_id", "-i", type=int, default=0)
parser.add_argument("--device", type=str, default="2")
parser.add_argument("--force", action="store_true")
args = parser.parse_args()
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()
scene_names = args.scene_names[args.job_id::args.job_num]
scene_names = ['47334096']
device = f"cuda:{args.device}" if torch.cuda.is_available() else "cpu"
model = model.to(device)
from datetime import datetime
errors_path = Path(f"logs/errors_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")
for scene_name in tqdm(scene_names):
print(f"Processing scene {scene_name}")
input_dir = Path(args.input_dir) / scene_name
output_dir = Path(args.output_dir) / scene_name
output_dir.mkdir(parents=True, exist_ok=True)
try:
process_scene(model, scene_name, input_dir, output_dir, device=device, max_images=args.max_images, force=args.force)
except Exception as e:
print(f"Error processing scene {scene_name}: {e}")
errors_path.parent.mkdir(parents=True, exist_ok=True)
with open(errors_path, "a") as f:
f.write(f"{scene_name}\n")
|