[WIP] Upload demo_colmap.py, which can (a) run both vggt and vggt+ba, (b) output in a colmap format, (c) in the original image resolution. The output can be directly used for nerf/gaussian training.

.gitignore

@@ -1,7 +1,6 @@
 .hydra/
 output/
 ckpt/
-dependency/
 # Byte-compiled / optimized / DLL files
 __pycache__/
 **/__pycache__/

demo_colmap.py

@@ -0,0 +1,327 @@
+# 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 random
+import numpy as np
+import glob
+import os
+import copy
+import torch
+import torch.nn.functional as F
+
+# Configure CUDA settings
+torch.backends.cudnn.enabled = True
+torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.deterministic = False
+
+import argparse
+from pathlib import Path
+import trimesh
+import pycolmap
+
+
+from vggt.models.vggt import VGGT
+from vggt.utils.load_fn import load_and_preprocess_images_square
+from vggt.utils.pose_enc import pose_encoding_to_extri_intri
+from vggt.utils.geometry import unproject_depth_map_to_point_map
+from vggt.utils.helper import create_pixel_coordinate_grid, randomly_limit_trues
+from vggt.dependency.track_predict import predict_tracks
+from vggt.dependency.np_to_pycolmap import batch_np_matrix_to_pycolmap, batch_np_matrix_to_pycolmap_wo_track
+
+
+# TODO: add support for masks
+# TODO: add iterative BA
+# TODO: add support for radial distortion, which needs extra_params
+# TODO: test with more cases
+# TODO: test different camera types
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="VGGT Demo")
+    parser.add_argument("--scene_dir", type=str, required=True, help="Directory containing the scene images")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
+    parser.add_argument("--use_ba", action="store_true", default=False, help="Use BA for reconstruction")
+    ######### BA parameters #########
+    parser.add_argument(
+        "--max_reproj_error", type=float, default=8.0, help="Maximum reprojection error for reconstruction"
+    )
+    parser.add_argument("--shared_camera", action="store_true", default=False, help="Use shared camera for all images")
+    parser.add_argument("--camera_type", type=str, default="SIMPLE_PINHOLE", help="Camera type for reconstruction")
+    parser.add_argument("--vis_thresh", type=float, default=0.2, help="Visibility threshold for tracks")
+    parser.add_argument("--query_frame_num", type=int, default=5, help="Number of frames to query")
+    parser.add_argument("--max_query_pts", type=int, default=2048, help="Maximum number of query points")
+    parser.add_argument(
+        "--fine_tracking", action="store_true", default=True, help="Use fine tracking (slower but more accurate)"
+    )
+    return parser.parse_args()
+
+
+def run_VGGT(model, images, dtype, resolution=518):
+    # images: [B, 3, H, W]
+
+    assert len(images.shape) == 4
+    assert images.shape[1] == 3
+
+    # hard-coded to use 518 for VGGT
+    images = F.interpolate(images, size=(resolution, resolution), mode="bilinear", align_corners=False)
+
+    with torch.no_grad():
+        with torch.cuda.amp.autocast(dtype=dtype):
+            images = images[None]  # add batch dimension
+            aggregated_tokens_list, ps_idx = model.aggregator(images)
+
+        # Predict Cameras
+        pose_enc = model.camera_head(aggregated_tokens_list)[-1]
+        # Extrinsic and intrinsic matrices, following OpenCV convention (camera from world)
+        extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc, images.shape[-2:])
+        # Predict Depth Maps
+        depth_map, depth_conf = model.depth_head(aggregated_tokens_list, images, ps_idx)
+
+    extrinsic = extrinsic.squeeze(0).cpu().numpy()
+    intrinsic = intrinsic.squeeze(0).cpu().numpy()
+    depth_map = depth_map.squeeze(0).cpu().numpy()
+    depth_conf = depth_conf.squeeze(0).cpu().numpy()
+    return extrinsic, intrinsic, depth_map, depth_conf
+
+
+def demo_fn(args):
+    # Print configuration
+    print("Arguments:", vars(args))
+
+    # Set seed for reproducibility
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    random.seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(args.seed)
+        torch.cuda.manual_seed_all(args.seed)  # for multi-GPU
+    print(f"Setting seed as: {args.seed}")
+
+    # Set device and dtype
+    dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+    print(f"Using dtype: {dtype}")
+
+    # Run VGGT for camera and depth estimation
+    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)
+    print(f"Model loaded")
+
+    # Get image paths and preprocess them
+    image_dir = os.path.join(args.scene_dir, "images")
+    image_path_list = glob.glob(os.path.join(image_dir, "*"))
+    if len(image_path_list) == 0:
+        raise ValueError(f"No images found in {image_dir}")
+    base_image_path_list = [os.path.basename(path) for path in image_path_list]
+
+    # Load images and original coordinates
+    # Load Image in 1024, while running VGGT with 518
+    vggt_fixed_resolution = 518
+    img_load_resolution = 1024
+
+    images, original_coords = load_and_preprocess_images_square(image_path_list, img_load_resolution)
+    images = images.to(device)
+    original_coords = original_coords.to(device)
+    print(f"Loaded {len(images)} images from {image_dir}")
+
+    # Run VGGT to estimate camera and depth
+    # Run with 518x518 images
+    extrinsic, intrinsic, depth_map, depth_conf = run_VGGT(model, images, dtype, vggt_fixed_resolution)
+    points_3d = unproject_depth_map_to_point_map(depth_map, extrinsic, intrinsic)
+
+    if args.use_ba:
+        image_size = np.array(images.shape[-2:])
+        scale = img_load_resolution / vggt_fixed_resolution
+        shared_camera = args.shared_camera
+
+        with torch.cuda.amp.autocast(dtype=dtype):
+            # Predicting Tracks
+            # Using VGGSfM tracker instead of VGGT tracker for efficiency
+            # VGGT tracker requires multiple backbone runs to query different frames (this is a problem caused by the training process)
+            # Will be fixed in VGGT v2
+
+            # You can also change the pred_tracks to tracks from any other methods
+            # e.g., from COLMAP, from CoTracker, or by chaining 2D matches from Lightglue/LoFTR.
+            pred_tracks, pred_vis_scores, pred_confs, points_3d, points_rgb = predict_tracks(
+                images,
+                conf=depth_conf,
+                points_3d=points_3d,
+                masks=None,
+                max_query_pts=args.max_query_pts,
+                query_frame_num=args.query_frame_num,
+                keypoint_extractor="aliked+sp",
+                fine_tracking=args.fine_tracking,
+            )
+
+            torch.cuda.empty_cache()
+
+        # rescale the intrinsic matrix from 518 to 1024
+        intrinsic[:, :2, :] *= scale
+        track_mask = pred_vis_scores > args.vis_thresh
+
+        # TODO: radial distortion, iterative BA, masks
+        reconstruction, valid_track_mask = batch_np_matrix_to_pycolmap(
+            points_3d,
+            extrinsic,
+            intrinsic,
+            pred_tracks,
+            image_size,
+            masks=track_mask,
+            max_reproj_error=args.max_reproj_error,
+            shared_camera=shared_camera,
+            camera_type=args.camera_type,
+            points_rgb=points_rgb,
+        )
+
+        if reconstruction is None:
+            raise ValueError("No reconstruction can be built with BA")
+
+        # Bundle Adjustment
+        ba_options = pycolmap.BundleAdjustmentOptions()
+        pycolmap.bundle_adjustment(reconstruction, ba_options)
+
+        reconstruction_resolution = img_load_resolution
+    else:
+        conf_thres_value = 5  # hard-coded to 5
+        max_points_for_colmap = 100000  # randomly sample 3D points
+        shared_camera = False  # in the feedforward manner, we do not support shared camera
+        camera_type = "PINHOLE"  # in the feedforward manner, we only support PINHOLE camera
+
+        image_size = np.array([vggt_fixed_resolution, vggt_fixed_resolution])
+        num_frames, height, width, _ = points_3d.shape
+
+        points_rgb = F.interpolate(
+            images, size=(vggt_fixed_resolution, vggt_fixed_resolution), mode="bilinear", align_corners=False
+        )
+        points_rgb = (points_rgb.cpu().numpy() * 255).astype(np.uint8)
+        points_rgb = points_rgb.transpose(0, 2, 3, 1)
+
+        # (S, H, W, 3), with x, y coordinates and frame indices
+        points_xyf = create_pixel_coordinate_grid(num_frames, height, width)
+
+        conf_mask = depth_conf >= conf_thres_value
+        # at most writing 100000 3d points to colmap reconstruction object
+        conf_mask = randomly_limit_trues(conf_mask, max_points_for_colmap)
+
+        points_3d = points_3d[conf_mask]
+        points_xyf = points_xyf[conf_mask]
+        points_rgb = points_rgb[conf_mask]
+
+        print("Converting to COLMAP format")
+        reconstruction = batch_np_matrix_to_pycolmap_wo_track(
+            points_3d,
+            points_xyf,
+            points_rgb,
+            extrinsic,
+            intrinsic,
+            image_size,
+            shared_camera=shared_camera,
+            camera_type=camera_type,
+        )
+
+        reconstruction_resolution = vggt_fixed_resolution
+
+    # Save point cloud for fast visualization
+    trimesh.PointCloud(points_3d, colors=points_rgb).export(os.path.join(args.scene_dir, "sparse/points.ply"))
+
+    reconstruction = rename_colmap_recons_and_rescale_camera(
+        reconstruction,
+        base_image_path_list,
+        original_coords.cpu().numpy(),
+        img_size=reconstruction_resolution,
+        shift_point2d_to_original_res=True,
+        shared_camera=shared_camera,
+    )
+
+    print(f"Saving reconstruction to {args.scene_dir}/sparse")
+    sparse_reconstruction_dir = os.path.join(args.scene_dir, "sparse")
+    os.makedirs(sparse_reconstruction_dir, exist_ok=True)
+    reconstruction.write(sparse_reconstruction_dir)
+
+    return True
+
+
+def rename_colmap_recons_and_rescale_camera(
+    reconstruction, image_paths, original_coords, img_size, shift_point2d_to_original_res=False, shared_camera=False
+):
+    rescale_camera = True
+
+    for pyimageid in reconstruction.images:
+        # Reshaped the padded&resized image to the original size
+        # Rename the images to the original names
+        pyimage = reconstruction.images[pyimageid]
+        pycamera = reconstruction.cameras[pyimage.camera_id]
+        pyimage.name = image_paths[pyimageid - 1]
+
+        if rescale_camera:
+            # Rescale the camera parameters
+            pred_params = copy.deepcopy(pycamera.params)
+
+            real_image_size = original_coords[pyimageid - 1, -2:]
+            resize_ratio = max(real_image_size) / img_size
+            pred_params = pred_params * resize_ratio
+            real_pp = real_image_size / 2
+            pred_params[-2:] = real_pp  # center of the image
+
+            pycamera.params = pred_params
+            pycamera.width = real_image_size[0]
+            pycamera.height = real_image_size[1]
+
+        if shift_point2d_to_original_res:
+            # Also shift the point2D to original resolution
+            top_left = original_coords[pyimageid - 1, :2]
+
+            for point2D in pyimage.points2D:
+                point2D.xy = (point2D.xy - top_left) * resize_ratio
+
+        if shared_camera:
+            # If shared_camera, all images share the same camera
+            # no need to rescale any more
+            rescale_camera = False
+
+    return reconstruction
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    with torch.no_grad():
+        demo_fn(args)
+
+
+# Work in Progress (WIP)
+
+"""
+VGGT Runner Script
+=================
+
+A script to run the VGGT model for 3D reconstruction from image sequences.
+
+Directory Structure
+------------------
+Input:
+    input_folder/
+    └── images/            # Source images for reconstruction
+
+Output:
+    output_folder/
+    ├── images/
+    ├── sparse/           # Reconstruction results
+    │   ├── cameras.bin   # Camera parameters (COLMAP format)
+    │   ├── images.bin    # Pose for each image (COLMAP format)
+    │   ├── points3D.bin  # 3D points (COLMAP format)
+    │   └── points.ply    # Point cloud visualization file 
+    └── visuals/          # Visualization outputs TODO
+
+Key Features
+-----------
+• Dual-mode Support: Run reconstructions using either VGGT or VGGT+BA
+• Resolution Preservation: Maintains original image resolution in camera parameters and tracks
+• COLMAP Compatibility: Exports results in standard COLMAP sparse reconstruction format
+"""

demo_gradio.py

@@ -289,21 +289,20 @@ def update_visualization(
         return None, f"No reconstruction available at {predictions_path}. Please run 'Reconstruct' first."
 
     key_list = [
-        'pose_enc',
-        'depth',
-        'depth_conf',
-        'world_points',
-        'world_points_conf',
-        'images',
-        'extrinsic',
-        'intrinsic',
-        'world_points_from_depth',
+        "pose_enc",
+        "depth",
+        "depth_conf",
+        "world_points",
+        "world_points_conf",
+        "images",
+        "extrinsic",
+        "intrinsic",
+        "world_points_from_depth",
     ]
 
     loaded = np.load(predictions_path)
     predictions = {key: np.array(loaded[key]) for key in key_list}
 
-    
     glbfile = os.path.join(
         target_dir,
         f"glbscene_{conf_thres}_{frame_filter.replace('.', '_').replace(':', '').replace(' ', '_')}_maskb{mask_black_bg}_maskw{mask_white_bg}_cam{show_cam}_sky{mask_sky}_pred{prediction_mode.replace(' ', '_')}.glb",
@@ -390,7 +389,6 @@ with gr.Blocks(
     }
     """,
 ) as demo:
-
     # 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")
@@ -534,13 +532,7 @@ with gr.Blocks(
             prediction_mode,
             is_example,
         ],
-        outputs=[
-            reconstruction_output,
-            log_output,
-            target_dir_output,
-            frame_filter,
-            image_gallery,
-        ],
+        outputs=[reconstruction_output, log_output, target_dir_output, frame_filter, image_gallery],
         fn=example_pipeline,
         cache_examples=False,
         examples_per_page=50,

demo_viser.py

@@ -112,24 +112,15 @@ def viser_wrapper(
     frame_indices = np.repeat(np.arange(S), H * W)
 
     # Build the viser GUI
-    gui_show_frames = server.gui.add_checkbox(
-        "Show Cameras",
-        initial_value=True,
-    )
+    gui_show_frames = server.gui.add_checkbox("Show Cameras", initial_value=True)
 
     # Now the slider represents percentage of points to filter out
     gui_points_conf = server.gui.add_slider(
-        "Confidence Percent",
-        min=0,
-        max=100,
-        step=0.1,
-        initial_value=init_conf_threshold,
+        "Confidence Percent", min=0, max=100, step=0.1, initial_value=init_conf_threshold
     )
 
     gui_frame_selector = server.gui.add_dropdown(
-        "Show Points from Frames",
-        options=["All"] + [str(i) for i in range(S)],
-        initial_value="All",
+        "Show Points from Frames", options=["All"] + [str(i) for i in range(S)], initial_value="All"
     )
 
     # Create the main point cloud handle
@@ -200,12 +191,7 @@ def viser_wrapper(
 
             # Add the frustum
             frustum_cam = server.scene.add_camera_frustum(
-                f"frame_{img_id}/frustum",
-                fov=fov,
-                aspect=w / h,
-                scale=0.05,
-                image=img,
-                line_width=1.0,
+                f"frame_{img_id}/frustum", fov=fov, aspect=w / h, scale=0.05, image=img, line_width=1.0
             )
             frustums.append(frustum_cam)
             attach_callback(frustum_cam, frame_axis)
@@ -373,7 +359,7 @@ 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():

requirements_demo.txt

@@ -9,3 +9,7 @@ onnxruntime
 requests
 trimesh
 matplotlib
+pycolmap==3.10.0 
+pyceres==2.3
+git+https://github.com/jytime/LightGlue.git#egg=lightglue
+

vggt/dependency/__init__.py

@@ -0,0 +1,3 @@
+from .track_modules.track_refine import refine_track
+from .track_modules.blocks import BasicEncoder, ShallowEncoder
+from .track_modules.base_track_predictor import BaseTrackerPredictor

vggt/dependency/distortion.py

@@ -0,0 +1,182 @@
+# 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 torch
+import numpy as np
+from typing import Union
+
+ArrayLike = Union[np.ndarray, torch.Tensor]
+
+
+def _is_numpy(x: ArrayLike) -> bool:
+    return isinstance(x, np.ndarray)
+
+
+def _is_torch(x: ArrayLike) -> bool:
+    return isinstance(x, torch.Tensor)
+
+
+def _ensure_torch(x: ArrayLike) -> torch.Tensor:
+    """Convert input to torch tensor if it's not already one."""
+    if _is_numpy(x):
+        return torch.from_numpy(x)
+    elif _is_torch(x):
+        return x
+    else:
+        return torch.tensor(x)
+
+
+def single_undistortion(params, tracks_normalized):
+    """
+    Apply undistortion to the normalized tracks using the given distortion parameters once.
+
+    Args:
+        params (torch.Tensor or numpy.ndarray): Distortion parameters of shape BxN.
+        tracks_normalized (torch.Tensor or numpy.ndarray): Normalized tracks tensor of shape [batch_size, num_tracks, 2].
+
+    Returns:
+        torch.Tensor: Undistorted normalized tracks tensor.
+    """
+    params = _ensure_torch(params)
+    tracks_normalized = _ensure_torch(tracks_normalized)
+
+    u, v = tracks_normalized[..., 0].clone(), tracks_normalized[..., 1].clone()
+    u_undist, v_undist = apply_distortion(params, u, v)
+    return torch.stack([u_undist, v_undist], dim=-1)
+
+
+def iterative_undistortion(params, tracks_normalized, max_iterations=100, max_step_norm=1e-10, rel_step_size=1e-6):
+    """
+    Iteratively undistort the normalized tracks using the given distortion parameters.
+
+    Args:
+        params (torch.Tensor or numpy.ndarray): Distortion parameters of shape BxN.
+        tracks_normalized (torch.Tensor or numpy.ndarray): Normalized tracks tensor of shape [batch_size, num_tracks, 2].
+        max_iterations (int): Maximum number of iterations for the undistortion process.
+        max_step_norm (float): Maximum step norm for convergence.
+        rel_step_size (float): Relative step size for numerical differentiation.
+
+    Returns:
+        torch.Tensor: Undistorted normalized tracks tensor.
+    """
+    params = _ensure_torch(params)
+    tracks_normalized = _ensure_torch(tracks_normalized)
+
+    B, N, _ = tracks_normalized.shape
+    u, v = tracks_normalized[..., 0].clone(), tracks_normalized[..., 1].clone()
+    original_u, original_v = u.clone(), v.clone()
+
+    eps = torch.finfo(u.dtype).eps
+    for idx in range(max_iterations):
+        u_undist, v_undist = apply_distortion(params, u, v)
+        dx = original_u - u_undist
+        dy = original_v - v_undist
+
+        step_u = torch.clamp(torch.abs(u) * rel_step_size, min=eps)
+        step_v = torch.clamp(torch.abs(v) * rel_step_size, min=eps)
+
+        J_00 = (apply_distortion(params, u + step_u, v)[0] - apply_distortion(params, u - step_u, v)[0]) / (2 * step_u)
+        J_01 = (apply_distortion(params, u, v + step_v)[0] - apply_distortion(params, u, v - step_v)[0]) / (2 * step_v)
+        J_10 = (apply_distortion(params, u + step_u, v)[1] - apply_distortion(params, u - step_u, v)[1]) / (2 * step_u)
+        J_11 = (apply_distortion(params, u, v + step_v)[1] - apply_distortion(params, u, v - step_v)[1]) / (2 * step_v)
+
+        J = torch.stack([torch.stack([J_00 + 1, J_01], dim=-1), torch.stack([J_10, J_11 + 1], dim=-1)], dim=-2)
+
+        delta = torch.linalg.solve(J, torch.stack([dx, dy], dim=-1))
+
+        u += delta[..., 0]
+        v += delta[..., 1]
+
+        if torch.max((delta**2).sum(dim=-1)) < max_step_norm:
+            break
+
+    return torch.stack([u, v], dim=-1)
+
+
+def apply_distortion(extra_params, u, v):
+    """
+    Applies radial or OpenCV distortion to the given 2D points.
+
+    Args:
+        extra_params (torch.Tensor or numpy.ndarray): Distortion parameters of shape BxN, where N can be 1, 2, or 4.
+        u (torch.Tensor or numpy.ndarray): Normalized x coordinates of shape Bxnum_tracks.
+        v (torch.Tensor or numpy.ndarray): Normalized y coordinates of shape Bxnum_tracks.
+
+    Returns:
+        points2D (torch.Tensor): Distorted 2D points of shape BxNx2.
+    """
+    extra_params = _ensure_torch(extra_params)
+    u = _ensure_torch(u)
+    v = _ensure_torch(v)
+
+    num_params = extra_params.shape[1]
+
+    if num_params == 1:
+        # Simple radial distortion
+        k = extra_params[:, 0]
+        u2 = u * u
+        v2 = v * v
+        r2 = u2 + v2
+        radial = k[:, None] * r2
+        du = u * radial
+        dv = v * radial
+
+    elif num_params == 2:
+        # RadialCameraModel distortion
+        k1, k2 = extra_params[:, 0], extra_params[:, 1]
+        u2 = u * u
+        v2 = v * v
+        r2 = u2 + v2
+        radial = k1[:, None] * r2 + k2[:, None] * r2 * r2
+        du = u * radial
+        dv = v * radial
+
+    elif num_params == 4:
+        # OpenCVCameraModel distortion
+        k1, k2, p1, p2 = (extra_params[:, 0], extra_params[:, 1], extra_params[:, 2], extra_params[:, 3])
+        u2 = u * u
+        v2 = v * v
+        uv = u * v
+        r2 = u2 + v2
+        radial = k1[:, None] * r2 + k2[:, None] * r2 * r2
+        du = u * radial + 2 * p1[:, None] * uv + p2[:, None] * (r2 + 2 * u2)
+        dv = v * radial + 2 * p2[:, None] * uv + p1[:, None] * (r2 + 2 * v2)
+    else:
+        raise ValueError("Unsupported number of distortion parameters")
+
+    u = u.clone() + du
+    v = v.clone() + dv
+
+    return u, v
+
+
+if __name__ == "__main__":
+    import random
+    import pycolmap
+
+    max_diff = 0
+    for i in range(1000):
+        # Define distortion parameters (assuming 1 parameter for simplicity)
+        B = random.randint(1, 500)
+        track_num = random.randint(100, 1000)
+        params = torch.rand((B, 1), dtype=torch.float32)  # Batch size 1, 4 parameters
+        tracks_normalized = torch.rand((B, track_num, 2), dtype=torch.float32)  # Batch size 1, 5 points
+
+        # Undistort the tracks
+        undistorted_tracks = iterative_undistortion(params, tracks_normalized)
+
+        for b in range(B):
+            pycolmap_intri = np.array([1, 0, 0, params[b].item()])
+            pycam = pycolmap.Camera(model="SIMPLE_RADIAL", width=1, height=1, params=pycolmap_intri, camera_id=0)
+
+            undistorted_tracks_pycolmap = pycam.cam_from_img(tracks_normalized[b].numpy())
+            diff = (undistorted_tracks[b] - undistorted_tracks_pycolmap).abs().median()
+            max_diff = max(max_diff, diff)
+            print(f"diff: {diff}, max_diff: {max_diff}")
+
+    import pdb
+
+    pdb.set_trace()

vggt/dependency/np_to_pycolmap.py

@@ -0,0 +1,318 @@
+# 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 numpy as np
+import pycolmap
+from .projection import project_3D_points_np
+
+
+def batch_np_matrix_to_pycolmap(
+    points3d,
+    extrinsics,
+    intrinsics,
+    tracks,
+    image_size,
+    masks=None,
+    max_reproj_error=None,
+    max_points3D_val=3000,
+    shared_camera=False,
+    camera_type="SIMPLE_PINHOLE",
+    extra_params=None,
+    min_inlier_per_frame=64,
+    points_rgb=None,
+):
+    """
+    Convert Batched NumPy Arrays to PyCOLMAP
+
+    Check https://github.com/colmap/pycolmap for more details about its format
+
+    NOTE that colmap expects images/cameras/points3D to be 1-indexed
+    so there is a +1 offset between colmap index and batch index
+
+
+    NOTE: different from VGGSfM, this function:
+    1. Use np instead of torch
+    2. Frame index and camera id starts from 1 rather than 0 (to fit the format of PyCOLMAP)
+    """
+    # points3d: Px3
+    # extrinsics: Nx3x4
+    # intrinsics: Nx3x3
+    # tracks: NxPx2
+    # masks: NxP
+    # image_size: 2, assume all the frames have been padded to the same size
+    # where N is the number of frames and P is the number of tracks
+
+    N, P, _ = tracks.shape
+    assert len(extrinsics) == N
+    assert len(intrinsics) == N
+    assert len(points3d) == P
+    assert image_size.shape[0] == 2
+
+    if max_reproj_error is not None:
+        projected_points_2d, projected_points_cam = project_3D_points_np(points3d, extrinsics, intrinsics)
+        projected_diff = np.linalg.norm(projected_points_2d - tracks, axis=-1)
+        projected_points_2d[projected_points_cam[:, -1] <= 0] = 1e6
+        reproj_mask = projected_diff < max_reproj_error
+
+    if masks is not None and reproj_mask is not None:
+        masks = np.logical_and(masks, reproj_mask)
+    elif masks is not None:
+        masks = masks
+    else:
+        masks = reproj_mask
+
+    assert masks is not None
+
+    if masks.sum(1).min() < min_inlier_per_frame:
+        print(f"Not enough inliers per frame, skip BA.")
+        return None, None
+
+    # Reconstruction object, following the format of PyCOLMAP/COLMAP
+    reconstruction = pycolmap.Reconstruction()
+
+    inlier_num = masks.sum(0)
+    valid_mask = inlier_num >= 2  # a track is invalid if without two inliers
+    valid_idx = np.nonzero(valid_mask)[0]
+
+    # Only add 3D points that have sufficient 2D points
+    for vidx in valid_idx:
+        # Use RGB colors if provided, otherwise use zeros
+        rgb = points_rgb[vidx] if points_rgb is not None else np.zeros(3)
+        reconstruction.add_point3D(points3d[vidx], pycolmap.Track(), rgb)
+
+    num_points3D = len(valid_idx)
+    camera = None
+    # frame idx
+    for fidx in range(N):
+        # set camera
+        if camera is None or (not shared_camera):
+            pycolmap_intri = _build_pycolmap_intri(fidx, intrinsics, camera_type, extra_params)
+
+            camera = pycolmap.Camera(
+                model=camera_type, width=image_size[0], height=image_size[1], params=pycolmap_intri, camera_id=fidx + 1
+            )
+
+            # add camera
+            reconstruction.add_camera(camera)
+
+        # set image
+        cam_from_world = pycolmap.Rigid3d(
+            pycolmap.Rotation3d(extrinsics[fidx][:3, :3]), extrinsics[fidx][:3, 3]
+        )  # Rot and Trans
+
+        image = pycolmap.Image(
+            id=fidx + 1, name=f"image_{fidx + 1}", camera_id=camera.camera_id, cam_from_world=cam_from_world
+        )
+
+        points2D_list = []
+
+        point2D_idx = 0
+
+        # NOTE point3D_id start by 1
+        for point3D_id in range(1, num_points3D + 1):
+            original_track_idx = valid_idx[point3D_id - 1]
+
+            if (reconstruction.points3D[point3D_id].xyz < max_points3D_val).all():
+                if masks[fidx][original_track_idx]:
+                    # It seems we don't need +0.5 for BA
+                    point2D_xy = tracks[fidx][original_track_idx]
+                    # Please note when adding the Point2D object
+                    # It not only requires the 2D xy location, but also the id to 3D point
+                    points2D_list.append(pycolmap.Point2D(point2D_xy, point3D_id))
+
+                    # add element
+                    track = reconstruction.points3D[point3D_id].track
+                    track.add_element(fidx + 1, point2D_idx)
+                    point2D_idx += 1
+
+        assert point2D_idx == len(points2D_list)
+
+        try:
+            image.points2D = pycolmap.ListPoint2D(points2D_list)
+            image.registered = True
+        except:
+            print(f"frame {fidx + 1} is out of BA")
+            image.registered = False
+
+        # add image
+        reconstruction.add_image(image)
+
+    return reconstruction, valid_mask
+
+
+def pycolmap_to_batch_np_matrix(reconstruction, device="cpu", camera_type="SIMPLE_PINHOLE"):
+    """
+    Convert a PyCOLMAP Reconstruction Object to batched NumPy arrays.
+
+    Args:
+        reconstruction (pycolmap.Reconstruction): The reconstruction object from PyCOLMAP.
+        device (str): Ignored in NumPy version (kept for API compatibility).
+        camera_type (str): The type of camera model used (default: "SIMPLE_PINHOLE").
+
+    Returns:
+        tuple: A tuple containing points3D, extrinsics, intrinsics, and optionally extra_params.
+    """
+
+    num_images = len(reconstruction.images)
+    max_points3D_id = max(reconstruction.point3D_ids())
+    points3D = np.zeros((max_points3D_id, 3))
+
+    for point3D_id in reconstruction.points3D:
+        points3D[point3D_id - 1] = reconstruction.points3D[point3D_id].xyz
+
+    extrinsics = []
+    intrinsics = []
+
+    extra_params = [] if camera_type == "SIMPLE_RADIAL" else None
+
+    for i in range(num_images):
+        # Extract and append extrinsics
+        pyimg = reconstruction.images[i + 1]
+        pycam = reconstruction.cameras[pyimg.camera_id]
+        matrix = pyimg.cam_from_world.matrix()
+        extrinsics.append(matrix)
+
+        # Extract and append intrinsics
+        calibration_matrix = pycam.calibration_matrix()
+        intrinsics.append(calibration_matrix)
+
+        if camera_type == "SIMPLE_RADIAL":
+            extra_params.append(pycam.params[-1])
+
+    # Convert lists to NumPy arrays instead of torch tensors
+    extrinsics = np.stack(extrinsics)
+    intrinsics = np.stack(intrinsics)
+
+    if camera_type == "SIMPLE_RADIAL":
+        extra_params = np.stack(extra_params)
+        extra_params = extra_params[:, None]
+
+    return points3D, extrinsics, intrinsics, extra_params
+
+
+########################################################
+
+
+def batch_np_matrix_to_pycolmap_wo_track(
+    points3d,
+    points_xyf,
+    points_rgb,
+    extrinsics,
+    intrinsics,
+    image_size,
+    shared_camera=False,
+    camera_type="SIMPLE_PINHOLE",
+):
+    """
+    Convert Batched NumPy Arrays to PyCOLMAP
+
+    Different from batch_np_matrix_to_pycolmap, this function does not use tracks.
+
+    It saves points3d to colmap reconstruction format only to serve as init for Gaussians or other nvs methods.
+
+    Do NOT use this for BA.
+    """
+    # points3d: Px3
+    # points_xyf: Px3, with x, y coordinates and frame indices
+    # points_rgb: Px3, rgb colors
+    # extrinsics: Nx3x4
+    # intrinsics: Nx3x3
+    # image_size: 2, assume all the frames have been padded to the same size
+    # where N is the number of frames and P is the number of tracks
+
+    N = len(extrinsics)
+    P = len(points3d)
+
+    # Reconstruction object, following the format of PyCOLMAP/COLMAP
+    reconstruction = pycolmap.Reconstruction()
+
+    for vidx in range(P):
+        reconstruction.add_point3D(points3d[vidx], pycolmap.Track(), points_rgb[vidx])
+
+    camera = None
+    # frame idx
+    for fidx in range(N):
+        # set camera
+        if camera is None or (not shared_camera):
+            pycolmap_intri = _build_pycolmap_intri(fidx, intrinsics, camera_type)
+
+            camera = pycolmap.Camera(
+                model=camera_type, width=image_size[0], height=image_size[1], params=pycolmap_intri, camera_id=fidx + 1
+            )
+
+            # add camera
+            reconstruction.add_camera(camera)
+
+        # set image
+        cam_from_world = pycolmap.Rigid3d(
+            pycolmap.Rotation3d(extrinsics[fidx][:3, :3]), extrinsics[fidx][:3, 3]
+        )  # Rot and Trans
+
+        image = pycolmap.Image(
+            id=fidx + 1, name=f"image_{fidx + 1}", camera_id=camera.camera_id, cam_from_world=cam_from_world
+        )
+
+        points2D_list = []
+
+        point2D_idx = 0
+
+        points_belong_to_fidx = points_xyf[:, 2].astype(np.int32) == fidx
+        points_belong_to_fidx = np.nonzero(points_belong_to_fidx)[0]
+
+        for point3D_batch_idx in points_belong_to_fidx:
+            point3D_id = point3D_batch_idx + 1
+            point2D_xyf = points_xyf[point3D_batch_idx]
+            point2D_xy = point2D_xyf[:2]
+            points2D_list.append(pycolmap.Point2D(point2D_xy, point3D_id))
+
+            # add element
+            track = reconstruction.points3D[point3D_id].track
+            track.add_element(fidx + 1, point2D_idx)
+            point2D_idx += 1
+
+        assert point2D_idx == len(points2D_list)
+
+        try:
+            image.points2D = pycolmap.ListPoint2D(points2D_list)
+            image.registered = True
+        except:
+            print(f"frame {fidx + 1} does not have any points")
+            image.registered = False
+
+        # add image
+        reconstruction.add_image(image)
+
+    return reconstruction
+
+
+def _build_pycolmap_intri(fidx, intrinsics, camera_type, extra_params=None):
+    """
+    Helper function to get camera parameters based on camera type.
+
+    Args:
+        fidx: Frame index
+        intrinsics: Camera intrinsic parameters
+        camera_type: Type of camera model
+        extra_params: Additional parameters for certain camera types
+
+    Returns:
+        pycolmap_intri: NumPy array of camera parameters
+    """
+    if camera_type == "PINHOLE":
+        pycolmap_intri = np.array(
+            [intrinsics[fidx][0, 0], intrinsics[fidx][1, 1], intrinsics[fidx][0, 2], intrinsics[fidx][1, 2]]
+        )
+    elif camera_type == "SIMPLE_PINHOLE":
+        focal = (intrinsics[fidx][0, 0] + intrinsics[fidx][1, 1]) / 2
+        pycolmap_intri = np.array([focal, intrinsics[fidx][0, 2], intrinsics[fidx][1, 2]])
+    elif camera_type == "SIMPLE_RADIAL":
+        raise NotImplementedError("SIMPLE_RADIAL is not supported yet")
+        focal = (intrinsics[fidx][0, 0] + intrinsics[fidx][1, 1]) / 2
+        pycolmap_intri = np.array([focal, intrinsics[fidx][0, 2], intrinsics[fidx][1, 2], extra_params[fidx][0]])
+    else:
+        raise ValueError(f"Camera type {camera_type} is not supported yet")
+
+    return pycolmap_intri

vggt/dependency/projection.py

@@ -0,0 +1,228 @@
+# 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 torch
+import numpy as np
+from .distortion import apply_distortion
+
+
+def img_from_cam_np(
+    intrinsics: np.ndarray, points_cam: np.ndarray, extra_params: np.ndarray | None = None, default: float = 0.0
+) -> np.ndarray:
+    """
+    Apply intrinsics (and optional radial distortion) to camera-space points.
+
+    Args
+    ----
+    intrinsics  : (B,3,3) camera matrix K.
+    points_cam  : (B,3,N) homogeneous camera coords  (x, y, z)ᵀ.
+    extra_params: (B, N) or (B, k) distortion params (k = 1,2,4) or None.
+    default     : value used for np.nan replacement.
+
+    Returns
+    -------
+    points2D : (B,N,2) pixel coordinates.
+    """
+    # 1. perspective divide  ───────────────────────────────────────
+    z = points_cam[:, 2:3, :]  # (B,1,N)
+    points_cam_norm = points_cam / z  # (B,3,N)
+    uv = points_cam_norm[:, :2, :]  # (B,2,N)
+
+    # 2. optional distortion ──────────────────────────────────────
+    if extra_params is not None:
+        uu, vv = apply_distortion(extra_params, uv[:, 0], uv[:, 1])
+        uv = np.stack([uu, vv], axis=1)  # (B,2,N)
+
+    # 3. homogeneous coords then K multiplication ─────────────────
+    ones = np.ones_like(uv[:, :1, :])  # (B,1,N)
+    points_cam_h = np.concatenate([uv, ones], axis=1)  # (B,3,N)
+
+    # batched mat-mul: K · [u v 1]ᵀ
+    points2D_h = np.einsum("bij,bjk->bik", intrinsics, points_cam_h)  # (B,3,N)
+    points2D = np.nan_to_num(points2D_h[:, :2, :], nan=default)  # (B,2,N)
+
+    return points2D.transpose(0, 2, 1)  # (B,N,2)
+
+
+def project_3D_points_np(
+    points3D: np.ndarray,
+    extrinsics: np.ndarray,
+    intrinsics: np.ndarray | None = None,
+    extra_params: np.ndarray | None = None,
+    *,
+    default: float = 0.0,
+    only_points_cam: bool = False,
+):
+    """
+    NumPy clone of ``project_3D_points``.
+
+    Parameters
+    ----------
+    points3D          : (N,3) world-space points.
+    extrinsics        : (B,3,4)  [R|t] matrix for each of B cameras.
+    intrinsics        : (B,3,3)  K matrix (optional if you only need cam-space).
+    extra_params      : (B,k) or (B,N) distortion parameters (k ∈ {1,2,4}) or None.
+    default           : value used to replace NaNs.
+    only_points_cam   : if True, skip the projection and return points_cam with points2D as None.
+
+    Returns
+    -------
+    (points2D, points_cam) : A tuple where points2D is (B,N,2) pixel coords or None if only_points_cam=True,
+                           and points_cam is (B,3,N) camera-space coordinates.
+    """
+    # ----- 0. prep sizes -----------------------------------------------------
+    N = points3D.shape[0]  # #points
+    B = extrinsics.shape[0]  # #cameras
+
+    # ----- 1. world → homogeneous -------------------------------------------
+    w_h = np.ones((N, 1), dtype=points3D.dtype)
+    points3D_h = np.concatenate([points3D, w_h], axis=1)  # (N,4)
+
+    # broadcast to every camera (no actual copying with np.broadcast_to) ------
+    points3D_h_B = np.broadcast_to(points3D_h, (B, N, 4))  # (B,N,4)
+
+    # ----- 2. apply extrinsics  (camera frame) ------------------------------
+    # X_cam = E · X_hom
+    # einsum:  E_(b i j)  ·  X_(b n j)  →  (b n i)
+    points_cam = np.einsum("bij,bnj->bni", extrinsics, points3D_h_B)  # (B,N,3)
+    points_cam = points_cam.transpose(0, 2, 1)  # (B,3,N)
+
+    if only_points_cam:
+        return None, points_cam
+
+    # ----- 3. intrinsics + distortion ---------------------------------------
+    if intrinsics is None:
+        raise ValueError("`intrinsics` must be provided unless only_points_cam=True")
+
+    points2D = img_from_cam_np(intrinsics, points_cam, extra_params=extra_params, default=default)
+
+    return points2D, points_cam
+
+
+def project_3D_points(points3D, extrinsics, intrinsics=None, extra_params=None, default=0, only_points_cam=False):
+    """
+    Transforms 3D points to 2D using extrinsic and intrinsic parameters.
+    Args:
+        points3D (torch.Tensor): 3D points of shape Px3.
+        extrinsics (torch.Tensor): Extrinsic parameters of shape Bx3x4.
+        intrinsics (torch.Tensor): Intrinsic parameters of shape Bx3x3.
+        extra_params (torch.Tensor): Extra parameters of shape BxN, used for radial distortion.
+        default (float): Default value to replace NaNs.
+        only_points_cam (bool): If True, skip the projection and return points2D as None.
+
+    Returns:
+        tuple: (points2D, points_cam) where points2D is of shape BxNx2 or None if only_points_cam=True,
+               and points_cam is of shape Bx3xN.
+    """
+    with torch.cuda.amp.autocast(dtype=torch.double):
+        N = points3D.shape[0]  # Number of points
+        B = extrinsics.shape[0]  # Batch size, i.e., number of cameras
+        points3D_homogeneous = torch.cat([points3D, torch.ones_like(points3D[..., 0:1])], dim=1)  # Nx4
+        # Reshape for batch processing
+        points3D_homogeneous = points3D_homogeneous.unsqueeze(0).expand(B, -1, -1)  # BxNx4
+
+        # Step 1: Apply extrinsic parameters
+        # Transform 3D points to camera coordinate system for all cameras
+        points_cam = torch.bmm(extrinsics, points3D_homogeneous.transpose(-1, -2))
+
+        if only_points_cam:
+            return None, points_cam
+
+        # Step 2: Apply intrinsic parameters and (optional) distortion
+        points2D = img_from_cam(intrinsics, points_cam, extra_params, default)
+
+        return points2D, points_cam
+
+
+def img_from_cam(intrinsics, points_cam, extra_params=None, default=0.0):
+    """
+    Applies intrinsic parameters and optional distortion to the given 3D points.
+
+    Args:
+        intrinsics (torch.Tensor): Intrinsic camera parameters of shape Bx3x3.
+        points_cam (torch.Tensor): 3D points in camera coordinates of shape Bx3xN.
+        extra_params (torch.Tensor, optional): Distortion parameters of shape BxN, where N can be 1, 2, or 4.
+        default (float, optional): Default value to replace NaNs in the output.
+
+    Returns:
+        points2D (torch.Tensor): 2D points in pixel coordinates of shape BxNx2.
+    """
+
+    # Normalize by the third coordinate (homogeneous division)
+    points_cam = points_cam / points_cam[:, 2:3, :]
+    # Extract uv
+    uv = points_cam[:, :2, :]
+
+    # Apply distortion if extra_params are provided
+    if extra_params is not None:
+        uu, vv = apply_distortion(extra_params, uv[:, 0], uv[:, 1])
+        uv = torch.stack([uu, vv], dim=1)
+
+    # Prepare points_cam for batch matrix multiplication
+    points_cam_homo = torch.cat((uv, torch.ones_like(uv[:, :1, :])), dim=1)  # Bx3xN
+    # Apply intrinsic parameters using batch matrix multiplication
+    points2D_homo = torch.bmm(intrinsics, points_cam_homo)  # Bx3xN
+
+    # Extract x and y coordinates
+    points2D = points2D_homo[:, :2, :]  # Bx2xN
+
+    # Replace NaNs with default value
+    points2D = torch.nan_to_num(points2D, nan=default)
+
+    return points2D.transpose(1, 2)  # BxNx2
+
+
+if __name__ == "__main__":
+    # Set up example input
+    B, N = 24, 10240
+
+    for _ in range(100):
+        points3D = np.random.rand(N, 3).astype(np.float64)
+        extrinsics = np.random.rand(B, 3, 4).astype(np.float64)
+        intrinsics = np.random.rand(B, 3, 3).astype(np.float64)
+
+        # Convert to torch tensors
+        points3D_torch = torch.tensor(points3D)
+        extrinsics_torch = torch.tensor(extrinsics)
+        intrinsics_torch = torch.tensor(intrinsics)
+
+        # Run NumPy implementation
+        points2D_np, points_cam_np = project_3D_points_np(points3D, extrinsics, intrinsics)
+
+        # Run torch implementation
+        points2D_torch, points_cam_torch = project_3D_points(points3D_torch, extrinsics_torch, intrinsics_torch)
+
+        # Convert torch output to numpy
+        points2D_torch_np = points2D_torch.detach().numpy()
+        points_cam_torch_np = points_cam_torch.detach().numpy()
+
+        # Compute difference
+        diff = np.abs(points2D_np - points2D_torch_np)
+        print("Difference between NumPy and PyTorch implementations:")
+        print(diff)
+
+        # Check max error
+        max_diff = np.max(diff)
+        print(f"Maximum difference: {max_diff}")
+
+        if np.allclose(points2D_np, points2D_torch_np, atol=1e-6):
+            print("Implementations match closely.")
+        else:
+            print("Significant differences detected.")
+
+        if points_cam_np is not None:
+            points_cam_diff = np.abs(points_cam_np - points_cam_torch_np)
+            print("Difference between NumPy and PyTorch camera-space coordinates:")
+            print(points_cam_diff)
+
+            # Check max error
+            max_cam_diff = np.max(points_cam_diff)
+            print(f"Maximum camera-space coordinate difference: {max_cam_diff}")
+
+            if np.allclose(points_cam_np, points_cam_torch_np, atol=1e-6):
+                print("Camera-space coordinates match closely.")
+            else:
+                print("Significant differences detected in camera-space coordinates.")

vggt/dependency/track_modules/base_track_predictor.py

@@ -0,0 +1,190 @@
+# 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 torch
+import torch.nn as nn
+from einops import rearrange, repeat
+
+from .blocks import EfficientUpdateFormer, CorrBlock
+from .utils import sample_features4d, get_2d_embedding, get_2d_sincos_pos_embed
+
+
+class BaseTrackerPredictor(nn.Module):
+    def __init__(
+        self,
+        stride=4,
+        corr_levels=5,
+        corr_radius=4,
+        latent_dim=128,
+        hidden_size=384,
+        use_spaceatt=True,
+        depth=6,
+        fine=False,
+    ):
+        super(BaseTrackerPredictor, self).__init__()
+        """
+        The base template to create a track predictor
+        
+        Modified from https://github.com/facebookresearch/co-tracker/
+        """
+
+        self.stride = stride
+        self.latent_dim = latent_dim
+        self.corr_levels = corr_levels
+        self.corr_radius = corr_radius
+        self.hidden_size = hidden_size
+        self.fine = fine
+
+        self.flows_emb_dim = latent_dim // 2
+        self.transformer_dim = self.corr_levels * (self.corr_radius * 2 + 1) ** 2 + self.latent_dim * 2
+
+        if self.fine:
+            # TODO this is the old dummy code, will remove this when we train next model
+            self.transformer_dim += 4 if self.transformer_dim % 2 == 0 else 5
+        else:
+            self.transformer_dim += (4 - self.transformer_dim % 4) % 4
+
+        space_depth = depth if use_spaceatt else 0
+        time_depth = depth
+
+        self.updateformer = EfficientUpdateFormer(
+            space_depth=space_depth,
+            time_depth=time_depth,
+            input_dim=self.transformer_dim,
+            hidden_size=self.hidden_size,
+            output_dim=self.latent_dim + 2,
+            mlp_ratio=4.0,
+            add_space_attn=use_spaceatt,
+        )
+
+        self.norm = nn.GroupNorm(1, self.latent_dim)
+
+        # A linear layer to update track feats at each iteration
+        self.ffeat_updater = nn.Sequential(nn.Linear(self.latent_dim, self.latent_dim), nn.GELU())
+
+        if not self.fine:
+            self.vis_predictor = nn.Sequential(nn.Linear(self.latent_dim, 1))
+
+    def forward(self, query_points, fmaps=None, iters=4, return_feat=False, down_ratio=1):
+        """
+        query_points: B x N x 2, the number of batches, tracks, and xy
+        fmaps: B x S x C x HH x WW, the number of batches, frames, and feature dimension.
+                note HH and WW is the size of feature maps instead of original images
+        """
+        B, N, D = query_points.shape
+        B, S, C, HH, WW = fmaps.shape
+
+        assert D == 2
+
+        # Scale the input query_points because we may downsample the images
+        # by down_ratio or self.stride
+        # e.g., if a 3x1024x1024 image is processed to a 128x256x256 feature map
+        # its query_points should be query_points/4
+        if down_ratio > 1:
+            query_points = query_points / float(down_ratio)
+        query_points = query_points / float(self.stride)
+
+        # Init with coords as the query points
+        # It means the search will start from the position of query points at the reference frames
+        coords = query_points.clone().reshape(B, 1, N, 2).repeat(1, S, 1, 1)
+
+        # Sample/extract the features of the query points in the query frame
+        query_track_feat = sample_features4d(fmaps[:, 0], coords[:, 0])
+
+        # init track feats by query feats
+        track_feats = query_track_feat.unsqueeze(1).repeat(1, S, 1, 1)  # B, S, N, C
+        # back up the init coords
+        coords_backup = coords.clone()
+
+        # Construct the correlation block
+
+        fcorr_fn = CorrBlock(fmaps, num_levels=self.corr_levels, radius=self.corr_radius)
+
+        coord_preds = []
+
+        # Iterative Refinement
+        for itr in range(iters):
+            # Detach the gradients from the last iteration
+            # (in my experience, not very important for performance)
+            coords = coords.detach()
+
+            # Compute the correlation (check the implementation of CorrBlock)
+
+            fcorr_fn.corr(track_feats)
+            fcorrs = fcorr_fn.sample(coords)  # B, S, N, corrdim
+
+            corrdim = fcorrs.shape[3]
+
+            fcorrs_ = fcorrs.permute(0, 2, 1, 3).reshape(B * N, S, corrdim)
+
+            # Movement of current coords relative to query points
+            flows = (coords - coords[:, 0:1]).permute(0, 2, 1, 3).reshape(B * N, S, 2)
+
+            flows_emb = get_2d_embedding(flows, self.flows_emb_dim, cat_coords=False)
+
+            # (In my trials, it is also okay to just add the flows_emb instead of concat)
+            flows_emb = torch.cat([flows_emb, flows], dim=-1)
+
+            track_feats_ = track_feats.permute(0, 2, 1, 3).reshape(B * N, S, self.latent_dim)
+
+            # Concatenate them as the input for the transformers
+            transformer_input = torch.cat([flows_emb, fcorrs_, track_feats_], dim=2)
+
+            if transformer_input.shape[2] < self.transformer_dim:
+                # pad the features to match the dimension
+                pad_dim = self.transformer_dim - transformer_input.shape[2]
+                pad = torch.zeros_like(flows_emb[..., 0:pad_dim])
+                transformer_input = torch.cat([transformer_input, pad], dim=2)
+
+            # 2D positional embed
+            # TODO: this can be much simplified
+            pos_embed = get_2d_sincos_pos_embed(self.transformer_dim, grid_size=(HH, WW)).to(query_points.device)
+            sampled_pos_emb = sample_features4d(pos_embed.expand(B, -1, -1, -1), coords[:, 0])
+            sampled_pos_emb = rearrange(sampled_pos_emb, "b n c -> (b n) c").unsqueeze(1)
+
+            x = transformer_input + sampled_pos_emb
+
+            # B, N, S, C
+            x = rearrange(x, "(b n) s d -> b n s d", b=B)
+
+            # Compute the delta coordinates and delta track features
+            delta = self.updateformer(x)
+            # BN, S, C
+            delta = rearrange(delta, " b n s d -> (b n) s d", b=B)
+            delta_coords_ = delta[:, :, :2]
+            delta_feats_ = delta[:, :, 2:]
+
+            track_feats_ = track_feats_.reshape(B * N * S, self.latent_dim)
+            delta_feats_ = delta_feats_.reshape(B * N * S, self.latent_dim)
+
+            # Update the track features
+            track_feats_ = self.ffeat_updater(self.norm(delta_feats_)) + track_feats_
+            track_feats = track_feats_.reshape(B, N, S, self.latent_dim).permute(0, 2, 1, 3)  # BxSxNxC
+
+            # B x S x N x 2
+            coords = coords + delta_coords_.reshape(B, N, S, 2).permute(0, 2, 1, 3)
+
+            # Force coord0 as query
+            # because we assume the query points should not be changed
+            coords[:, 0] = coords_backup[:, 0]
+
+            # The predicted tracks are in the original image scale
+            if down_ratio > 1:
+                coord_preds.append(coords * self.stride * down_ratio)
+            else:
+                coord_preds.append(coords * self.stride)
+
+        # B, S, N
+        if not self.fine:
+            vis_e = self.vis_predictor(track_feats.reshape(B * S * N, self.latent_dim)).reshape(B, S, N)
+            vis_e = torch.sigmoid(vis_e)
+        else:
+            vis_e = None
+
+        if return_feat:
+            return coord_preds, vis_e, track_feats, query_track_feat
+        else:
+            return coord_preds, vis_e

vggt/dependency/track_modules/blocks.py

@@ -0,0 +1,329 @@
+# 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.
+
+
+# Modified from https://github.com/facebookresearch/co-tracker/
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+from typing import Callable
+import collections
+from torch import Tensor
+from itertools import repeat
+
+from .utils import bilinear_sampler
+
+from .modules import Mlp, AttnBlock, CrossAttnBlock, ResidualBlock
+
+
+class BasicEncoder(nn.Module):
+    def __init__(self, input_dim=3, output_dim=128, stride=4):
+        super(BasicEncoder, self).__init__()
+
+        self.stride = stride
+        self.norm_fn = "instance"
+        self.in_planes = output_dim // 2
+
+        self.norm1 = nn.InstanceNorm2d(self.in_planes)
+        self.norm2 = nn.InstanceNorm2d(output_dim * 2)
+
+        self.conv1 = nn.Conv2d(input_dim, self.in_planes, kernel_size=7, stride=2, padding=3, padding_mode="zeros")
+        self.relu1 = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(output_dim // 2, stride=1)
+        self.layer2 = self._make_layer(output_dim // 4 * 3, stride=2)
+        self.layer3 = self._make_layer(output_dim, stride=2)
+        self.layer4 = self._make_layer(output_dim, stride=2)
+
+        self.conv2 = nn.Conv2d(
+            output_dim * 3 + output_dim // 4, output_dim * 2, kernel_size=3, padding=1, padding_mode="zeros"
+        )
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(output_dim * 2, output_dim, kernel_size=1)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.InstanceNorm2d)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        _, _, H, W = x.shape
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        a = self.layer1(x)
+        b = self.layer2(a)
+        c = self.layer3(b)
+        d = self.layer4(c)
+
+        a = _bilinear_intepolate(a, self.stride, H, W)
+        b = _bilinear_intepolate(b, self.stride, H, W)
+        c = _bilinear_intepolate(c, self.stride, H, W)
+        d = _bilinear_intepolate(d, self.stride, H, W)
+
+        x = self.conv2(torch.cat([a, b, c, d], dim=1))
+        x = self.norm2(x)
+        x = self.relu2(x)
+        x = self.conv3(x)
+        return x
+
+
+class ShallowEncoder(nn.Module):
+    def __init__(self, input_dim=3, output_dim=32, stride=1, norm_fn="instance"):
+        super(ShallowEncoder, self).__init__()
+        self.stride = stride
+        self.norm_fn = norm_fn
+        self.in_planes = output_dim
+
+        if self.norm_fn == "group":
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=self.in_planes)
+            self.norm2 = nn.GroupNorm(num_groups=8, num_channels=output_dim * 2)
+        elif self.norm_fn == "batch":
+            self.norm1 = nn.BatchNorm2d(self.in_planes)
+            self.norm2 = nn.BatchNorm2d(output_dim * 2)
+        elif self.norm_fn == "instance":
+            self.norm1 = nn.InstanceNorm2d(self.in_planes)
+            self.norm2 = nn.InstanceNorm2d(output_dim * 2)
+        elif self.norm_fn == "none":
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(input_dim, self.in_planes, kernel_size=3, stride=2, padding=1, padding_mode="zeros")
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.layer1 = self._make_layer(output_dim, stride=2)
+
+        self.layer2 = self._make_layer(output_dim, stride=2)
+        self.conv2 = nn.Conv2d(output_dim, output_dim, kernel_size=1)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        self.in_planes = dim
+
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        return layer1
+
+    def forward(self, x):
+        _, _, H, W = x.shape
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        tmp = self.layer1(x)
+        x = x + F.interpolate(tmp, (x.shape[-2:]), mode="bilinear", align_corners=True)
+        tmp = self.layer2(tmp)
+        x = x + F.interpolate(tmp, (x.shape[-2:]), mode="bilinear", align_corners=True)
+        tmp = None
+        x = self.conv2(x) + x
+
+        x = F.interpolate(x, (H // self.stride, W // self.stride), mode="bilinear", align_corners=True)
+
+        return x
+
+
+def _bilinear_intepolate(x, stride, H, W):
+    return F.interpolate(x, (H // stride, W // stride), mode="bilinear", align_corners=True)
+
+
+class EfficientUpdateFormer(nn.Module):
+    """
+    Transformer model that updates track estimates.
+    """
+
+    def __init__(
+        self,
+        space_depth=6,
+        time_depth=6,
+        input_dim=320,
+        hidden_size=384,
+        num_heads=8,
+        output_dim=130,
+        mlp_ratio=4.0,
+        add_space_attn=True,
+        num_virtual_tracks=64,
+    ):
+        super().__init__()
+
+        self.out_channels = 2
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.add_space_attn = add_space_attn
+        self.input_transform = torch.nn.Linear(input_dim, hidden_size, bias=True)
+        self.flow_head = torch.nn.Linear(hidden_size, output_dim, bias=True)
+        self.num_virtual_tracks = num_virtual_tracks
+
+        if self.add_space_attn:
+            self.virual_tracks = nn.Parameter(torch.randn(1, num_virtual_tracks, 1, hidden_size))
+        else:
+            self.virual_tracks = None
+
+        self.time_blocks = nn.ModuleList(
+            [
+                AttnBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attn_class=nn.MultiheadAttention)
+                for _ in range(time_depth)
+            ]
+        )
+
+        if add_space_attn:
+            self.space_virtual_blocks = nn.ModuleList(
+                [
+                    AttnBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attn_class=nn.MultiheadAttention)
+                    for _ in range(space_depth)
+                ]
+            )
+            self.space_point2virtual_blocks = nn.ModuleList(
+                [CrossAttnBlock(hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(space_depth)]
+            )
+            self.space_virtual2point_blocks = nn.ModuleList(
+                [CrossAttnBlock(hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(space_depth)]
+            )
+            assert len(self.time_blocks) >= len(self.space_virtual2point_blocks)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        def init_weights_vit_timm(module: nn.Module, name: str = ""):
+            """ViT weight initialization, original timm impl (for reproducibility)"""
+            if isinstance(module, nn.Linear):
+                trunc_normal_(module.weight, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+    def forward(self, input_tensor, mask=None):
+        tokens = self.input_transform(input_tensor)
+
+        init_tokens = tokens
+
+        B, _, T, _ = tokens.shape
+
+        if self.add_space_attn:
+            virtual_tokens = self.virual_tracks.repeat(B, 1, T, 1)
+            tokens = torch.cat([tokens, virtual_tokens], dim=1)
+
+        _, N, _, _ = tokens.shape
+
+        j = 0
+        for i in range(len(self.time_blocks)):
+            time_tokens = tokens.contiguous().view(B * N, T, -1)  # B N T C -> (B N) T C
+            time_tokens = self.time_blocks[i](time_tokens)
+
+            tokens = time_tokens.view(B, N, T, -1)  # (B N) T C -> B N T C
+            if self.add_space_attn and (i % (len(self.time_blocks) // len(self.space_virtual_blocks)) == 0):
+                space_tokens = tokens.permute(0, 2, 1, 3).contiguous().view(B * T, N, -1)  # B N T C -> (B T) N C
+                point_tokens = space_tokens[:, : N - self.num_virtual_tracks]
+                virtual_tokens = space_tokens[:, N - self.num_virtual_tracks :]
+
+                virtual_tokens = self.space_virtual2point_blocks[j](virtual_tokens, point_tokens, mask=mask)
+                virtual_tokens = self.space_virtual_blocks[j](virtual_tokens)
+                point_tokens = self.space_point2virtual_blocks[j](point_tokens, virtual_tokens, mask=mask)
+                space_tokens = torch.cat([point_tokens, virtual_tokens], dim=1)
+                tokens = space_tokens.view(B, T, N, -1).permute(0, 2, 1, 3)  # (B T) N C -> B N T C
+                j += 1
+
+        if self.add_space_attn:
+            tokens = tokens[:, : N - self.num_virtual_tracks]
+
+        tokens = tokens + init_tokens
+
+        flow = self.flow_head(tokens)
+        return flow
+
+
+class CorrBlock:
+    def __init__(self, fmaps, num_levels=4, radius=4, multiple_track_feats=False, padding_mode="zeros"):
+        B, S, C, H, W = fmaps.shape
+        self.S, self.C, self.H, self.W = S, C, H, W
+        self.padding_mode = padding_mode
+        self.num_levels = num_levels
+        self.radius = radius
+        self.fmaps_pyramid = []
+        self.multiple_track_feats = multiple_track_feats
+
+        self.fmaps_pyramid.append(fmaps)
+        for i in range(self.num_levels - 1):
+            fmaps_ = fmaps.reshape(B * S, C, H, W)
+            fmaps_ = F.avg_pool2d(fmaps_, 2, stride=2)
+            _, _, H, W = fmaps_.shape
+            fmaps = fmaps_.reshape(B, S, C, H, W)
+            self.fmaps_pyramid.append(fmaps)
+
+    def sample(self, coords):
+        r = self.radius
+        B, S, N, D = coords.shape
+        assert D == 2
+
+        H, W = self.H, self.W
+        out_pyramid = []
+        for i in range(self.num_levels):
+            corrs = self.corrs_pyramid[i]  # B, S, N, H, W
+            *_, H, W = corrs.shape
+
+            dx = torch.linspace(-r, r, 2 * r + 1)
+            dy = torch.linspace(-r, r, 2 * r + 1)
+            delta = torch.stack(torch.meshgrid(dy, dx, indexing="ij"), axis=-1).to(coords.device)
+
+            centroid_lvl = coords.reshape(B * S * N, 1, 1, 2) / 2**i
+            delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)
+            coords_lvl = centroid_lvl + delta_lvl
+
+            corrs = bilinear_sampler(corrs.reshape(B * S * N, 1, H, W), coords_lvl, padding_mode=self.padding_mode)
+            corrs = corrs.view(B, S, N, -1)
+
+            out_pyramid.append(corrs)
+
+        out = torch.cat(out_pyramid, dim=-1).contiguous()  # B, S, N, LRR*2
+        return out
+
+    def corr(self, targets):
+        B, S, N, C = targets.shape
+        if self.multiple_track_feats:
+            targets_split = targets.split(C // self.num_levels, dim=-1)
+            B, S, N, C = targets_split[0].shape
+
+        assert C == self.C
+        assert S == self.S
+
+        fmap1 = targets
+
+        self.corrs_pyramid = []
+        for i, fmaps in enumerate(self.fmaps_pyramid):
+            *_, H, W = fmaps.shape
+            fmap2s = fmaps.view(B, S, C, H * W)  # B S C H W ->  B S C (H W)
+            if self.multiple_track_feats:
+                fmap1 = targets_split[i]
+            corrs = torch.matmul(fmap1, fmap2s)
+            corrs = corrs.view(B, S, N, H, W)  # B S N (H W) -> B S N H W
+            corrs = corrs / torch.sqrt(torch.tensor(C).float())
+            self.corrs_pyramid.append(corrs)

vggt/dependency/track_modules/modules.py

@@ -0,0 +1,202 @@
+# 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 torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+from typing import Callable
+import collections
+from torch import Tensor
+from itertools import repeat
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+
+    return parse
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+to_2tuple = _ntuple(2)
+
+
+class ResidualBlock(nn.Module):
+    """
+    ResidualBlock: construct a block of two conv layers with residual connections
+    """
+
+    def __init__(self, in_planes, planes, norm_fn="group", stride=1, kernel_size=3):
+        super(ResidualBlock, self).__init__()
+
+        self.conv1 = nn.Conv2d(
+            in_planes, planes, kernel_size=kernel_size, padding=1, stride=stride, padding_mode="zeros"
+        )
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, padding=1, padding_mode="zeros")
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == "group":
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+
+        elif norm_fn == "batch":
+            self.norm1 = nn.BatchNorm2d(planes)
+            self.norm2 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.BatchNorm2d(planes)
+
+        elif norm_fn == "instance":
+            self.norm1 = nn.InstanceNorm2d(planes)
+            self.norm2 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == "none":
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            if not stride == 1:
+                self.norm3 = nn.Sequential()
+        else:
+            raise NotImplementedError
+
+        if stride == 1:
+            self.downsample = None
+        else:
+            self.downsample = nn.Sequential(nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x + y)
+
+
+class Mlp(nn.Module):
+    """MLP as used in Vision Transformer, MLP-Mixer and related networks"""
+
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        norm_layer=None,
+        bias=True,
+        drop=0.0,
+        use_conv=False,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        bias = to_2tuple(bias)
+        drop_probs = to_2tuple(drop)
+        linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
+
+        self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0])
+        self.act = act_layer()
+        self.drop1 = nn.Dropout(drop_probs[0])
+        self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1])
+        self.drop2 = nn.Dropout(drop_probs[1])
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop1(x)
+        x = self.fc2(x)
+        x = self.drop2(x)
+        return x
+
+
+class AttnBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        attn_class: Callable[..., nn.Module] = nn.MultiheadAttention,
+        mlp_ratio=4.0,
+        **block_kwargs,
+    ):
+        """
+        Self attention block
+        """
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.attn = attn_class(embed_dim=hidden_size, num_heads=num_heads, batch_first=True, **block_kwargs)
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, drop=0)
+
+    def forward(self, x, mask=None):
+        # Prepare the mask for PyTorch's attention (it expects a different format)
+        # attn_mask = mask if mask is not None else None
+        # Normalize before attention
+        x = self.norm1(x)
+
+        # PyTorch's MultiheadAttention returns attn_output, attn_output_weights
+        # attn_output, _ = self.attn(x, x, x, attn_mask=attn_mask)
+
+        attn_output, _ = self.attn(x, x, x)
+
+        # Add & Norm
+        x = x + attn_output
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class CrossAttnBlock(nn.Module):
+    def __init__(self, hidden_size, context_dim, num_heads=1, mlp_ratio=4.0, **block_kwargs):
+        """
+        Cross attention block
+        """
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_context = nn.LayerNorm(hidden_size)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.cross_attn = nn.MultiheadAttention(
+            embed_dim=hidden_size, num_heads=num_heads, batch_first=True, **block_kwargs
+        )
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, drop=0)
+
+    def forward(self, x, context, mask=None):
+        # Normalize inputs
+        x = self.norm1(x)
+        context = self.norm_context(context)
+
+        # Apply cross attention
+        # Note: nn.MultiheadAttention returns attn_output, attn_output_weights
+        attn_output, _ = self.cross_attn(x, context, context, attn_mask=mask)
+
+        # Add & Norm
+        x = x + attn_output
+        x = x + self.mlp(self.norm2(x))
+        return x

vggt/dependency/track_modules/track_refine.py

@@ -0,0 +1,419 @@
+# 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 numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+from torch import nn, einsum
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange, Reduce
+
+from PIL import Image
+import os
+from typing import Union, Tuple
+
+
+def refine_track(
+    images, fine_fnet, fine_tracker, coarse_pred, compute_score=False, pradius=15, sradius=2, fine_iters=6, chunk=40960
+):
+    """
+    Refines the tracking of images using a fine track predictor and a fine feature network.
+    Check https://arxiv.org/abs/2312.04563 for more details.
+
+    Args:
+        images (torch.Tensor): The images to be tracked.
+        fine_fnet (nn.Module): The fine feature network.
+        fine_tracker (nn.Module): The fine track predictor.
+        coarse_pred (torch.Tensor): The coarse predictions of tracks.
+        compute_score (bool, optional): Whether to compute the score. Defaults to False.
+        pradius (int, optional): The radius of a patch. Defaults to 15.
+        sradius (int, optional): The search radius. Defaults to 2.
+
+    Returns:
+        torch.Tensor: The refined tracks.
+        torch.Tensor, optional: The score.
+    """
+
+    # coarse_pred shape: BxSxNx2,
+    # where B is the batch, S is the video/images length, and N is the number of tracks
+    # now we are going to extract patches with the center at coarse_pred
+    # Please note that the last dimension indicates x and y, and hence has a dim number of 2
+    B, S, N, _ = coarse_pred.shape
+    _, _, _, H, W = images.shape
+
+    # Given the raidus of a patch, compute the patch size
+    psize = pradius * 2 + 1
+
+    # Note that we assume the first frame is the query frame
+    # so the 2D locations of the first frame are the query points
+    query_points = coarse_pred[:, 0]
+
+    # Given 2D positions, we can use grid_sample to extract patches
+    # but it takes too much memory.
+    # Instead, we use the floored track xy to sample patches.
+
+    # For example, if the query point xy is (128.16, 252.78),
+    # and the patch size is (31, 31),
+    # our goal is to extract the content of a rectangle
+    # with left top: (113.16, 237.78)
+    # and right bottom: (143.16, 267.78).
+    # However, we record the floored left top: (113, 237)
+    # and the offset (0.16, 0.78)
+    # Then what we need is just unfolding the images like in CNN,
+    # picking the content at [(113, 237), (143, 267)].
+    # Such operations are highly optimized at pytorch
+    # (well if you really want to use interpolation, check the function extract_glimpse() below)
+
+    with torch.no_grad():
+        content_to_extract = images.reshape(B * S, 3, H, W)
+        C_in = content_to_extract.shape[1]
+
+        # Please refer to https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html
+        # for the detailed explanation of unfold()
+        # Here it runs sliding windows (psize x psize) to build patches
+        # The shape changes from
+        # (B*S)x C_in x H x W to (B*S)x C_in x H_new x W_new x Psize x Psize
+        # where Psize is the size of patch
+        content_to_extract = content_to_extract.unfold(2, psize, 1).unfold(3, psize, 1)
+
+    # Floor the coarse predictions to get integers and save the fractional/decimal
+    track_int = coarse_pred.floor().int()
+    track_frac = coarse_pred - track_int
+
+    # Note the points represent the center of patches
+    # now we get the location of the top left corner of patches
+    # because the ouput of pytorch unfold are indexed by top left corner
+    topleft = track_int - pradius
+    topleft_BSN = topleft.clone()
+
+    # clamp the values so that we will not go out of indexes
+    # NOTE: (VERY IMPORTANT: This operation ASSUMES H=W).
+    # You need to seperately clamp x and y if H!=W
+    topleft = topleft.clamp(0, H - psize)
+
+    # Reshape from BxSxNx2 -> (B*S)xNx2
+    topleft = topleft.reshape(B * S, N, 2)
+
+    # Prepare batches for indexing, shape: (B*S)xN
+    batch_indices = torch.arange(B * S)[:, None].expand(-1, N).to(content_to_extract.device)
+
+    # extracted_patches: (B*S) x N x C_in x Psize x Psize
+    extracted_patches = content_to_extract[batch_indices, :, topleft[..., 1], topleft[..., 0]]
+
+    if chunk < 0:
+        # Extract image patches based on top left corners
+        # Feed patches to fine fent for features
+        patch_feat = fine_fnet(extracted_patches.reshape(B * S * N, C_in, psize, psize))
+    else:
+        patches = extracted_patches.reshape(B * S * N, C_in, psize, psize)
+
+        patch_feat_list = []
+        for p in torch.split(patches, chunk):
+            patch_feat_list += [fine_fnet(p)]
+        patch_feat = torch.cat(patch_feat_list, 0)
+
+    C_out = patch_feat.shape[1]
+
+    # Refine the coarse tracks by fine_tracker
+    # reshape back to B x S x N x C_out x Psize x Psize
+    patch_feat = patch_feat.reshape(B, S, N, C_out, psize, psize)
+    patch_feat = rearrange(patch_feat, "b s n c p q -> (b n) s c p q")
+
+    # Prepare for the query points for fine tracker
+    # They are relative to the patch left top corner,
+    # instead of the image top left corner now
+    # patch_query_points: N x 1 x 2
+    # only 1 here because for each patch we only have 1 query point
+    patch_query_points = track_frac[:, 0] + pradius
+    patch_query_points = patch_query_points.reshape(B * N, 2).unsqueeze(1)
+
+    # Feed the PATCH query points and tracks into fine tracker
+    fine_pred_track_lists, _, _, query_point_feat = fine_tracker(
+        query_points=patch_query_points, fmaps=patch_feat, iters=fine_iters, return_feat=True
+    )
+
+    # relative the patch top left
+    fine_pred_track = fine_pred_track_lists[-1].clone()
+
+    # From (relative to the patch top left) to (relative to the image top left)
+    for idx in range(len(fine_pred_track_lists)):
+        fine_level = rearrange(fine_pred_track_lists[idx], "(b n) s u v -> b s n u v", b=B, n=N)
+        fine_level = fine_level.squeeze(-2)
+        fine_level = fine_level + topleft_BSN
+        fine_pred_track_lists[idx] = fine_level
+
+    # relative to the image top left
+    refined_tracks = fine_pred_track_lists[-1].clone()
+    refined_tracks[:, 0] = query_points
+
+    score = None
+
+    if compute_score:
+        score = compute_score_fn(query_point_feat, patch_feat, fine_pred_track, sradius, psize, B, N, S, C_out)
+
+    return refined_tracks, score
+
+
+def refine_track_v0(
+    images, fine_fnet, fine_tracker, coarse_pred, compute_score=False, pradius=15, sradius=2, fine_iters=6
+):
+    """
+    COPIED FROM VGGSfM
+
+    Refines the tracking of images using a fine track predictor and a fine feature network.
+    Check https://arxiv.org/abs/2312.04563 for more details.
+
+    Args:
+        images (torch.Tensor): The images to be tracked.
+        fine_fnet (nn.Module): The fine feature network.
+        fine_tracker (nn.Module): The fine track predictor.
+        coarse_pred (torch.Tensor): The coarse predictions of tracks.
+        compute_score (bool, optional): Whether to compute the score. Defaults to False.
+        pradius (int, optional): The radius of a patch. Defaults to 15.
+        sradius (int, optional): The search radius. Defaults to 2.
+
+    Returns:
+        torch.Tensor: The refined tracks.
+        torch.Tensor, optional: The score.
+    """
+
+    # coarse_pred shape: BxSxNx2,
+    # where B is the batch, S is the video/images length, and N is the number of tracks
+    # now we are going to extract patches with the center at coarse_pred
+    # Please note that the last dimension indicates x and y, and hence has a dim number of 2
+    B, S, N, _ = coarse_pred.shape
+    _, _, _, H, W = images.shape
+
+    # Given the raidus of a patch, compute the patch size
+    psize = pradius * 2 + 1
+
+    # Note that we assume the first frame is the query frame
+    # so the 2D locations of the first frame are the query points
+    query_points = coarse_pred[:, 0]
+
+    # Given 2D positions, we can use grid_sample to extract patches
+    # but it takes too much memory.
+    # Instead, we use the floored track xy to sample patches.
+
+    # For example, if the query point xy is (128.16, 252.78),
+    # and the patch size is (31, 31),
+    # our goal is to extract the content of a rectangle
+    # with left top: (113.16, 237.78)
+    # and right bottom: (143.16, 267.78).
+    # However, we record the floored left top: (113, 237)
+    # and the offset (0.16, 0.78)
+    # Then what we need is just unfolding the images like in CNN,
+    # picking the content at [(113, 237), (143, 267)].
+    # Such operations are highly optimized at pytorch
+    # (well if you really want to use interpolation, check the function extract_glimpse() below)
+
+    with torch.no_grad():
+        content_to_extract = images.reshape(B * S, 3, H, W)
+        C_in = content_to_extract.shape[1]
+
+        # Please refer to https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html
+        # for the detailed explanation of unfold()
+        # Here it runs sliding windows (psize x psize) to build patches
+        # The shape changes from
+        # (B*S)x C_in x H x W to (B*S)x C_in x H_new x W_new x Psize x Psize
+        # where Psize is the size of patch
+        content_to_extract = content_to_extract.unfold(2, psize, 1).unfold(3, psize, 1)
+
+    # Floor the coarse predictions to get integers and save the fractional/decimal
+    track_int = coarse_pred.floor().int()
+    track_frac = coarse_pred - track_int
+
+    # Note the points represent the center of patches
+    # now we get the location of the top left corner of patches
+    # because the ouput of pytorch unfold are indexed by top left corner
+    topleft = track_int - pradius
+    topleft_BSN = topleft.clone()
+
+    # clamp the values so that we will not go out of indexes
+    # NOTE: (VERY IMPORTANT: This operation ASSUMES H=W).
+    # You need to seperately clamp x and y if H!=W
+    topleft = topleft.clamp(0, H - psize)
+
+    # Reshape from BxSxNx2 -> (B*S)xNx2
+    topleft = topleft.reshape(B * S, N, 2)
+
+    # Prepare batches for indexing, shape: (B*S)xN
+    batch_indices = torch.arange(B * S)[:, None].expand(-1, N).to(content_to_extract.device)
+
+    # Extract image patches based on top left corners
+    # extracted_patches: (B*S) x N x C_in x Psize x Psize
+    extracted_patches = content_to_extract[batch_indices, :, topleft[..., 1], topleft[..., 0]]
+
+    # Feed patches to fine fent for features
+    patch_feat = fine_fnet(extracted_patches.reshape(B * S * N, C_in, psize, psize))
+
+    C_out = patch_feat.shape[1]
+
+    # Refine the coarse tracks by fine_tracker
+
+    # reshape back to B x S x N x C_out x Psize x Psize
+    patch_feat = patch_feat.reshape(B, S, N, C_out, psize, psize)
+    patch_feat = rearrange(patch_feat, "b s n c p q -> (b n) s c p q")
+
+    # Prepare for the query points for fine tracker
+    # They are relative to the patch left top corner,
+    # instead of the image top left corner now
+    # patch_query_points: N x 1 x 2
+    # only 1 here because for each patch we only have 1 query point
+    patch_query_points = track_frac[:, 0] + pradius
+    patch_query_points = patch_query_points.reshape(B * N, 2).unsqueeze(1)
+
+    # Feed the PATCH query points and tracks into fine tracker
+    fine_pred_track_lists, _, _, query_point_feat = fine_tracker(
+        query_points=patch_query_points, fmaps=patch_feat, iters=fine_iters, return_feat=True
+    )
+
+    # relative the patch top left
+    fine_pred_track = fine_pred_track_lists[-1].clone()
+
+    # From (relative to the patch top left) to (relative to the image top left)
+    for idx in range(len(fine_pred_track_lists)):
+        fine_level = rearrange(fine_pred_track_lists[idx], "(b n) s u v -> b s n u v", b=B, n=N)
+        fine_level = fine_level.squeeze(-2)
+        fine_level = fine_level + topleft_BSN
+        fine_pred_track_lists[idx] = fine_level
+
+    # relative to the image top left
+    refined_tracks = fine_pred_track_lists[-1].clone()
+    refined_tracks[:, 0] = query_points
+
+    score = None
+
+    if compute_score:
+        score = compute_score_fn(query_point_feat, patch_feat, fine_pred_track, sradius, psize, B, N, S, C_out)
+
+    return refined_tracks, score
+
+
+################################## NOTE: NOT USED ##################################
+
+
+def compute_score_fn(query_point_feat, patch_feat, fine_pred_track, sradius, psize, B, N, S, C_out):
+    """
+    Compute the scores, i.e., the standard deviation of the 2D similarity heatmaps,
+    given the query point features and reference frame feature maps
+    """
+
+    from kornia.utils.grid import create_meshgrid
+    from kornia.geometry.subpix import dsnt
+
+    # query_point_feat initial shape: B x N x C_out,
+    # query_point_feat indicates the feat at the coorponsing query points
+    # Therefore we don't have S dimension here
+    query_point_feat = query_point_feat.reshape(B, N, C_out)
+    # reshape and expand to B x (S-1) x N x C_out
+    query_point_feat = query_point_feat.unsqueeze(1).expand(-1, S - 1, -1, -1)
+    # and reshape to (B*(S-1)*N) x C_out
+    query_point_feat = query_point_feat.reshape(B * (S - 1) * N, C_out)
+
+    # Radius and size for computing the score
+    ssize = sradius * 2 + 1
+
+    # Reshape, you know it, so many reshaping operations
+    patch_feat = rearrange(patch_feat, "(b n) s c p q -> b s n c p q", b=B, n=N)
+
+    # Again, we unfold the patches to smaller patches
+    # so that we can then focus on smaller patches
+    # patch_feat_unfold shape:
+    # B x S x N x C_out x (psize - 2*sradius) x (psize - 2*sradius) x ssize x ssize
+    # well a bit scary, but actually not
+    patch_feat_unfold = patch_feat.unfold(4, ssize, 1).unfold(5, ssize, 1)
+
+    # Do the same stuffs above, i.e., the same as extracting patches
+    fine_prediction_floor = fine_pred_track.floor().int()
+    fine_level_floor_topleft = fine_prediction_floor - sradius
+
+    # Clamp to ensure the smaller patch is valid
+    fine_level_floor_topleft = fine_level_floor_topleft.clamp(0, psize - ssize)
+    fine_level_floor_topleft = fine_level_floor_topleft.squeeze(2)
+
+    # Prepare the batch indices and xy locations
+
+    batch_indices_score = torch.arange(B)[:, None, None].expand(-1, S, N)  # BxSxN
+    batch_indices_score = batch_indices_score.reshape(-1).to(patch_feat_unfold.device)  # B*S*N
+    y_indices = fine_level_floor_topleft[..., 0].flatten()  # Flatten H indices
+    x_indices = fine_level_floor_topleft[..., 1].flatten()  # Flatten W indices
+
+    reference_frame_feat = patch_feat_unfold.reshape(
+        B * S * N, C_out, psize - sradius * 2, psize - sradius * 2, ssize, ssize
+    )
+
+    # Note again, according to pytorch convention
+    # x_indices cooresponds to [..., 1] and y_indices cooresponds to [..., 0]
+    reference_frame_feat = reference_frame_feat[batch_indices_score, :, x_indices, y_indices]
+    reference_frame_feat = reference_frame_feat.reshape(B, S, N, C_out, ssize, ssize)
+    # pick the frames other than the first one, so we have S-1 frames here
+    reference_frame_feat = reference_frame_feat[:, 1:].reshape(B * (S - 1) * N, C_out, ssize * ssize)
+
+    # Compute similarity
+    sim_matrix = torch.einsum("mc,mcr->mr", query_point_feat, reference_frame_feat)
+    softmax_temp = 1.0 / C_out**0.5
+    heatmap = torch.softmax(softmax_temp * sim_matrix, dim=1)
+    # 2D heatmaps
+    heatmap = heatmap.reshape(B * (S - 1) * N, ssize, ssize)  # * x ssize x ssize
+
+    coords_normalized = dsnt.spatial_expectation2d(heatmap[None], True)[0]
+    grid_normalized = create_meshgrid(ssize, ssize, normalized_coordinates=True, device=heatmap.device).reshape(
+        1, -1, 2
+    )
+
+    var = torch.sum(grid_normalized**2 * heatmap.view(-1, ssize * ssize, 1), dim=1) - coords_normalized**2
+    std = torch.sum(torch.sqrt(torch.clamp(var, min=1e-10)), -1)  # clamp needed for numerical stability
+
+    score = std.reshape(B, S - 1, N)
+    # set score as 1 for the query frame
+    score = torch.cat([torch.ones_like(score[:, 0:1]), score], dim=1)
+
+    return score
+
+
+def extract_glimpse(
+    tensor: torch.Tensor, size: Tuple[int, int], offsets, mode="bilinear", padding_mode="zeros", debug=False, orib=None
+):
+    B, C, W, H = tensor.shape
+
+    h, w = size
+    xs = torch.arange(0, w, dtype=tensor.dtype, device=tensor.device) - (w - 1) / 2.0
+    ys = torch.arange(0, h, dtype=tensor.dtype, device=tensor.device) - (h - 1) / 2.0
+
+    vy, vx = torch.meshgrid(ys, xs)
+    grid = torch.stack([vx, vy], dim=-1)  # h, w, 2
+    grid = grid[None]
+
+    B, N, _ = offsets.shape
+
+    offsets = offsets.reshape((B * N), 1, 1, 2)
+    offsets_grid = offsets + grid
+
+    # normalised grid  to [-1, 1]
+    offsets_grid = (offsets_grid - offsets_grid.new_tensor([W / 2, H / 2])) / offsets_grid.new_tensor([W / 2, H / 2])
+
+    # BxCxHxW -> Bx1xCxHxW
+    tensor = tensor[:, None]
+
+    # Bx1xCxHxW -> BxNxCxHxW
+    tensor = tensor.expand(-1, N, -1, -1, -1)
+
+    # BxNxCxHxW -> (B*N)xCxHxW
+    tensor = tensor.reshape((B * N), C, W, H)
+
+    sampled = torch.nn.functional.grid_sample(
+        tensor, offsets_grid, mode=mode, align_corners=False, padding_mode=padding_mode
+    )
+
+    # NOTE: I am not sure it should be h, w or w, h here
+    # but okay for sqaures
+    sampled = sampled.reshape(B, N, C, h, w)
+
+    return sampled

vggt/dependency/track_modules/utils.py

@@ -0,0 +1,216 @@
+# 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.
+
+# Modified from https://github.com/facebookresearch/PoseDiffusion
+# and https://github.com/facebookresearch/co-tracker/tree/main
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Optional, Tuple, Union
+from einops import rearrange, repeat
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: Union[int, Tuple[int, int]], return_grid=False) -> torch.Tensor:
+    """
+    This function initializes a grid and generates a 2D positional embedding using sine and cosine functions.
+    It is a wrapper of get_2d_sincos_pos_embed_from_grid.
+    Args:
+    - embed_dim: The embedding dimension.
+    - grid_size: The grid size.
+    Returns:
+    - pos_embed: The generated 2D positional embedding.
+    """
+    if isinstance(grid_size, tuple):
+        grid_size_h, grid_size_w = grid_size
+    else:
+        grid_size_h = grid_size_w = grid_size
+    grid_h = torch.arange(grid_size_h, dtype=torch.float)
+    grid_w = torch.arange(grid_size_w, dtype=torch.float)
+    grid = torch.meshgrid(grid_w, grid_h, indexing="xy")
+    grid = torch.stack(grid, dim=0)
+    grid = grid.reshape([2, 1, grid_size_h, grid_size_w])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if return_grid:
+        return (pos_embed.reshape(1, grid_size_h, grid_size_w, -1).permute(0, 3, 1, 2), grid)
+    return pos_embed.reshape(1, grid_size_h, grid_size_w, -1).permute(0, 3, 1, 2)
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: torch.Tensor) -> torch.Tensor:
+    """
+    This function generates a 2D positional embedding from a given grid using sine and cosine functions.
+
+    Args:
+    - embed_dim: The embedding dimension.
+    - grid: The grid to generate the embedding from.
+
+    Returns:
+    - emb: The generated 2D positional embedding.
+    """
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = torch.cat([emb_h, emb_w], dim=2)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: torch.Tensor) -> torch.Tensor:
+    """
+    This function generates a 1D positional embedding from a given grid using sine and cosine functions.
+
+    Args:
+    - embed_dim: The embedding dimension.
+    - pos: The position to generate the embedding from.
+
+    Returns:
+    - emb: The generated 1D positional embedding.
+    """
+    assert embed_dim % 2 == 0
+    omega = torch.arange(embed_dim // 2, dtype=torch.double)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = torch.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = torch.sin(out)  # (M, D/2)
+    emb_cos = torch.cos(out)  # (M, D/2)
+
+    emb = torch.cat([emb_sin, emb_cos], dim=1)  # (M, D)
+    return emb[None].float()
+
+
+def get_2d_embedding(xy: torch.Tensor, C: int, cat_coords: bool = True) -> torch.Tensor:
+    """
+    This function generates a 2D positional embedding from given coordinates using sine and cosine functions.
+
+    Args:
+    - xy: The coordinates to generate the embedding from.
+    - C: The size of the embedding.
+    - cat_coords: A flag to indicate whether to concatenate the original coordinates to the embedding.
+
+    Returns:
+    - pe: The generated 2D positional embedding.
+    """
+    B, N, D = xy.shape
+    assert D == 2
+
+    x = xy[:, :, 0:1]
+    y = xy[:, :, 1:2]
+    div_term = (torch.arange(0, C, 2, device=xy.device, dtype=torch.float32) * (1000.0 / C)).reshape(1, 1, int(C / 2))
+
+    pe_x = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+    pe_y = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+
+    pe_x[:, :, 0::2] = torch.sin(x * div_term)
+    pe_x[:, :, 1::2] = torch.cos(x * div_term)
+
+    pe_y[:, :, 0::2] = torch.sin(y * div_term)
+    pe_y[:, :, 1::2] = torch.cos(y * div_term)
+
+    pe = torch.cat([pe_x, pe_y], dim=2)  # (B, N, C*3)
+    if cat_coords:
+        pe = torch.cat([xy, pe], dim=2)  # (B, N, C*3+3)
+    return pe
+
+
+def bilinear_sampler(input, coords, align_corners=True, padding_mode="border"):
+    r"""Sample a tensor using bilinear interpolation
+
+    `bilinear_sampler(input, coords)` samples a tensor :attr:`input` at
+    coordinates :attr:`coords` using bilinear interpolation. It is the same
+    as `torch.nn.functional.grid_sample()` but with a different coordinate
+    convention.
+
+    The input tensor is assumed to be of shape :math:`(B, C, H, W)`, where
+    :math:`B` is the batch size, :math:`C` is the number of channels,
+    :math:`H` is the height of the image, and :math:`W` is the width of the
+    image. The tensor :attr:`coords` of shape :math:`(B, H_o, W_o, 2)` is
+    interpreted as an array of 2D point coordinates :math:`(x_i,y_i)`.
+
+    Alternatively, the input tensor can be of size :math:`(B, C, T, H, W)`,
+    in which case sample points are triplets :math:`(t_i,x_i,y_i)`. Note
+    that in this case the order of the components is slightly different
+    from `grid_sample()`, which would expect :math:`(x_i,y_i,t_i)`.
+
+    If `align_corners` is `True`, the coordinate :math:`x` is assumed to be
+    in the range :math:`[0,W-1]`, with 0 corresponding to the center of the
+    left-most image pixel :math:`W-1` to the center of the right-most
+    pixel.
+
+    If `align_corners` is `False`, the coordinate :math:`x` is assumed to
+    be in the range :math:`[0,W]`, with 0 corresponding to the left edge of
+    the left-most pixel :math:`W` to the right edge of the right-most
+    pixel.
+
+    Similar conventions apply to the :math:`y` for the range
+    :math:`[0,H-1]` and :math:`[0,H]` and to :math:`t` for the range
+    :math:`[0,T-1]` and :math:`[0,T]`.
+
+    Args:
+        input (Tensor): batch of input images.
+        coords (Tensor): batch of coordinates.
+        align_corners (bool, optional): Coordinate convention. Defaults to `True`.
+        padding_mode (str, optional): Padding mode. Defaults to `"border"`.
+
+    Returns:
+        Tensor: sampled points.
+    """
+
+    sizes = input.shape[2:]
+
+    assert len(sizes) in [2, 3]
+
+    if len(sizes) == 3:
+        # t x y -> x y t to match dimensions T H W in grid_sample
+        coords = coords[..., [1, 2, 0]]
+
+    if align_corners:
+        coords = coords * torch.tensor([2 / max(size - 1, 1) for size in reversed(sizes)], device=coords.device)
+    else:
+        coords = coords * torch.tensor([2 / size for size in reversed(sizes)], device=coords.device)
+
+    coords -= 1
+
+    return F.grid_sample(input, coords, align_corners=align_corners, padding_mode=padding_mode)
+
+
+def sample_features4d(input, coords):
+    r"""Sample spatial features
+
+    `sample_features4d(input, coords)` samples the spatial features
+    :attr:`input` represented by a 4D tensor :math:`(B, C, H, W)`.
+
+    The field is sampled at coordinates :attr:`coords` using bilinear
+    interpolation. :attr:`coords` is assumed to be of shape :math:`(B, R,
+    2)`, where each sample has the format :math:`(x_i, y_i)`. This uses the
+    same convention as :func:`bilinear_sampler` with `align_corners=True`.
+
+    The output tensor has one feature per point, and has shape :math:`(B,
+    R, C)`.
+
+    Args:
+        input (Tensor): spatial features.
+        coords (Tensor): points.
+
+    Returns:
+        Tensor: sampled features.
+    """
+
+    B, _, _, _ = input.shape
+
+    # B R 2 -> B R 1 2
+    coords = coords.unsqueeze(2)
+
+    # B C R 1
+    feats = bilinear_sampler(input, coords)
+
+    return feats.permute(0, 2, 1, 3).view(B, -1, feats.shape[1] * feats.shape[3])  # B C R 1 -> B R C

vggt/dependency/track_predict.py

@@ -0,0 +1,326 @@
+# 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 torch
+import numpy as np
+from .vggsfm_utils import *
+
+
+def predict_tracks(
+    images,
+    conf=None,
+    points_3d=None,
+    masks=None,
+    max_query_pts=2048,
+    query_frame_num=5,
+    keypoint_extractor="aliked+sp",
+    max_points_num=163840,
+    fine_tracking=True,
+    complete_non_vis=True,
+):
+    """
+    Predict tracks for the given images and masks.
+
+    TODO: support non-square images
+    TODO: support masks
+
+
+    This function predicts the tracks for the given images and masks using the specified query method
+    and track predictor. It finds query points, and predicts the tracks, visibility, and scores for the query frames.
+
+    Args:
+        images: Tensor of shape [S, 3, H, W] containing the input images.
+        conf: Tensor of shape [S, 1, H, W] containing the confidence scores. Default is None.
+        points_3d: Tensor containing 3D points. Default is None.
+        masks: Optional tensor of shape [S, 1, H, W] containing masks. Default is None.
+        max_query_pts: Maximum number of query points. Default is 2048.
+        query_frame_num: Number of query frames to use. Default is 5.
+        keypoint_extractor: Method for keypoint extraction. Default is "aliked+sp".
+        max_points_num: Maximum number of points to process at once. Default is 163840.
+        fine_tracking: Whether to use fine tracking. Default is True.
+        complete_non_vis: Whether to augment non-visible frames. Default is True.
+
+    Returns:
+        pred_tracks: Numpy array containing the predicted tracks.
+        pred_vis_scores: Numpy array containing the visibility scores for the tracks.
+        pred_confs: Numpy array containing the confidence scores for the tracks.
+        pred_points_3d: Numpy array containing the 3D points for the tracks.
+        pred_colors: Numpy array containing the point colors for the tracks. (0, 255)
+    """
+
+    device = images.device
+    dtype = images.dtype
+    tracker = build_vggsfm_tracker().to(device, dtype)
+
+    # Find query frames
+    query_frame_indexes = generate_rank_by_dino(images, query_frame_num=query_frame_num, device=device)
+
+    # Add the first image to the front if not already present
+    if 0 in query_frame_indexes:
+        query_frame_indexes.remove(0)
+    query_frame_indexes = [0, *query_frame_indexes]
+
+    # TODO: add the functionality to handle the masks
+    keypoint_extractors = initialize_feature_extractors(
+        max_query_pts, extractor_method=keypoint_extractor, device=device
+    )
+
+    pred_tracks = []
+    pred_vis_scores = []
+    pred_confs = []
+    pred_points_3d = []
+    pred_colors = []
+
+    fmaps_for_tracker = tracker.process_images_to_fmaps(images)
+
+    if fine_tracking:
+        print("For faster inference, consider disabling fine_tracking")
+
+    for query_index in query_frame_indexes:
+        print(f"Predicting tracks for query frame {query_index}")
+        pred_track, pred_vis, pred_conf, pred_point_3d, pred_color = _forward_on_query(
+            query_index,
+            images,
+            conf,
+            points_3d,
+            fmaps_for_tracker,
+            keypoint_extractors,
+            tracker,
+            max_points_num,
+            fine_tracking,
+            device,
+        )
+
+        pred_tracks.append(pred_track)
+        pred_vis_scores.append(pred_vis)
+        pred_confs.append(pred_conf)
+        pred_points_3d.append(pred_point_3d)
+        pred_colors.append(pred_color)
+
+    if complete_non_vis:
+        pred_tracks, pred_vis_scores, pred_confs, pred_points_3d, pred_colors = _augment_non_visible_frames(
+            pred_tracks,
+            pred_vis_scores,
+            pred_confs,
+            pred_points_3d,
+            pred_colors,
+            images,
+            conf,
+            points_3d,
+            fmaps_for_tracker,
+            keypoint_extractors,
+            tracker,
+            max_points_num,
+            fine_tracking,
+            min_vis=500,
+            non_vis_thresh=0.1,
+            device=device,
+        )
+
+    pred_tracks = np.concatenate(pred_tracks, axis=1)
+    pred_vis_scores = np.concatenate(pred_vis_scores, axis=1)
+    pred_confs = np.concatenate(pred_confs, axis=0) if pred_confs else None
+    pred_points_3d = np.concatenate(pred_points_3d, axis=0) if pred_points_3d else None
+    pred_colors = np.concatenate(pred_colors, axis=0) if pred_colors else None
+
+    # from vggt.utils.visual_track import visualize_tracks_on_images
+    # visualize_tracks_on_images(images[None], torch.from_numpy(pred_tracks[None]), torch.from_numpy(pred_vis_scores[None])>0.2, out_dir="track_visuals")
+
+    return pred_tracks, pred_vis_scores, pred_confs, pred_points_3d, pred_colors
+
+
+def _forward_on_query(
+    query_index,
+    images,
+    conf,
+    points_3d,
+    fmaps_for_tracker,
+    keypoint_extractors,
+    tracker,
+    max_points_num,
+    fine_tracking,
+    device,
+):
+    """
+    Process a single query frame for track prediction.
+
+    Args:
+        query_index: Index of the query frame
+        images: Tensor of shape [S, 3, H, W] containing the input images
+        conf: Confidence tensor
+        points_3d: 3D points tensor
+        fmaps_for_tracker: Feature maps for the tracker
+        keypoint_extractors: Initialized feature extractors
+        tracker: VGG-SFM tracker
+        max_points_num: Maximum number of points to process at once
+        fine_tracking: Whether to use fine tracking
+        device: Device to use for computation
+
+    Returns:
+        pred_track: Predicted tracks
+        pred_vis: Visibility scores for the tracks
+        pred_conf: Confidence scores for the tracks
+        pred_point_3d: 3D points for the tracks
+        pred_color: Point colors for the tracks (0, 255)
+    """
+    frame_num, _, height, width = images.shape
+
+    query_image = images[query_index]
+    query_points = extract_keypoints(query_image, keypoint_extractors, round_keypoints=False)
+    query_points = query_points[:, torch.randperm(query_points.shape[1], device=device)]
+
+    # Extract the color at the keypoint locations
+    query_points_long = query_points.squeeze(0).round().long()
+    pred_color = images[query_index][:, query_points_long[:, 1], query_points_long[:, 0]]
+    pred_color = (pred_color.permute(1, 0).cpu().numpy() * 255).astype(np.uint8)
+
+    # Query the confidence and points_3d at the keypoint locations
+    if (conf is not None) and (points_3d is not None):
+        assert height == width
+        assert conf.shape[-2] == conf.shape[-1]
+        assert conf.shape[:3] == points_3d.shape[:3]
+        scale = conf.shape[-1] / width
+
+        query_points_scaled = (query_points.squeeze(0) * scale).round().long()
+        query_points_scaled = query_points_scaled.cpu().numpy()
+
+        pred_conf = conf[query_index][query_points_scaled[:, 1], query_points_scaled[:, 0]]
+        pred_point_3d = points_3d[query_index][query_points_scaled[:, 1], query_points_scaled[:, 0]]
+
+        # heuristic to remove low confidence points
+        # should I export this as an input parameter?
+        valid_mask = pred_conf > 1.2
+        if valid_mask.sum() > 512:
+            query_points = query_points[:, valid_mask]  # Make sure shape is compatible
+            pred_conf = pred_conf[valid_mask]
+            pred_point_3d = pred_point_3d[valid_mask]
+            pred_color = pred_color[valid_mask]
+    else:
+        pred_conf = None
+        pred_point_3d = None
+
+    reorder_index = calculate_index_mappings(query_index, frame_num, device=device)
+
+    images_feed, fmaps_feed = switch_tensor_order([images, fmaps_for_tracker], reorder_index, dim=0)
+    images_feed = images_feed[None]  # add batch dimension
+    fmaps_feed = fmaps_feed[None]  # add batch dimension
+
+    all_points_num = images_feed.shape[1] * query_points.shape[1]
+
+    # Don't need to be scared, this is just chunking to make GPU happy
+    if all_points_num > max_points_num:
+        num_splits = (all_points_num + max_points_num - 1) // max_points_num
+        query_points = torch.chunk(query_points, num_splits, dim=1)
+    else:
+        query_points = [query_points]
+
+    pred_track, pred_vis, _ = predict_tracks_in_chunks(
+        tracker, images_feed, query_points, fmaps_feed, fine_tracking=fine_tracking
+    )
+
+    pred_track, pred_vis = switch_tensor_order([pred_track, pred_vis], reorder_index, dim=1)
+
+    pred_track = pred_track.squeeze(0).float().cpu().numpy()
+    pred_vis = pred_vis.squeeze(0).float().cpu().numpy()
+
+    return pred_track, pred_vis, pred_conf, pred_point_3d, pred_color
+
+
+def _augment_non_visible_frames(
+    pred_tracks: list,  # ← running list of np.ndarrays
+    pred_vis_scores: list,  # ← running list of np.ndarrays
+    pred_confs: list,  # ← running list of np.ndarrays for confidence scores
+    pred_points_3d: list,  # ← running list of np.ndarrays for 3D points
+    pred_colors: list,  # ← running list of np.ndarrays for colors
+    images: torch.Tensor,
+    conf,
+    points_3d,
+    fmaps_for_tracker,
+    keypoint_extractors,
+    tracker,
+    max_points_num: int,
+    fine_tracking: bool,
+    *,
+    min_vis: int = 500,
+    non_vis_thresh: float = 0.1,
+    device: torch.device = None,
+):
+    """
+    Augment tracking for frames with insufficient visibility.
+
+    Args:
+        pred_tracks: List of numpy arrays containing predicted tracks.
+        pred_vis_scores: List of numpy arrays containing visibility scores.
+        pred_confs: List of numpy arrays containing confidence scores.
+        pred_points_3d: List of numpy arrays containing 3D points.
+        pred_colors: List of numpy arrays containing point colors.
+        images: Tensor of shape [S, 3, H, W] containing the input images.
+        conf: Tensor of shape [S, 1, H, W] containing confidence scores
+        points_3d: Tensor containing 3D points
+        fmaps_for_tracker: Feature maps for the tracker
+        keypoint_extractors: Initialized feature extractors
+        tracker: VGG-SFM tracker
+        max_points_num: Maximum number of points to process at once
+        fine_tracking: Whether to use fine tracking
+        min_vis: Minimum visibility threshold
+        non_vis_thresh: Non-visibility threshold
+        device: Device to use for computation
+
+    Returns:
+        Updated pred_tracks, pred_vis_scores, pred_confs, pred_points_3d, and pred_colors lists.
+    """
+    last_query = -1
+    final_trial = False
+    cur_extractors = keypoint_extractors  # may be replaced on the final trial
+
+    while True:
+        # Visibility per frame
+        vis_array = np.concatenate(pred_vis_scores, axis=1)
+
+        # Count frames with sufficient visibility using numpy
+        sufficient_vis_count = (vis_array > non_vis_thresh).sum(axis=-1)
+        non_vis_frames = np.where(sufficient_vis_count < min_vis)[0].tolist()
+
+        if len(non_vis_frames) == 0:
+            break
+
+        print("Processing non visible frames:", non_vis_frames)
+
+        # Decide the frames & extractor for this round
+        if non_vis_frames[0] == last_query:
+            # Same frame failed twice - final "all-in" attempt
+            final_trial = True
+            cur_extractors = initialize_feature_extractors(2048, extractor_method="sp+sift+aliked", device=device)
+            query_frame_list = non_vis_frames  # blast them all at once
+        else:
+            query_frame_list = [non_vis_frames[0]]  # Process one at a time
+
+        last_query = non_vis_frames[0]
+
+        # Run the tracker for every selected frame
+        for query_index in query_frame_list:
+            new_track, new_vis, new_conf, new_point_3d, new_color = _forward_on_query(
+                query_index,
+                images,
+                conf,
+                points_3d,
+                fmaps_for_tracker,
+                cur_extractors,
+                tracker,
+                max_points_num,
+                fine_tracking,
+                device,
+            )
+            pred_tracks.append(new_track)
+            pred_vis_scores.append(new_vis)
+            pred_confs.append(new_conf)
+            pred_points_3d.append(new_point_3d)
+            pred_colors.append(new_color)
+
+        if final_trial:
+            break  # Stop after final attempt
+
+    return pred_tracks, pred_vis_scores, pred_confs, pred_points_3d, pred_colors

vggt/dependency/vggsfm_tracker.py

@@ -0,0 +1,124 @@
+# 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 numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+from torch import nn, einsum
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange, Reduce
+
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+
+from .track_modules.track_refine import refine_track
+from .track_modules.blocks import BasicEncoder, ShallowEncoder
+from .track_modules.base_track_predictor import BaseTrackerPredictor
+
+
+class TrackerPredictor(nn.Module):
+    def __init__(self, **extra_args):
+        super(TrackerPredictor, self).__init__()
+        """
+        Initializes the tracker predictor.
+
+        Both coarse_predictor and fine_predictor are constructed as a BaseTrackerPredictor,
+        check track_modules/base_track_predictor.py
+
+        Both coarse_fnet and fine_fnet are constructed as a 2D CNN network
+        check track_modules/blocks.py for BasicEncoder and ShallowEncoder
+        """
+        # Define coarse predictor configuration
+        coarse_stride = 4
+        self.coarse_down_ratio = 2
+
+        # Create networks directly instead of using instantiate
+        self.coarse_fnet = BasicEncoder(stride=coarse_stride)
+        self.coarse_predictor = BaseTrackerPredictor(stride=coarse_stride)
+
+        # Create fine predictor with stride = 1
+        self.fine_fnet = ShallowEncoder(stride=1)
+        self.fine_predictor = BaseTrackerPredictor(
+            stride=1,
+            depth=4,
+            corr_levels=3,
+            corr_radius=3,
+            latent_dim=32,
+            hidden_size=256,
+            fine=True,
+            use_spaceatt=False,
+        )
+
+    def forward(
+        self, images, query_points, fmaps=None, coarse_iters=6, inference=True, fine_tracking=True, fine_chunk=40960
+    ):
+        """
+        Args:
+            images (torch.Tensor): Images as RGB, in the range of [0, 1], with a shape of B x S x 3 x H x W.
+            query_points (torch.Tensor): 2D xy of query points, relative to top left, with a shape of B x N x 2.
+            fmaps (torch.Tensor, optional): Precomputed feature maps. Defaults to None.
+            coarse_iters (int, optional): Number of iterations for coarse prediction. Defaults to 6.
+            inference (bool, optional): Whether to perform inference. Defaults to True.
+            fine_tracking (bool, optional): Whether to perform fine tracking. Defaults to True.
+
+        Returns:
+            tuple: A tuple containing fine_pred_track, coarse_pred_track, pred_vis, and pred_score.
+        """
+
+        if fmaps is None:
+            batch_num, frame_num, image_dim, height, width = images.shape
+            reshaped_image = images.reshape(batch_num * frame_num, image_dim, height, width)
+            fmaps = self.process_images_to_fmaps(reshaped_image)
+            fmaps = fmaps.reshape(batch_num, frame_num, -1, fmaps.shape[-2], fmaps.shape[-1])
+
+            if inference:
+                torch.cuda.empty_cache()
+
+        # Coarse prediction
+        coarse_pred_track_lists, pred_vis = self.coarse_predictor(
+            query_points=query_points, fmaps=fmaps, iters=coarse_iters, down_ratio=self.coarse_down_ratio
+        )
+        coarse_pred_track = coarse_pred_track_lists[-1]
+
+        if inference:
+            torch.cuda.empty_cache()
+
+        if fine_tracking:
+            # Refine the coarse prediction
+            fine_pred_track, pred_score = refine_track(
+                images, self.fine_fnet, self.fine_predictor, coarse_pred_track, compute_score=False, chunk=fine_chunk
+            )
+
+            if inference:
+                torch.cuda.empty_cache()
+        else:
+            fine_pred_track = coarse_pred_track
+            pred_score = torch.ones_like(pred_vis)
+
+        return fine_pred_track, coarse_pred_track, pred_vis, pred_score
+
+    def process_images_to_fmaps(self, images):
+        """
+        This function processes images for inference.
+
+        Args:
+            images (torch.Tensor): The images to be processed with shape S x 3 x H x W.
+
+        Returns:
+            torch.Tensor: The processed feature maps.
+        """
+        if self.coarse_down_ratio > 1:
+            # whether or not scale down the input images to save memory
+            fmaps = self.coarse_fnet(
+                F.interpolate(images, scale_factor=1 / self.coarse_down_ratio, mode="bilinear", align_corners=True)
+            )
+        else:
+            fmaps = self.coarse_fnet(images)
+
+        return fmaps

vggt/dependency/vggsfm_utils.py

@@ -0,0 +1,305 @@
+# 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 logging
+import warnings
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import pycolmap
+import torch
+import torch.nn.functional as F
+from lightglue import ALIKED, SIFT, SuperPoint
+
+from .vggsfm_tracker import TrackerPredictor
+
+# Suppress verbose logging from dependencies
+logging.getLogger("dinov2").setLevel(logging.WARNING)
+warnings.filterwarnings("ignore", message="xFormers is available")
+warnings.filterwarnings("ignore", message="dinov2")
+
+# Constants
+_RESNET_MEAN = [0.485, 0.456, 0.406]
+_RESNET_STD = [0.229, 0.224, 0.225]
+
+
+def build_vggsfm_tracker(model_path=None):
+    """
+    Build and initialize the VGGSfM tracker.
+
+    Args:
+        model_path: Path to the model weights file. If None, weights are downloaded from HuggingFace.
+
+    Returns:
+        Initialized tracker model in eval mode.
+    """
+    tracker = TrackerPredictor()
+
+    if model_path is None:
+        default_url = "https://huggingface.co/facebook/VGGSfM/resolve/main/vggsfm_v2_tracker.pt"
+        tracker.load_state_dict(torch.hub.load_state_dict_from_url(default_url))
+    else:
+        tracker.load_state_dict(torch.load(model_path))
+
+    tracker.eval()
+    return tracker
+
+
+def generate_rank_by_dino(
+    images, query_frame_num, image_size=336, model_name="dinov2_vitb14_reg", device="cuda", spatial_similarity=False
+):
+    """
+    Generate a ranking of frames using DINO ViT features.
+
+    Args:
+        images: Tensor of shape (S, 3, H, W) with values in range [0, 1]
+        query_frame_num: Number of frames to select
+        image_size: Size to resize images to before processing
+        model_name: Name of the DINO model to use
+        device: Device to run the model on
+        spatial_similarity: Whether to use spatial token similarity or CLS token similarity
+
+    Returns:
+        List of frame indices ranked by their representativeness
+    """
+    # Resize images to the target size
+    images = F.interpolate(images, (image_size, image_size), mode="bilinear", align_corners=False)
+
+    # Load DINO model
+    dino_v2_model = torch.hub.load("facebookresearch/dinov2", model_name)
+    dino_v2_model.eval()
+    dino_v2_model = dino_v2_model.to(device)
+
+    # Normalize images using ResNet normalization
+    resnet_mean = torch.tensor(_RESNET_MEAN, device=device).view(1, 3, 1, 1)
+    resnet_std = torch.tensor(_RESNET_STD, device=device).view(1, 3, 1, 1)
+    images_resnet_norm = (images - resnet_mean) / resnet_std
+
+    with torch.no_grad():
+        frame_feat = dino_v2_model(images_resnet_norm, is_training=True)
+
+    # Process features based on similarity type
+    if spatial_similarity:
+        frame_feat = frame_feat["x_norm_patchtokens"]
+        frame_feat_norm = F.normalize(frame_feat, p=2, dim=1)
+
+        # Compute the similarity matrix
+        frame_feat_norm = frame_feat_norm.permute(1, 0, 2)
+        similarity_matrix = torch.bmm(frame_feat_norm, frame_feat_norm.transpose(-1, -2))
+        similarity_matrix = similarity_matrix.mean(dim=0)
+    else:
+        frame_feat = frame_feat["x_norm_clstoken"]
+        frame_feat_norm = F.normalize(frame_feat, p=2, dim=1)
+        similarity_matrix = torch.mm(frame_feat_norm, frame_feat_norm.transpose(-1, -2))
+
+    distance_matrix = 100 - similarity_matrix.clone()
+
+    # Ignore self-pairing
+    similarity_matrix.fill_diagonal_(-100)
+    similarity_sum = similarity_matrix.sum(dim=1)
+
+    # Find the most common frame
+    most_common_frame_index = torch.argmax(similarity_sum).item()
+
+    # Conduct FPS sampling starting from the most common frame
+    fps_idx = farthest_point_sampling(distance_matrix, query_frame_num, most_common_frame_index)
+
+    # Clean up all tensors and models to free memory
+    del frame_feat, frame_feat_norm, similarity_matrix, distance_matrix
+    del dino_v2_model
+    torch.cuda.empty_cache()
+
+    return fps_idx
+
+
+def farthest_point_sampling(distance_matrix, num_samples, most_common_frame_index=0):
+    """
+    Farthest point sampling algorithm to select diverse frames.
+
+    Args:
+        distance_matrix: Matrix of distances between frames
+        num_samples: Number of frames to select
+        most_common_frame_index: Index of the first frame to select
+
+    Returns:
+        List of selected frame indices
+    """
+    distance_matrix = distance_matrix.clamp(min=0)
+    N = distance_matrix.size(0)
+
+    # Initialize with the most common frame
+    selected_indices = [most_common_frame_index]
+    check_distances = distance_matrix[selected_indices]
+
+    while len(selected_indices) < num_samples:
+        # Find the farthest point from the current set of selected points
+        farthest_point = torch.argmax(check_distances)
+        selected_indices.append(farthest_point.item())
+
+        check_distances = distance_matrix[farthest_point]
+        # Mark already selected points to avoid selecting them again
+        check_distances[selected_indices] = 0
+
+        # Break if all points have been selected
+        if len(selected_indices) == N:
+            break
+
+    return selected_indices
+
+
+def calculate_index_mappings(query_index, S, device=None):
+    """
+    Construct an order that switches [query_index] and [0]
+    so that the content of query_index would be placed at [0].
+
+    Args:
+        query_index: Index to swap with 0
+        S: Total number of elements
+        device: Device to place the tensor on
+
+    Returns:
+        Tensor of indices with the swapped order
+    """
+    new_order = torch.arange(S)
+    new_order[0] = query_index
+    new_order[query_index] = 0
+    if device is not None:
+        new_order = new_order.to(device)
+    return new_order
+
+
+def switch_tensor_order(tensors, order, dim=1):
+    """
+    Reorder tensors along a specific dimension according to the given order.
+
+    Args:
+        tensors: List of tensors to reorder
+        order: Tensor of indices specifying the new order
+        dim: Dimension along which to reorder
+
+    Returns:
+        List of reordered tensors
+    """
+    return [torch.index_select(tensor, dim, order) if tensor is not None else None for tensor in tensors]
+
+
+def initialize_feature_extractors(max_query_num, det_thres=0.005, extractor_method="aliked", device="cuda"):
+    """
+    Initialize feature extractors that can be reused based on a method string.
+
+    Args:
+        max_query_num: Maximum number of keypoints to extract
+        det_thres: Detection threshold for keypoint extraction
+        extractor_method: String specifying which extractors to use (e.g., "aliked", "sp+sift", "aliked+sp+sift")
+        device: Device to run extraction on
+
+    Returns:
+        Dictionary of initialized extractors
+    """
+    extractors = {}
+    methods = extractor_method.lower().split("+")
+
+    for method in methods:
+        method = method.strip()
+        if method == "aliked":
+            aliked_extractor = ALIKED(max_num_keypoints=max_query_num, detection_threshold=det_thres)
+            extractors["aliked"] = aliked_extractor.to(device).eval()
+        elif method == "sp":
+            sp_extractor = SuperPoint(max_num_keypoints=max_query_num, detection_threshold=det_thres)
+            extractors["sp"] = sp_extractor.to(device).eval()
+        elif method == "sift":
+            sift_extractor = SIFT(max_num_keypoints=max_query_num)
+            extractors["sift"] = sift_extractor.to(device).eval()
+        else:
+            print(f"Warning: Unknown feature extractor '{method}', ignoring.")
+
+    if not extractors:
+        print(f"Warning: No valid extractors found in '{extractor_method}'. Using ALIKED by default.")
+        aliked_extractor = ALIKED(max_num_keypoints=max_query_num, detection_threshold=det_thres)
+        extractors["aliked"] = aliked_extractor.to(device).eval()
+
+    return extractors
+
+
+def extract_keypoints(query_image, extractors, round_keypoints=True):
+    """
+    Extract keypoints using pre-initialized feature extractors.
+
+    Args:
+        query_image: Input image tensor (3xHxW, range [0, 1])
+        extractors: Dictionary of initialized extractors
+
+    Returns:
+        Tensor of keypoint coordinates (1xNx2)
+    """
+    query_points = None
+
+    with torch.no_grad():
+        for extractor_name, extractor in extractors.items():
+            query_points_data = extractor.extract(query_image, invalid_mask=None)
+            extractor_points = query_points_data["keypoints"]
+            if round_keypoints:
+                extractor_points = extractor_points.round()
+
+            if query_points is not None:
+                query_points = torch.cat([query_points, extractor_points], dim=1)
+            else:
+                query_points = extractor_points
+
+    return query_points
+
+
+def predict_tracks_in_chunks(
+    track_predictor, images_feed, query_points_list, fmaps_feed, fine_tracking, num_splits=None, fine_chunk=40960
+):
+    """
+    Process a list of query points to avoid memory issues.
+
+    Args:
+        track_predictor (object): The track predictor object used for predicting tracks.
+        images_feed (torch.Tensor): A tensor of shape (B, T, C, H, W) representing a batch of images.
+        query_points_list (list or tuple): A list/tuple of tensors, each of shape (B, Ni, 2) representing chunks of query points.
+        fmaps_feed (torch.Tensor): A tensor of feature maps for the tracker.
+        fine_tracking (bool): Whether to perform fine tracking.
+        num_splits (int, optional): Ignored when query_points_list is provided. Kept for backward compatibility.
+
+    Returns:
+        tuple: A tuple containing the concatenated predicted tracks, visibility, and scores.
+    """
+    # If query_points_list is not a list or tuple but a single tensor, handle it like the old version for backward compatibility
+    if not isinstance(query_points_list, (list, tuple)):
+        query_points = query_points_list
+        if num_splits is None:
+            num_splits = 1
+        query_points_list = torch.chunk(query_points, num_splits, dim=1)
+
+    # Ensure query_points_list is a list for iteration (as torch.chunk returns a tuple)
+    if isinstance(query_points_list, tuple):
+        query_points_list = list(query_points_list)
+
+    fine_pred_track_list = []
+    pred_vis_list = []
+    pred_score_list = []
+
+    for split_points in query_points_list:
+        # Feed into track predictor for each split
+        fine_pred_track, _, pred_vis, pred_score = track_predictor(
+            images_feed, split_points, fmaps=fmaps_feed, fine_tracking=fine_tracking, fine_chunk=fine_chunk
+        )
+        fine_pred_track_list.append(fine_pred_track)
+        pred_vis_list.append(pred_vis)
+        pred_score_list.append(pred_score)
+
+    # Concatenate the results from all splits
+    fine_pred_track = torch.cat(fine_pred_track_list, dim=2)
+    pred_vis = torch.cat(pred_vis_list, dim=2)
+
+    if pred_score is not None:
+        pred_score = torch.cat(pred_score_list, dim=2)
+    else:
+        pred_score = None
+
+    return fine_pred_track, pred_vis, pred_score

vggt/heads/camera_head.py

@@ -50,12 +50,7 @@ class CameraHead(nn.Module):
         # Build the trunk using a sequence of transformer blocks.
         self.trunk = nn.Sequential(
             *[
-                Block(
-                    dim=dim_in,
-                    num_heads=num_heads,
-                    mlp_ratio=mlp_ratio,
-                    init_values=init_values,
-                )
+                Block(dim=dim_in, num_heads=num_heads, mlp_ratio=mlp_ratio, init_values=init_values)
                 for _ in range(trunk_depth)
             ]
         )
@@ -73,12 +68,7 @@ class CameraHead(nn.Module):
 
         # Adaptive layer normalization without affine parameters.
         self.adaln_norm = nn.LayerNorm(dim_in, elementwise_affine=False, eps=1e-6)
-        self.pose_branch = Mlp(
-            in_features=dim_in,
-            hidden_features=dim_in // 2,
-            out_features=self.target_dim,
-            drop=0,
-        )
+        self.pose_branch = Mlp(in_features=dim_in, hidden_features=dim_in // 2, out_features=self.target_dim, drop=0)
 
     def forward(self, aggregated_tokens_list: list, num_iterations: int = 4) -> list:
         """
@@ -144,10 +134,7 @@ class CameraHead(nn.Module):
 
             # Apply final activation functions for translation, quaternion, and field-of-view.
             activated_pose = activate_pose(
-                pred_pose_enc,
-                trans_act=self.trans_act,
-                quat_act=self.quat_act,
-                fl_act=self.fl_act,
+                pred_pose_enc, trans_act=self.trans_act, quat_act=self.quat_act, fl_act=self.fl_act
             )
             pred_pose_enc_list.append(activated_pose)
 

vggt/heads/dpt_head.py

@@ -67,16 +67,7 @@ class DPTHead(nn.Module):
 
         # Projection layers for each output channel from tokens.
         self.projects = nn.ModuleList(
-            [
-                nn.Conv2d(
-                    in_channels=dim_in,
-                    out_channels=oc,
-                    kernel_size=1,
-                    stride=1,
-                    padding=0,
-                )
-                for oc in out_channels
-            ]
+            [nn.Conv2d(in_channels=dim_in, out_channels=oc, kernel_size=1, stride=1, padding=0) for oc in out_channels]
         )
 
         # Resize layers for upsampling feature maps.
@@ -95,11 +86,7 @@ class DPTHead(nn.Module):
             ]
         )
 
-        self.scratch = _make_scratch(
-            out_channels,
-            features,
-            expand=False,
-        )
+        self.scratch = _make_scratch(out_channels, features, expand=False)
 
         # Attach additional modules to scratch.
         self.scratch.stem_transpose = None

vggt/heads/track_head.py

@@ -99,10 +99,6 @@ class TrackHead(nn.Module):
             iters = self.iters
 
         # Perform tracking using the extracted features
-        coord_preds, vis_scores, conf_scores = self.tracker(
-            query_points=query_points,
-            fmaps=feature_maps,
-            iters=iters,
-        )
+        coord_preds, vis_scores, conf_scores = self.tracker(query_points=query_points, fmaps=feature_maps, iters=iters)
 
         return coord_preds, vis_scores, conf_scores

vggt/heads/track_modules/blocks.py

@@ -56,12 +56,7 @@ class EfficientUpdateFormer(nn.Module):
 
         self.time_blocks = nn.ModuleList(
             [
-                AttnBlock(
-                    hidden_size,
-                    num_heads,
-                    mlp_ratio=mlp_ratio,
-                    attn_class=nn.MultiheadAttention,
-                )
+                AttnBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attn_class=nn.MultiheadAttention)
                 for _ in range(time_depth)
             ]
         )
@@ -69,12 +64,7 @@ class EfficientUpdateFormer(nn.Module):
         if add_space_attn:
             self.space_virtual_blocks = nn.ModuleList(
                 [
-                    AttnBlock(
-                        hidden_size,
-                        num_heads,
-                        mlp_ratio=mlp_ratio,
-                        attn_class=nn.MultiheadAttention,
-                    )
+                    AttnBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attn_class=nn.MultiheadAttention)
                     for _ in range(space_depth)
                 ]
             )

vggt/heads/track_modules/modules.py

@@ -45,20 +45,9 @@ class ResidualBlock(nn.Module):
         super(ResidualBlock, self).__init__()
 
         self.conv1 = nn.Conv2d(
-            in_planes,
-            planes,
-            kernel_size=kernel_size,
-            padding=1,
-            stride=stride,
-            padding_mode="zeros",
-        )
-        self.conv2 = nn.Conv2d(
-            planes,
-            planes,
-            kernel_size=kernel_size,
-            padding=1,
-            padding_mode="zeros",
+            in_planes, planes, kernel_size=kernel_size, padding=1, stride=stride, padding_mode="zeros"
         )
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, padding=1, padding_mode="zeros")
         self.relu = nn.ReLU(inplace=True)
 
         num_groups = planes // 8
@@ -92,10 +81,7 @@ class ResidualBlock(nn.Module):
         if stride == 1:
             self.downsample = None
         else:
-            self.downsample = nn.Sequential(
-                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride),
-                self.norm3,
-            )
+            self.downsample = nn.Sequential(nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
 
     def forward(self, x):
         y = x
@@ -151,7 +137,7 @@ class AttnBlock(nn.Module):
         num_heads,
         attn_class: Callable[..., nn.Module] = nn.MultiheadAttention,
         mlp_ratio=4.0,
-        **block_kwargs
+        **block_kwargs,
     ):
         """
         Self attention block

vggt/heads/track_modules/utils.py

@@ -36,10 +36,7 @@ def get_2d_sincos_pos_embed(embed_dim: int, grid_size: Union[int, Tuple[int, int
     grid = grid.reshape([2, 1, grid_size_h, grid_size_w])
     pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
     if return_grid:
-        return (
-            pos_embed.reshape(1, grid_size_h, grid_size_w, -1).permute(0, 3, 1, 2),
-            grid,
-        )
+        return (pos_embed.reshape(1, grid_size_h, grid_size_w, -1).permute(0, 3, 1, 2), grid)
     return pos_embed.reshape(1, grid_size_h, grid_size_w, -1).permute(0, 3, 1, 2)
 
 

vggt/layers/attention.py

@@ -58,12 +58,7 @@ class Attention(nn.Module):
             k = self.rope(k, pos)
 
         if self.fused_attn:
-            x = F.scaled_dot_product_attention(
-                q,
-                k,
-                v,
-                dropout_p=self.attn_drop.p if self.training else 0.0,
-            )
+            x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0)
         else:
             q = q * self.scale
             attn = q @ k.transpose(-2, -1)

vggt/layers/block.py

@@ -67,11 +67,7 @@ class Block(nn.Module):
         self.norm2 = norm_layer(dim)
         mlp_hidden_dim = int(dim * mlp_ratio)
         self.mlp = ffn_layer(
-            in_features=dim,
-            hidden_features=mlp_hidden_dim,
-            act_layer=act_layer,
-            drop=drop,
-            bias=ffn_bias,
+            in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop, bias=ffn_bias
         )
         self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
         self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
@@ -88,15 +84,10 @@ class Block(nn.Module):
         if self.training and self.sample_drop_ratio > 0.1:
             # the overhead is compensated only for a drop path rate larger than 0.1
             x = drop_add_residual_stochastic_depth(
-                x,
-                pos=pos,
-                residual_func=attn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
+                x, pos=pos, residual_func=attn_residual_func, sample_drop_ratio=self.sample_drop_ratio
             )
             x = drop_add_residual_stochastic_depth(
-                x,
-                residual_func=ffn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
+                x, residual_func=ffn_residual_func, sample_drop_ratio=self.sample_drop_ratio
             )
         elif self.training and self.sample_drop_ratio > 0.0:
             x = x + self.drop_path1(attn_residual_func(x, pos=pos))
@@ -108,10 +99,7 @@ class Block(nn.Module):
 
 
 def drop_add_residual_stochastic_depth(
-    x: Tensor,
-    residual_func: Callable[[Tensor], Tensor],
-    sample_drop_ratio: float = 0.0,
-    pos=None,
+    x: Tensor, residual_func: Callable[[Tensor], Tensor], sample_drop_ratio: float = 0.0, pos=None
 ) -> Tensor:
     # 1) extract subset using permutation
     b, n, d = x.shape
@@ -226,13 +214,13 @@ class NestedTensorBlock(Block):
                 x_list,
                 residual_func=attn_residual_func,
                 sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
+                scaling_vector=(self.ls1.gamma if isinstance(self.ls1, LayerScale) else None),
             )
             x_list = drop_add_residual_stochastic_depth_list(
                 x_list,
                 residual_func=ffn_residual_func,
                 sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
+                scaling_vector=(self.ls2.gamma if isinstance(self.ls1, LayerScale) else None),
             )
             return x_list
         else:

vggt/layers/layer_scale.py

@@ -13,12 +13,7 @@ from torch import nn
 
 
 class LayerScale(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        init_values: Union[float, Tensor] = 1e-5,
-        inplace: bool = False,
-    ) -> None:
+    def __init__(self, dim: int, init_values: Union[float, Tensor] = 1e-5, inplace: bool = False) -> None:
         super().__init__()
         self.inplace = inplace
         self.gamma = nn.Parameter(init_values * torch.ones(dim))

vggt/layers/patch_embed.py

@@ -47,10 +47,7 @@ class PatchEmbed(nn.Module):
 
         image_HW = make_2tuple(img_size)
         patch_HW = make_2tuple(patch_size)
-        patch_grid_size = (
-            image_HW[0] // patch_HW[0],
-            image_HW[1] // patch_HW[1],
-        )
+        patch_grid_size = (image_HW[0] // patch_HW[0], image_HW[1] // patch_HW[1])
 
         self.img_size = image_HW
         self.patch_size = patch_HW

vggt/layers/swiglu_ffn.py

@@ -64,9 +64,4 @@ class SwiGLUFFNFused(SwiGLU):
         out_features = out_features or in_features
         hidden_features = hidden_features or in_features
         hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
-        super().__init__(
-            in_features=in_features,
-            hidden_features=hidden_features,
-            out_features=out_features,
-            bias=bias,
-        )
+        super().__init__(in_features=in_features, hidden_features=hidden_features, out_features=out_features, bias=bias)

vggt/layers/vision_transformer.py

@@ -224,14 +224,7 @@ class DinoVisionTransformer(nn.Module):
         x = x + self.interpolate_pos_encoding(x, w, h)
 
         if self.register_tokens is not None:
-            x = torch.cat(
-                (
-                    x[:, :1],
-                    self.register_tokens.expand(x.shape[0], -1, -1),
-                    x[:, 1:],
-                ),
-                dim=1,
-            )
+            x = torch.cat((x[:, :1], self.register_tokens.expand(x.shape[0], -1, -1), x[:, 1:]), dim=1)
 
         return x
 

vggt/models/aggregator.py

@@ -133,15 +133,8 @@ class Aggregator(nn.Module):
         nn.init.normal_(self.register_token, std=1e-6)
 
         # Register normalization constants as buffers
-        for name, value in (
-            ("_resnet_mean", _RESNET_MEAN),
-            ("_resnet_std", _RESNET_STD),
-        ):
-            self.register_buffer(
-                name,
-                torch.FloatTensor(value).view(1, 1, 3, 1, 1),
-                persistent=False,
-            )
+        for name, value in (("_resnet_mean", _RESNET_MEAN), ("_resnet_std", _RESNET_STD)):
+            self.register_buffer(name, torch.FloatTensor(value).view(1, 1, 3, 1, 1), persistent=False)
 
     def __build_patch_embed__(
         self,
@@ -184,10 +177,7 @@ class Aggregator(nn.Module):
             if hasattr(self.patch_embed, "mask_token"):
                 self.patch_embed.mask_token.requires_grad_(False)
 
-    def forward(
-        self,
-        images: torch.Tensor,
-    ) -> Tuple[List[torch.Tensor], int]:
+    def forward(self, images: torch.Tensor) -> Tuple[List[torch.Tensor], int]:
         """
         Args:
             images (torch.Tensor): Input images with shape [B, S, 3, H, W], in range [0, 1].

vggt/models/vggt.py

@@ -24,11 +24,7 @@ class VGGT(nn.Module, PyTorchModelHubMixin):
         self.depth_head = DPTHead(dim_in=2 * embed_dim, output_dim=2, activation="exp", conf_activation="expp1")
         self.track_head = TrackHead(dim_in=2 * embed_dim, patch_size=patch_size)
 
-    def forward(
-        self,
-        images: torch.Tensor,
-        query_points: torch.Tensor = None,
-    ):
+    def forward(self, images: torch.Tensor, query_points: torch.Tensor = None):
         """
         Forward pass of the VGGT model.
 

vggt/utils/geometry.py

@@ -42,10 +42,7 @@ def unproject_depth_map_to_point_map(
 
 
 def depth_to_world_coords_points(
-    depth_map: np.ndarray,
-    extrinsic: np.ndarray,
-    intrinsic: np.ndarray,
-    eps=1e-8,
+    depth_map: np.ndarray, extrinsic: np.ndarray, intrinsic: np.ndarray, eps=1e-8
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """
     Convert a depth map to world coordinates.

vggt/utils/helper.py

@@ -0,0 +1,60 @@
+# 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 numpy as np
+
+
+def randomly_limit_trues(mask: np.ndarray, max_trues: int) -> np.ndarray:
+    """
+    If mask has more than max_trues True values,
+    randomly keep only max_trues of them and set the rest to False.
+    """
+    # 1D positions of all True entries
+    true_indices = np.flatnonzero(mask)  # shape = (N_true,)
+
+    # if already within budget, return as-is
+    if true_indices.size <= max_trues:
+        return mask
+
+    # randomly pick which True positions to keep
+    sampled_indices = np.random.choice(true_indices, size=max_trues, replace=False)  # shape = (max_trues,)
+
+    # build new flat mask: True only at sampled positions
+    limited_flat_mask = np.zeros(mask.size, dtype=bool)
+    limited_flat_mask[sampled_indices] = True
+
+    # restore original shape
+    return limited_flat_mask.reshape(mask.shape)
+
+
+def create_pixel_coordinate_grid(num_frames, height, width):
+    """
+    Creates a grid of pixel coordinates and frame indices for all frames.
+    Returns:
+        tuple: A tuple containing:
+            - points_xyf (numpy.ndarray): Array of shape (num_frames, height, width, 3)
+                                            with x, y coordinates and frame indices
+            - y_coords (numpy.ndarray): Array of y coordinates for all frames
+            - x_coords (numpy.ndarray): Array of x coordinates for all frames
+            - f_coords (numpy.ndarray): Array of frame indices for all frames
+    """
+    # Create coordinate grids for a single frame
+    y_grid, x_grid = np.indices((height, width), dtype=np.float32)
+    x_grid = x_grid[np.newaxis, :, :]
+    y_grid = y_grid[np.newaxis, :, :]
+
+    # Broadcast to all frames
+    x_coords = np.broadcast_to(x_grid, (num_frames, height, width))
+    y_coords = np.broadcast_to(y_grid, (num_frames, height, width))
+
+    # Create frame indices and broadcast
+    f_idx = np.arange(num_frames, dtype=np.float32)[:, np.newaxis, np.newaxis]
+    f_coords = np.broadcast_to(f_idx, (num_frames, height, width))
+
+    # Stack coordinates and frame indices
+    points_xyf = np.stack((x_coords, y_coords, f_coords), axis=-1)
+
+    return points_xyf

vggt/utils/load_fn.py

@@ -7,6 +7,91 @@
 import torch
 from PIL import Image
 from torchvision import transforms as TF
+import numpy as np
+
+
+def load_and_preprocess_images_square(image_path_list, target_size=1024):
+    """
+    Load and preprocess images by center padding to square and resizing to target size.
+    Also returns the position information of original pixels after transformation.
+
+    Args:
+        image_path_list (list): List of paths to image files
+        target_size (int, optional): Target size for both width and height. Defaults to 518.
+
+    Returns:
+        tuple: (
+            torch.Tensor: Batched tensor of preprocessed images with shape (N, 3, target_size, target_size),
+            torch.Tensor: Array of shape (N, 5) containing [x1, y1, x2, y2, width, height] for each image
+        )
+
+    Raises:
+        ValueError: If the input list is empty
+    """
+    # Check for empty list
+    if len(image_path_list) == 0:
+        raise ValueError("At least 1 image is required")
+
+    images = []
+    original_coords = []  # Renamed from position_info to be more descriptive
+    to_tensor = TF.ToTensor()
+
+    for image_path in image_path_list:
+        # Open image
+        img = Image.open(image_path)
+
+        # If there's an alpha channel, blend onto white background
+        if img.mode == "RGBA":
+            background = Image.new("RGBA", img.size, (255, 255, 255, 255))
+            img = Image.alpha_composite(background, img)
+
+        # Convert to RGB
+        img = img.convert("RGB")
+
+        # Get original dimensions
+        width, height = img.size
+
+        # Make the image square by padding the shorter dimension
+        max_dim = max(width, height)
+
+        # Calculate padding
+        left = (max_dim - width) // 2
+        top = (max_dim - height) // 2
+
+        # Calculate scale factor for resizing
+        scale = target_size / max_dim
+
+        # Calculate final coordinates of original image in target space
+        x1 = left * scale
+        y1 = top * scale
+        x2 = (left + width) * scale
+        y2 = (top + height) * scale
+
+        # Store original image coordinates and scale
+        original_coords.append(np.array([x1, y1, x2, y2, width, height]))
+
+        # Create a new black square image and paste original
+        square_img = Image.new("RGB", (max_dim, max_dim), (0, 0, 0))
+        square_img.paste(img, (left, top))
+
+        # Resize to target size
+        square_img = square_img.resize((target_size, target_size), Image.Resampling.BICUBIC)
+
+        # Convert to tensor
+        img_tensor = to_tensor(square_img)
+        images.append(img_tensor)
+
+    # Stack all images
+    images = torch.stack(images)
+    original_coords = torch.from_numpy(np.array(original_coords)).float()
+
+    # Add additional dimension if single image to ensure correct shape
+    if len(image_path_list) == 1:
+        if images.dim() == 3:
+            images = images.unsqueeze(0)
+            original_coords = original_coords.unsqueeze(0)
+
+    return images, original_coords
 
 
 def load_and_preprocess_images(image_path_list, mode="crop"):
@@ -39,7 +124,7 @@ def load_and_preprocess_images(image_path_list, mode="crop"):
     # Check for empty list
     if len(image_path_list) == 0:
         raise ValueError("At least 1 image is required")
-    
+
     # Validate mode
     if mode not in ["crop", "pad"]:
         raise ValueError("Mode must be either 'crop' or 'pad'")
@@ -51,7 +136,6 @@ def load_and_preprocess_images(image_path_list, mode="crop"):
 
     # First process all images and collect their shapes
     for image_path in image_path_list:
-
         # Open image
         img = Image.open(image_path)
 
@@ -66,7 +150,7 @@ def load_and_preprocess_images(image_path_list, mode="crop"):
         img = img.convert("RGB")
 
         width, height = img.size
-        
+
         if mode == "pad":
             # Make the largest dimension 518px while maintaining aspect ratio
             if width >= height:
@@ -89,18 +173,18 @@ def load_and_preprocess_images(image_path_list, mode="crop"):
         if mode == "crop" and new_height > target_size:
             start_y = (new_height - target_size) // 2
             img = img[:, start_y : start_y + target_size, :]
-        
+
         # For pad mode, pad to make a square of target_size x target_size
         if mode == "pad":
             h_padding = target_size - img.shape[1]
             w_padding = target_size - img.shape[2]
-            
+
             if h_padding > 0 or w_padding > 0:
                 pad_top = h_padding // 2
                 pad_bottom = h_padding - pad_top
                 pad_left = w_padding // 2
                 pad_right = w_padding - pad_left
-                
+
                 # Pad with white (value=1.0)
                 img = torch.nn.functional.pad(
                     img, (pad_left, pad_right, pad_top, pad_bottom), mode="constant", value=1.0

vggt/utils/pose_enc.py

@@ -9,10 +9,7 @@ from .rotation import quat_to_mat, mat_to_quat
 
 
 def extri_intri_to_pose_encoding(
-    extrinsics,
-    intrinsics,
-    image_size_hw=None,  # e.g., (256, 512)
-    pose_encoding_type="absT_quaR_FoV",
+    extrinsics, intrinsics, image_size_hw=None, pose_encoding_type="absT_quaR_FoV"  # e.g., (256, 512)
 ):
     """Convert camera extrinsics and intrinsics to a compact pose encoding.
 
@@ -63,10 +60,7 @@ def extri_intri_to_pose_encoding(
 
 
 def pose_encoding_to_extri_intri(
-    pose_encoding,
-    image_size_hw=None,  # e.g., (256, 512)
-    pose_encoding_type="absT_quaR_FoV",
-    build_intrinsics=True,
+    pose_encoding, image_size_hw=None, pose_encoding_type="absT_quaR_FoV", build_intrinsics=True  # e.g., (256, 512)
 ):
     """Convert a pose encoding back to camera extrinsics and intrinsics.
 

vggt/utils/rotation.py

@@ -63,13 +63,7 @@ def mat_to_quat(matrix: torch.Tensor) -> torch.Tensor:
 
     q_abs = _sqrt_positive_part(
         torch.stack(
-            [
-                1.0 + m00 + m11 + m22,
-                1.0 + m00 - m11 - m22,
-                1.0 - m00 + m11 - m22,
-                1.0 - m00 - m11 + m22,
-            ],
-            dim=-1,
+            [1.0 + m00 + m11 + m22, 1.0 + m00 - m11 - m22, 1.0 - m00 + m11 - m22, 1.0 - m00 - m11 + m22], dim=-1
         )
     )
 

vggt_to_colmap.py

@@ -1,596 +0,0 @@
-import os
-import argparse
-import numpy as np
-import torch
-import glob
-import struct
-from scipy.spatial.transform import Rotation
-import sys
-from PIL import Image
-import cv2
-import requests
-import tempfile
-
-sys.path.append("vggt/")
-
-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
-
-def load_model(device=None):
-    """Load and initialize the VGGT model."""
-    if device is 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):
-    """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'))])
-    print(f"Found {len(image_names)} images")
-    
-    if len(image_names) == 0:
-        raise ValueError(f"No images found in {image_dir}")
-
-    original_images = []
-    for img_path in image_names:
-        img = Image.open(img_path).convert('RGB')
-        original_images.append(np.array(img))
-    
-    images = load_and_preprocess_images(image_names).to(device)
-    print(f"Preprocessed images shape: {images.shape}")
-    
-    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)
-
-    print("Converting pose encoding to camera parameters...")
-    extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images.shape[-2:])
-    predictions["extrinsic"] = extrinsic
-    predictions["intrinsic"] = intrinsic
-
-    for key in predictions.keys():
-        if isinstance(predictions[key], torch.Tensor):
-            predictions[key] = predictions[key].cpu().numpy().squeeze(0)  # remove batch dimension
-    
-    print("Computing 3D points from depth maps...")
-    depth_map = predictions["depth"]  # (S, H, W, 1)
-    world_points = unproject_depth_map_to_point_map(depth_map, predictions["extrinsic"], predictions["intrinsic"])
-    predictions["world_points_from_depth"] = world_points
-    
-    predictions["original_images"] = original_images
-    
-    S, H, W = world_points.shape[:3]
-    normalized_images = np.zeros((S, H, W, 3), dtype=np.float32)
-    
-    for i, img in enumerate(original_images):
-        resized_img = cv2.resize(img, (W, H))
-        normalized_images[i] = resized_img / 255.0
-    
-    predictions["images"] = normalized_images
-    
-    return predictions, image_names
-
-def extrinsic_to_colmap_format(extrinsics):
-    """Convert extrinsic matrices to COLMAP format (quaternion + translation)."""
-    num_cameras = extrinsics.shape[0]
-    quaternions = []
-    translations = []
-    
-    for i in range(num_cameras):
-        # VGGT's extrinsic is camera-to-world (R|t) format
-        R = extrinsics[i, :3, :3]
-        t = extrinsics[i, :3, 3]
-        
-        # Convert rotation matrix to quaternion
-        # COLMAP quaternion format is [qw, qx, qy, qz]
-        rot = Rotation.from_matrix(R)
-        quat = rot.as_quat()  # scipy returns [x, y, z, w]
-        quat = np.array([quat[3], quat[0], quat[1], quat[2]])  # Convert to [w, x, y, z]
-        
-        quaternions.append(quat)
-        translations.append(t)
-    
-    return np.array(quaternions), np.array(translations)
-
-def download_file_from_url(url, filename):
-    """Downloads a file from a URL, handling redirects."""
-    try:
-        response = requests.get(url, allow_redirects=False)
-        response.raise_for_status() 
-
-        if response.status_code == 302:  
-            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()
-
-        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.
-    """
-    image = cv2.imread(image_path)
-
-    result_map = run_skyseg(onnx_session, [320, 320], image)
-    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
-
-    if mask_filename is not None:
-        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
-    
-    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")
-
-    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})
-
-    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=1, prediction_mode="Depthmap and Camera Branch"):
-    """
-    Filter points based on confidence and prepare for COLMAP format.
-    Implementation matches the conventions in the original VGGT code.
-    """
-    
-    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]))
-
-    colors_rgb = predictions["images"] 
-    
-    S, H, W = pred_world_points.shape[:3]
-    if colors_rgb.shape[:3] != (S, H, W):
-        print(f"Reshaping colors_rgb from {colors_rgb.shape} to match {(S, H, W, 3)}")
-        reshaped_colors = np.zeros((S, H, W, 3), dtype=np.float32)
-        for i in range(S):
-            if i < len(colors_rgb):
-                reshaped_colors[i] = cv2.resize(colors_rgb[i], (W, H))
-        colors_rgb = reshaped_colors
-    
-    colors_rgb = (colors_rgb * 255).astype(np.uint8)
-    
-    if mask_sky:
-        print("Applying sky segmentation mask")
-        try:
-            import onnxruntime
-         
-            with tempfile.TemporaryDirectory() as temp_dir:
-                print(f"Created temporary directory for sky segmentation: {temp_dir}")
-                temp_images_dir = os.path.join(temp_dir, "images")
-                sky_masks_dir = os.path.join(temp_dir, "sky_masks")
-                os.makedirs(temp_images_dir, exist_ok=True)
-                os.makedirs(sky_masks_dir, exist_ok=True)
-                
-                image_list = []
-                for i, img in enumerate(colors_rgb):
-                    img_path = os.path.join(temp_images_dir, f"image_{i:04d}.png")
-                    image_list.append(img_path)
-                    cv2.imwrite(img_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
-                
-           
-                skyseg_path = os.path.join(temp_dir, "skyseg.onnx")
-                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_path
-                    )
-                    if not download_success:
-                        print("Failed to download skyseg model, skipping sky filtering")
-                        mask_sky = False
-                else:
-            
-                    import shutil
-                    shutil.copy("skyseg.onnx", skyseg_path)
-                
-                if mask_sky:  
-                    skyseg_session = onnxruntime.InferenceSession(skyseg_path)
-                    sky_mask_list = []
-                    
-                    for img_path in image_list:
-                        mask_path = os.path.join(sky_masks_dir, os.path.basename(img_path))
-                        sky_mask = segment_sky(img_path, skyseg_session, mask_path)
-           
-                        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)
-                    
-                    sky_mask_array = np.array(sky_mask_list)
-                    
-                    sky_mask_binary = (sky_mask_array > 0.1).astype(np.float32)
-                    pred_world_points_conf = pred_world_points_conf * sky_mask_binary
-                    print(f"Applied sky mask, shape: {sky_mask_binary.shape}")
-                
-        except (ImportError, Exception) as e:
-            print(f"Error in sky segmentation: {e}")
-            mask_sky = False
-    
-    vertices_3d = pred_world_points.reshape(-1, 3)
-    conf = pred_world_points_conf.reshape(-1)
-    colors_rgb_flat = colors_rgb.reshape(-1, 3)
-
-    
-
-    if len(conf) != len(colors_rgb_flat):
-        print(f"WARNING: Shape mismatch between confidence ({len(conf)}) and colors ({len(colors_rgb_flat)})")
-        min_size = min(len(conf), len(colors_rgb_flat))
-        conf = conf[:min_size]
-        vertices_3d = vertices_3d[:min_size]
-        colors_rgb_flat = colors_rgb_flat[:min_size]
-    
-    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)
-    
-    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
-    
-    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
-    
-    filtered_vertices = vertices_3d[conf_mask]
-    filtered_colors = colors_rgb_flat[conf_mask]
-    
-    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")
-    
-    points3D = []
-    point_indices = {}
-    image_points2D = [[] for _ in range(len(pred_world_points))]
-    
-    print(f"Preparing points for COLMAP format with stride {stride}...")
-    
-    total_points = 0
-    for img_idx in range(S):
-        for y in range(0, H, stride):
-            for x in range(0, W, stride):
-                flat_idx = img_idx * H * W + y * W + x
-                
-                if flat_idx >= len(conf):
-                    continue
-                
-                if conf[flat_idx] < conf_thres_value or conf[flat_idx] <= 1e-5:
-                    continue
-                
-                if mask_black_bg and colors_rgb_flat[flat_idx].sum() < 16:
-                    continue
-                
-                if mask_white_bg and all(colors_rgb_flat[flat_idx] > 240):
-                    continue
-                
-                point3D = vertices_3d[flat_idx]
-                rgb = colors_rgb_flat[flat_idx]
-                
-                if not np.all(np.isfinite(point3D)):
-                    continue
-                
-                point_hash = hash_point(point3D, scale=100)
-                
-                if point_hash not in point_indices:
-                    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)
-                    total_points += 1
-                else:
-                    point_idx = point_indices[point_hash]
-                    points3D[point_idx]["track"].append((img_idx, len(image_points2D[img_idx])))
-                
-                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):
-    """Create a hash for a 3D point by quantizing coordinates."""
-    quantized = tuple(np.round(point * scale).astype(int))
-    return hash(quantized)
-
-def write_colmap_cameras_txt(file_path, intrinsics, image_width, image_height):
-    """Write camera intrinsics to COLMAP cameras.txt format."""
-    with open(file_path, '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, intrinsic in enumerate(intrinsics):
-            camera_id = i + 1  # COLMAP uses 1-indexed camera IDs
-            model = "PINHOLE" 
-            
-            fx = intrinsic[0, 0]
-            fy = intrinsic[1, 1]
-            cx = intrinsic[0, 2]
-            cy = intrinsic[1, 2]
-            
-            f.write(f"{camera_id} {model} {image_width} {image_height} {fx} {fy} {cx} {cy}\n")
-
-def write_colmap_images_txt(file_path, quaternions, translations, image_points2D, image_names):
-    """Write camera poses and keypoints to COLMAP images.txt format."""
-    with open(file_path, '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")
-        
-        num_points = sum(len(points) for points in image_points2D)
-        avg_points = num_points / len(image_points2D) if image_points2D else 0
-        f.write(f"# Number of images: {len(quaternions)}, mean observations per image: {avg_points:.1f}\n")
-        
-        for i in range(len(quaternions)):
-            image_id = i + 1 
-            camera_id = i + 1  
-          
-            qw, qx, qy, qz = quaternions[i]
-            tx, ty, tz = translations[i]
-            
-            f.write(f"{image_id} {qw} {qx} {qy} {qz} {tx} {ty} {tz} {camera_id} {os.path.basename(image_names[i])}\n")
-            
-            points_line = " ".join([f"{x} {y} {point3d_id+1}" for x, y, point3d_id in image_points2D[i]])
-            f.write(f"{points_line}\n")
-
-def write_colmap_points3D_txt(file_path, points3D):
-    """Write 3D points and tracks to COLMAP points3D.txt format."""
-    with open(file_path, '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")
-        
-        avg_track_length = sum(len(point["track"]) for point in points3D) / len(points3D) if points3D else 0
-        f.write(f"# Number of points: {len(points3D)}, mean track length: {avg_track_length:.4f}\n")
-        
-        for point in points3D:
-            point_id = point["id"] + 1  
-            x, y, z = point["xyz"]
-            r, g, b = point["rgb"]
-            error = point["error"]
-            
-            track = " ".join([f"{img_id+1} {point2d_idx}" for img_id, point2d_idx in point["track"]])
-            
-            f.write(f"{point_id} {x} {y} {z} {int(r)} {int(g)} {int(b)} {error} {track}\n")
-
-def write_colmap_cameras_bin(file_path, intrinsics, image_width, image_height):
-    """Write camera intrinsics to COLMAP cameras.bin format."""
-    with open(file_path, 'wb') as fid:
-        # Write number of cameras (uint64)
-        fid.write(struct.pack('<Q', len(intrinsics)))
-        
-        for i, intrinsic in enumerate(intrinsics):
-            camera_id = i + 1
-            model_id = 1 
-            
-            fx = float(intrinsic[0, 0])
-            fy = float(intrinsic[1, 1])
-            cx = float(intrinsic[0, 2])
-            cy = float(intrinsic[1, 2])
-            
-            # Camera ID (uint32)
-            fid.write(struct.pack('<I', camera_id))
-            # Model ID (uint32)
-            fid.write(struct.pack('<I', model_id))
-            # Width (uint64)
-            fid.write(struct.pack('<Q', image_width))
-            # Height (uint64)
-            fid.write(struct.pack('<Q', image_height))
-            
-            # Parameters (double)
-            fid.write(struct.pack('<dddd', fx, fy, cx, cy))
-
-def write_colmap_images_bin(file_path, quaternions, translations, image_points2D, image_names):
-    """Write camera poses and keypoints to COLMAP images.bin format."""
-    with open(file_path, 'wb') as fid:
-        # Write number of images (uint64)
-        fid.write(struct.pack('<Q', len(quaternions)))
-        
-        for i in range(len(quaternions)):
-            image_id = i + 1
-            camera_id = i + 1
-            
-            qw, qx, qy, qz = quaternions[i].astype(float)
-            tx, ty, tz = translations[i].astype(float)
-            
-            image_name = os.path.basename(image_names[i]).encode()
-            points = image_points2D[i]
-            
-            # Image ID (uint32)
-            fid.write(struct.pack('<I', image_id))
-            # Quaternion (double): qw, qx, qy, qz
-            fid.write(struct.pack('<dddd', qw, qx, qy, qz))
-            # Translation (double): tx, ty, tz
-            fid.write(struct.pack('<ddd', tx, ty, tz))
-            # Camera ID (uint32)
-            fid.write(struct.pack('<I', camera_id))
-            # Image name
-            for char in image_name:
-                fid.write(struct.pack('<c', bytes([char])))
-            fid.write(struct.pack('<c', b'\x00'))
-            
-            # Write number of 2D points (uint64)
-            fid.write(struct.pack('<Q', len(points)))
-            
-            # Write 2D points: x, y, point3D_id
-            for x, y, point3d_id in points:
-                fid.write(struct.pack('<dd', float(x), float(y)))
-                fid.write(struct.pack('<Q', point3d_id + 1))
-
-def write_colmap_points3D_bin(file_path, points3D):
-    """Write 3D points and tracks to COLMAP points3D.bin format."""
-    with open(file_path, 'wb') as fid:
-        # Write number of points (uint64)
-        fid.write(struct.pack('<Q', len(points3D)))
-        
-        for point in points3D:
-            point_id = point["id"] + 1
-            x, y, z = point["xyz"].astype(float)
-            r, g, b = point["rgb"].astype(np.uint8)
-            error = float(point["error"])
-            track = point["track"]
-            
-            # Point ID (uint64)
-            fid.write(struct.pack('<Q', point_id))
-            # Position (double): x, y, z
-            fid.write(struct.pack('<ddd', x, y, z))
-            # Color (uint8): r, g, b
-            fid.write(struct.pack('<BBB', int(r), int(g), int(b)))
-            # Error (double)
-            fid.write(struct.pack('<d', error))
-            
-            # Track: list of (image_id, point2D_idx)
-            fid.write(struct.pack('<Q', len(track)))
-            for img_id, point2d_idx in track:
-                fid.write(struct.pack('<II', img_id + 1, point2d_idx))
-
-def main():
-    parser = argparse.ArgumentParser(description="Convert images to COLMAP format using VGGT")
-    parser.add_argument("--image_dir", type=str, required=True, 
-                        help="Directory containing input images")
-    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")
-    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",
-                        help="Filter out points with very dark/black color")
-    parser.add_argument("--mask_white_bg", action="store_true",
-                        help="Filter out points with very bright/white color")
-    parser.add_argument("--binary", action="store_true", 
-                        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()
-    
-    os.makedirs(args.output_dir, exist_ok=True)
-    
-    model, device = load_model()
-    
-    predictions, image_names = process_images(args.image_dir, model, device)
-    
-    print("Converting camera parameters to COLMAP format...")
-    quaternions, translations = extrinsic_to_colmap_format(predictions["extrinsic"])
-    
-    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, 
-        mask_sky=args.mask_sky, 
-        mask_black_bg=args.mask_black_bg,
-        mask_white_bg=args.mask_white_bg,
-        stride=args.stride,
-        prediction_mode=args.prediction_mode
-    )
-    
-    height, width = predictions["depth"].shape[1:3]
-
-    print(f"Writing {'binary' if args.binary else 'text'} COLMAP files to {args.output_dir}...")
-    if args.binary:
-        write_colmap_cameras_bin(
-            os.path.join(args.output_dir, "cameras.bin"), 
-            predictions["intrinsic"], width, height)
-        write_colmap_images_bin(
-            os.path.join(args.output_dir, "images.bin"), 
-            quaternions, translations, image_points2D, image_names)
-        write_colmap_points3D_bin(
-            os.path.join(args.output_dir, "points3D.bin"), 
-            points3D)
-    else:
-        write_colmap_cameras_txt(
-            os.path.join(args.output_dir, "cameras.txt"), 
-            predictions["intrinsic"], width, height)
-        write_colmap_images_txt(
-            os.path.join(args.output_dir, "images.txt"), 
-            quaternions, translations, image_points2D, image_names)
-        write_colmap_points3D_txt(
-            os.path.join(args.output_dir, "points3D.txt"), 
-            points3D)
-    
-    print(f"COLMAP files successfully written to {args.output_dir}")
-
-if __name__ == "__main__":
-    main()

visual_util.py

@@ -215,12 +215,7 @@ def predictions_to_glb(
     return scene_3d
 
 
-def integrate_camera_into_scene(
-    scene: trimesh.Scene,
-    transform: np.ndarray,
-    face_colors: tuple,
-    scene_scale: float,
-):
+def integrate_camera_into_scene(scene: trimesh.Scene, transform: np.ndarray, face_colors: tuple, scene_scale: float):
     """
     Integrates a fake camera mesh into the 3D scene.