diff --git a/README.md b/README.md index 6dc106a..6dbeba4 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,8 @@ from vggt.models.vggt import VGGT from vggt.utils.load_fn import load_and_preprocess_images device = "cuda" if torch.cuda.is_available() else "cpu" -dtype = torch.bfloat16 # or torch.float16 +# bfloat16 is supported on Ampere GPUs (Compute Capability 8.0+) +dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16 # Initialize the model and load the pretrained weights. # This will automatically download the model weights the first time it's run, which may take a while. @@ -57,9 +58,10 @@ model = VGGT.from_pretrained("facebook/VGGT-1B").to(device) image_names = ["path/to/imageA.png", "path/to/imageB.png", "path/to/imageC.png"] images = load_and_preprocess_images(image_names).to(device) -with torch.no_grad() and torch.cuda.amp.autocast(dtype=dtype): - # Predict attributes including cameras, depth maps, and point maps. - predictions = model(images) +with torch.no_grad(): + with torch.cuda.amp.autocast(dtype=dtype): + # Predict attributes including cameras, depth maps, and point maps. + predictions = model(images) ``` The model weights will be automatically downloaded from Hugging Face. If you encounter issues such as slow loading, you can manually download them [here](https://huggingface.co/facebook/VGGT-1B/blob/main/model.pt) and load, or: @@ -165,6 +167,10 @@ This plots the tracks on the images and saves them to the specified output direc Our model shows surprisingly good performance on single-view reconstruction, although it was never trained for this task. The model does not need to duplicate the single-view image to a pair, instead, it can directly infer the 3D structure from the tokens of the single view image. Feel free to try it with our demos above, which naturally works for single-view reconstruction. +We did not quantitatively test monocular depth estimation performance ourselves, but [@kabouzeid](https://github.com/kabouzeid) generously provided a comparison of VGGT to recent methods [here](https://github.com/facebookresearch/vggt/issues/36). VGGT shows competitive or better results compared to state-of-the-art monocular approaches such as DepthAnything v2 or MoGe, despite never being explicitly trained for single-view tasks. + + + ## Runtime and GPU Memory We benchmark the runtime and GPU memory usage of VGGT's aggregator on a single NVIDIA H100 GPU across various input sizes. diff --git a/demo_gradio.py b/demo_gradio.py index 38eb4b2..8dcef99 100644 --- a/demo_gradio.py +++ b/demo_gradio.py @@ -68,8 +68,10 @@ def run_model(target_dir, model) -> dict: # 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=torch.bfloat16): + with torch.cuda.amp.autocast(dtype=dtype): predictions = model(images) # Convert pose encoding to extrinsic and intrinsic matrices @@ -217,7 +219,7 @@ def gradio_demo( # Handle None frame_filter if frame_filter is None: frame_filter = "All" - + # Build a GLB file name glbfile = os.path.join( target_dir, @@ -244,7 +246,7 @@ def gradio_demo( torch.cuda.empty_cache() end_time = time.time() - print(f"Total time: {end_time - start_time:.2f} seconds") + print(f"Total time: {end_time - start_time:.2f} seconds (including IO)") log_msg = f"Reconstruction Success ({len(all_files)} frames). Waiting for visualization." return glbfile, log_msg, gr.Dropdown(choices=frame_filter_choices, value=frame_filter, interactive=True) @@ -379,9 +381,9 @@ with gr.Blocks( # Instead of gr.State, we use a hidden Textbox: is_example = gr.Textbox(label="is_example", visible=False, value="None") num_images = gr.Textbox(label="num_images", visible=False, value="None") - + gr.HTML( - """ + """

🏛️ VGGT: Visual Geometry Grounded Transformer

🐙 GitHub Repository | @@ -542,7 +544,16 @@ with gr.Blocks( fn=update_log, inputs=[], outputs=[log_output] ).then( fn=gradio_demo, - inputs=[target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode], + inputs=[ + target_dir_output, + conf_thres, + frame_filter, + mask_black_bg, + mask_white_bg, + show_cam, + mask_sky, + prediction_mode, + ], outputs=[reconstruction_output, log_output, frame_filter], ).then( fn=lambda: "False", inputs=[], outputs=[is_example] # set is_example to "False" @@ -553,37 +564,107 @@ with gr.Blocks( # ------------------------------------------------------------------------- conf_thres.change( update_visualization, - [target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example], + [ + target_dir_output, + conf_thres, + frame_filter, + mask_black_bg, + mask_white_bg, + show_cam, + mask_sky, + prediction_mode, + is_example, + ], [reconstruction_output, log_output], ) frame_filter.change( update_visualization, - [target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example], + [ + target_dir_output, + conf_thres, + frame_filter, + mask_black_bg, + mask_white_bg, + show_cam, + mask_sky, + prediction_mode, + is_example, + ], [reconstruction_output, log_output], ) mask_black_bg.change( update_visualization, - [target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example], + [ + target_dir_output, + conf_thres, + frame_filter, + mask_black_bg, + mask_white_bg, + show_cam, + mask_sky, + prediction_mode, + is_example, + ], [reconstruction_output, log_output], ) mask_white_bg.change( update_visualization, - [target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example], + [ + target_dir_output, + conf_thres, + frame_filter, + mask_black_bg, + mask_white_bg, + show_cam, + mask_sky, + prediction_mode, + is_example, + ], [reconstruction_output, log_output], ) show_cam.change( update_visualization, - [target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example], + [ + target_dir_output, + conf_thres, + frame_filter, + mask_black_bg, + mask_white_bg, + show_cam, + mask_sky, + prediction_mode, + is_example, + ], [reconstruction_output, log_output], ) mask_sky.change( update_visualization, - [target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example], + [ + target_dir_output, + conf_thres, + frame_filter, + mask_black_bg, + mask_white_bg, + show_cam, + mask_sky, + prediction_mode, + is_example, + ], [reconstruction_output, log_output], ) prediction_mode.change( update_visualization, - [target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example], + [ + target_dir_output, + conf_thres, + frame_filter, + mask_black_bg, + mask_white_bg, + show_cam, + mask_sky, + prediction_mode, + is_example, + ], [reconstruction_output, log_output], ) diff --git a/demo_viser.py b/demo_viser.py index d25e91b..c52726b 100644 --- a/demo_viser.py +++ b/demo_viser.py @@ -30,6 +30,7 @@ from vggt.utils.load_fn import load_and_preprocess_images from vggt.utils.geometry import closed_form_inverse_se3, unproject_depth_map_to_point_map from vggt.utils.pose_enc import pose_encoding_to_extri_intri + def viser_wrapper( pred_dict: dict, port: int = 8080, @@ -214,9 +215,9 @@ def viser_wrapper( # Here we compute the threshold value based on the current percentage current_percentage = gui_points_conf.value threshold_val = np.percentile(conf_flat, current_percentage) - + print(f"Threshold absolute value: {threshold_val}, percentage: {current_percentage}%") - + conf_mask = (conf_flat >= threshold_val) & (conf_flat > 1e-5) if gui_frame_selector.value == "All": @@ -280,51 +281,44 @@ def apply_sky_segmentation(conf: np.ndarray, image_folder: str) -> np.ndarray: np.ndarray: Updated confidence scores with sky regions masked out """ S, H, W = conf.shape - sky_masks_dir = image_folder.rstrip('/') + "_sky_masks" + sky_masks_dir = image_folder.rstrip("/") + "_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_file_from_url( - "https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx", "skyseg.onnx" - ) - - skyseg_session = onnxruntime.InferenceSession("skyseg.onnx") + download_file_from_url("https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx", "skyseg.onnx") + + skyseg_session = onnxruntime.InferenceSession("skyseg.onnx") image_files = sorted(glob.glob(os.path.join(image_folder, "*"))) sky_mask_list = [] - + print("Generating sky masks...") for i, image_path in enumerate(tqdm(image_files[:S])): # Limit to the number of images in the batch image_name = os.path.basename(image_path) mask_filepath = os.path.join(sky_masks_dir, image_name) - + if os.path.exists(mask_filepath): sky_mask = cv2.imread(mask_filepath, cv2.IMREAD_GRAYSCALE) else: sky_mask = segment_sky(image_path, skyseg_session, mask_filepath) - + # 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) conf = conf * sky_mask_binary - + print("Sky segmentation applied successfully") return conf - - - - parser = argparse.ArgumentParser(description="VGGT demo with viser for 3D visualization") parser.add_argument( "--image_folder", type=str, default="examples/kitchen/images/", help="Path to folder containing images" @@ -341,14 +335,14 @@ parser.add_argument("--mask_sky", action="store_true", help="Apply sky segmentat def main(): """ Main function for the VGGT demo with viser for 3D visualization. - + This function: 1. Loads the VGGT model 2. Processes input images from the specified folder 3. Runs inference to generate 3D points and camera poses 4. Optionally applies sky segmentation to filter out sky points 5. Visualizes the results using viser - + Command-line arguments: --image_folder: Path to folder containing input images --use_point_map: Use point map instead of depth-based points @@ -363,7 +357,7 @@ def main(): print("Initializing and loading VGGT model...") # 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)) @@ -379,9 +373,11 @@ def main(): images = load_and_preprocess_images(image_names).to(device) print(f"Preprocessed images shape: {images.shape}") - print("Running 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=torch.bfloat16): + with torch.cuda.amp.autocast(dtype=dtype): predictions = model(images) print("Converting pose encoding to extrinsic and intrinsic matrices...") @@ -417,4 +413,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/requirements.txt b/requirements.txt index e2a6558..1cb61a2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,4 @@ numpy==1.26.1 Pillow huggingface_hub einops +safetensors diff --git a/requirements_demo.txt b/requirements_demo.txt index 56d8198..efd6f5e 100644 --- a/requirements_demo.txt +++ b/requirements_demo.txt @@ -7,4 +7,5 @@ opencv-python scipy onnxruntime requests -trimesh \ No newline at end of file +trimesh +matplotlib diff --git a/vggt/heads/dpt_head.py b/vggt/heads/dpt_head.py index c8c8af9..fa5cf6e 100644 --- a/vggt/heads/dpt_head.py +++ b/vggt/heads/dpt_head.py @@ -206,7 +206,7 @@ class DPTHead(nn.Module): Tensor or Tuple[Tensor, Tensor]: Feature maps or (predictions, confidence). """ if frames_start_idx is not None and frames_end_idx is not None: - images = images[:, frames_start_idx:frames_end_idx] + images = images[:, frames_start_idx:frames_end_idx].contiguous() B, S, _, H, W = images.shape diff --git a/vggt/models/vggt.py b/vggt/models/vggt.py index 1e36818..75587dc 100644 --- a/vggt/models/vggt.py +++ b/vggt/models/vggt.py @@ -38,7 +38,22 @@ class VGGT(nn.Module, PyTorchModelHubMixin): query_points (torch.Tensor, optional): Query points for tracking, in pixel coordinates. Shape: [N, 2] or [B, N, 2], where N is the number of query points. Default: None + + Returns: + dict: A dictionary containing the following predictions: + - pose_enc (torch.Tensor): Camera pose encoding with shape [B, S, 9] (from the last iteration) + - depth (torch.Tensor): Predicted depth maps with shape [B, S, H, W, 1] + - depth_conf (torch.Tensor): Confidence scores for depth predictions with shape [B, S, H, W] + - world_points (torch.Tensor): 3D world coordinates for each pixel with shape [B, S, H, W, 3] + - world_points_conf (torch.Tensor): Confidence scores for world points with shape [B, S, H, W] + - images (torch.Tensor): Original input images, preserved for visualization + + If query_points is provided, also includes: + - track (torch.Tensor): Point tracks with shape [B, S, N, 2] (from the last iteration), in pixel coordinates + - vis (torch.Tensor): Visibility scores for tracked points with shape [B, S, N] + - conf (torch.Tensor): Confidence scores for tracked points with shape [B, S, N] """ + # If without batch dimension, add it if len(images.shape) == 4: images = images.unsqueeze(0) diff --git a/vggt/utils/load_fn.py b/vggt/utils/load_fn.py index 99857c5..eab793d 100644 --- a/vggt/utils/load_fn.py +++ b/vggt/utils/load_fn.py @@ -44,9 +44,9 @@ def load_and_preprocess_images(image_path_list): img = Image.open(image_path) # If there's an alpha channel, blend onto white background: - if img.mode == 'RGBA': + if img.mode == "RGBA": # Create white background - background = Image.new('RGBA', img.size, (255, 255, 255, 255)) + background = Image.new("RGBA", img.size, (255, 255, 255, 255)) # Alpha composite onto the white background img = Image.alpha_composite(background, img) diff --git a/vggt/utils/pose_enc.py b/vggt/utils/pose_enc.py index 93956ec..2f98b08 100644 --- a/vggt/utils/pose_enc.py +++ b/vggt/utils/pose_enc.py @@ -13,9 +13,35 @@ def extri_intri_to_pose_encoding( intrinsics, image_size_hw=None, # e.g., (256, 512) pose_encoding_type="absT_quaR_FoV", - min_focal_length=0.1, - max_focal_length=10, ): + """Convert camera extrinsics and intrinsics to a compact pose encoding. + + This function transforms camera parameters into a unified pose encoding format, + which can be used for various downstream tasks like pose prediction or representation. + + Args: + extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4, + where B is batch size and S is sequence length. + In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world transformation. + The format is [R|t] where R is a 3x3 rotation matrix and t is a 3x1 translation vector. + intrinsics (torch.Tensor): Camera intrinsic parameters with shape BxSx3x3. + Defined in pixels, with format: + [[fx, 0, cx], + [0, fy, cy], + [0, 0, 1]] + where fx, fy are focal lengths and (cx, cy) is the principal point + image_size_hw (tuple): Tuple of (height, width) of the image in pixels. + Required for computing field of view values. For example: (256, 512). + pose_encoding_type (str): Type of pose encoding to use. Currently only + supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view). + + Returns: + torch.Tensor: Encoded camera pose parameters with shape BxSx9. + For "absT_quaR_FoV" type, the 9 dimensions are: + - [:3] = absolute translation vector T (3D) + - [3:7] = rotation as quaternion quat (4D) + - [7:] = field of view (2D) + """ # extrinsics: BxSx3x4 # intrinsics: BxSx3x3 @@ -25,7 +51,6 @@ def extri_intri_to_pose_encoding( T = extrinsics[:, :, :3, 3] # BxSx3 quat = mat_to_quat(R) - # R_reverse = quat_to_mat(quat) # Note the order of h and w here H, W = image_size_hw fov_h = 2 * torch.atan((H / 2) / intrinsics[..., 1, 1]) @@ -40,11 +65,43 @@ def extri_intri_to_pose_encoding( def pose_encoding_to_extri_intri( pose_encoding, image_size_hw=None, # e.g., (256, 512) - min_focal_length=0.1, - max_focal_length=10, pose_encoding_type="absT_quaR_FoV", build_intrinsics=True, ): + """Convert a pose encoding back to camera extrinsics and intrinsics. + + This function performs the inverse operation of extri_intri_to_pose_encoding, + reconstructing the full camera parameters from the compact encoding. + + Args: + pose_encoding (torch.Tensor): Encoded camera pose parameters with shape BxSx9, + where B is batch size and S is sequence length. + For "absT_quaR_FoV" type, the 9 dimensions are: + - [:3] = absolute translation vector T (3D) + - [3:7] = rotation as quaternion quat (4D) + - [7:] = field of view (2D) + image_size_hw (tuple): Tuple of (height, width) of the image in pixels. + Required for reconstructing intrinsics from field of view values. + For example: (256, 512). + pose_encoding_type (str): Type of pose encoding used. Currently only + supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view). + build_intrinsics (bool): Whether to reconstruct the intrinsics matrix. + If False, only extrinsics are returned and intrinsics will be None. + + Returns: + tuple: (extrinsics, intrinsics) + - extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4. + In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world + transformation. The format is [R|t] where R is a 3x3 rotation matrix and t is + a 3x1 translation vector. + - intrinsics (torch.Tensor or None): Camera intrinsic parameters with shape BxSx3x3, + or None if build_intrinsics is False. Defined in pixels, with format: + [[fx, 0, cx], + [0, fy, cy], + [0, 0, 1]] + where fx, fy are focal lengths and (cx, cy) is the principal point, + assumed to be at the center of the image (W/2, H/2). + """ intrinsics = None diff --git a/visual_util.py b/visual_util.py index 074d2a8..3c624c9 100644 --- a/visual_util.py +++ b/visual_util.py @@ -131,7 +131,7 @@ def predictions_to_glb( 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) + sky_mask_binary = (sky_mask_array > 0.1).astype(np.float32) pred_world_points_conf = pred_world_points_conf * sky_mask_binary if selected_frame_idx is not None: @@ -160,7 +160,7 @@ def predictions_to_glb( if mask_black_bg: black_bg_mask = colors_rgb.sum(axis=1) >= 16 conf_mask = conf_mask & black_bg_mask - + if mask_white_bg: # Filter out white background pixels (RGB values close to white) # Consider pixels white if all RGB values are above 240 @@ -392,7 +392,7 @@ def segment_sky(image_path, onnx_session, mask_filename=None): # 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