🐙 GitHub Repository | @@ -542,7 +542,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 +562,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..eef6fc8 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)) @@ -417,4 +411,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() 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