Upload folder using huggingface_hub

colmap_matcher.sh

@@ -144,9 +144,9 @@ then
 fi
 
 # LightGlue Feature Matcher
-if [ "${matcher_type}" == "lightglue" ]
+if [ "${matcher_type}" == "custom" ]
 then
-  pixi run -e lightglue python3 Baselines/colmap/lightglue_matcher.py --database ${database} --rgb_path ${rgb_path} --feature superpoint
+  pixi run -e lightglue python3 Baselines/colmap/feature_matcher.py --database ${database} --rgb_path ${rgb_path} --feature superpoint --matcher lightglue --use_gpu ${use_gpu}
   colmap matches_importer \
       --database_path ${database} \
       --match_list_path "${exp_folder_colmap}/matches.txt" \

feature_matcher.py

@@ -0,0 +1,359 @@
+import sqlite3
+from feature_matcher_utilities import extract_keypoints, feature_matching, unrotate_kps_W
+import os
+import torch
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+import numpy as np
+import cv2
+import argparse
+from pathlib import Path
+
+from lightglue import LightGlue
+
+# ==========================================
+# ==========================================
+# DATABASE UTILITIES
+# ==========================================
+def load_colmap_db(db_path): 
+    if not os.path.exists(db_path):
+        raise FileNotFoundError(f"Database file not found: {db_path}")
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+    return conn, cursor
+
+def create_pair_id(image_id1, image_id2):
+    if image_id1 > image_id2:
+        image_id1, image_id2 = image_id2, image_id1
+    return image_id1 * 2147483647 + image_id2
+
+def clean_database(cursor):
+    """Removes existing features and matches to ensure a clean overwrite."""
+    tables = ["keypoints", "descriptors"]#, "matches"], "two_view_geometry"]
+    for table in tables:
+        cursor.execute(f"DELETE FROM {table};")
+    print("Database cleaned (keypoints, descriptors, matches removed).")
+
+def insert_keypoints(cursor, image_id, keypoints, descriptors):
+    """
+    keypoints: (N, 2) numpy array, float32
+    descriptors: (N, D) numpy array, float32
+    """
+    keypoints_blob = keypoints.tobytes()
+    descriptors_blob = descriptors.tobytes()
+    
+    # Keypoints
+    cursor.execute(
+        "INSERT INTO keypoints(image_id, rows, cols, data) VALUES(?, ?, ?, ?)",
+        (image_id, keypoints.shape[0], keypoints.shape[1], keypoints_blob)
+    )
+    
+    # Descriptors (Optional but good practice)
+    cursor.execute(
+        "INSERT INTO descriptors(image_id, rows, cols, data) VALUES(?, ?, ?, ?)",
+        (image_id, descriptors.shape[0], descriptors.shape[1], descriptors_blob)
+    )
+
+def insert_matches(cursor, image_id1, image_id2, matches):
+    """
+    matches: (K, 2) numpy array, uint32. 
+             Col 0 is index in image1, Col 1 is index in image2
+    """
+    pair_id = create_pair_id(image_id1, image_id2)
+    matches_blob = matches.tobytes()
+    
+    cursor.execute(
+        "INSERT INTO matches(pair_id, rows, cols, data) VALUES(?, ?, ?, ?)",
+        (pair_id, matches.shape[0], matches.shape[1], matches_blob)
+    )
+
+def verify_matches_visual(cursor, image_id1, image_id2, image_dir):
+    """
+    Reads matches and keypoints from the COLMAP db and plots them.
+    
+    Args:
+        cursor: SQLite cursor connected to the database.
+        image_id1: ID of the first image.
+        image_id2: ID of the second image.
+        image_dir: Path to the directory containing the images.
+    """
+    
+    # 1. Helper to ensure image_id1 < image_id2 for pair_id calculation
+    if image_id1 > image_id2:
+        image_id1, image_id2 = image_id2, image_id1
+        swapped = True
+    else:
+        swapped = False
+        
+    pair_id = image_id1 * 2147483647 + image_id2
+
+    # 2. Fetch Matches
+    cursor.execute("SELECT data FROM matches WHERE pair_id = ?", (pair_id,))
+    match_row = cursor.fetchone()
+    
+    if match_row is None:
+        print(f"No matches found in DB for pair {image_id1}-{image_id2}")
+        return
+
+    # Decode Matches: UINT32 (N, 2)
+    matches = np.frombuffer(match_row[0], dtype=np.uint32).reshape(-1, 2)
+
+    # If we swapped inputs to generate pair_id, we must swap columns in matches
+    # so matches[:,0] corresponds to the requested image_id1
+    if swapped:
+        matches = matches[:, [1, 0]]
+
+    # 3. Fetch Keypoints for both images
+    def get_keypoints_and_name(img_id):
+        # Get Name
+        cursor.execute("SELECT name FROM images WHERE image_id = ?", (img_id,))
+        name = cursor.fetchone()[0]
+        
+        # Get Keypoints
+        cursor.execute("SELECT data FROM keypoints WHERE image_id = ?", (img_id,))
+        kp_row = cursor.fetchone()
+        # Decode Keypoints: FLOAT32 (N, 2)
+        kpts = np.frombuffer(kp_row[0], dtype=np.float32).reshape(-1, 2)
+        return name, kpts
+
+    name1, kpts1 = get_keypoints_and_name(image_id1)
+    name2, kpts2 = get_keypoints_and_name(image_id2)
+
+    # 4. Filter Keypoints using the Matches indices
+    # matches[:, 0] are indices into kpts1
+    # matches[:, 1] are indices into kpts2
+    valid_kpts1 = kpts1[matches[:, 0]]
+    valid_kpts2 = kpts2[matches[:, 1]]
+
+    # 5. Load Images
+    path1 = os.path.join(image_dir, name1)
+    path2 = os.path.join(image_dir, name2)
+    
+    img1 = cv2.imread(path1)
+    img2 = cv2.imread(path2)
+    
+    # Convert BGR (OpenCV) to RGB (Matplotlib)
+    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
+    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
+
+    # 6. Plotting
+    # Concatenate images side-by-side
+    h1, w1, _ = img1.shape
+    h2, w2, _ = img2.shape
+    
+    # Create a canvas large enough for both
+    height = max(h1, h2)
+    width = w1 + w2
+    canvas = np.zeros((height, width, 3), dtype=np.uint8)
+    
+    canvas[:h1, :w1, :] = img1
+    canvas[:h2, w1:w1+w2, :] = img2
+
+    plt.figure(figsize=(15, 10))
+    plt.imshow(canvas)
+    
+    # Plot lines
+    # Shift x-coordinates of image2 by w1
+    for (x1, y1), (x2, y2) in zip(valid_kpts1, valid_kpts2):
+        plt.plot([x1, x2 + w1], [y1, y2], 'c-', alpha=0.6, linewidth=0.5)
+        plt.plot(x1, y1, 'r.', markersize=2)
+        plt.plot(x2 + w1, y2, 'r.', markersize=2)
+
+    plt.title(f"DB Verification: {name1} (ID:{image_id1}) <-> {name2} (ID:{image_id2}) | Matches: {len(matches)}")
+    plt.axis('off')
+    plt.tight_layout()
+    plt.show()
+
+import numpy as np
+import matplotlib.pyplot as plt
+import cv2
+import os
+import sqlite3
+
+def plot_matches_from_db(cursor, image_id1, image_id2, image_dir):
+    """
+    Reads matches and keypoints for a specific pair from the COLMAP DB and plots them.
+    
+    Args:
+        cursor: SQLite cursor.
+        image_id1, image_id2: The IDs of the two images to plot.
+        image_dir: Path to the directory containing the actual image files.
+    """
+    
+    # 1. Resolve Pair ID (Colmap requires id1 < id2 for unique pair_id)
+    if image_id1 > image_id2:
+        id_a, id_b = image_id2, image_id1
+        swapped = True
+    else:
+        id_a, id_b = image_id1, image_id2
+        swapped = False
+        
+    pair_id = id_a * 2147483647 + id_b
+
+    # 2. Fetch Matches
+    print(f"Fetching matches for pair {image_id1}-{image_id2} (PairID: {pair_id})...")
+    cursor.execute("SELECT data, rows, cols FROM matches WHERE pair_id = ?", (pair_id,))
+    match_row = cursor.fetchone()
+    
+    if match_row is None:
+        print(f"No matches found in database for Pair {image_id1}-{image_id2}")
+        return
+
+    # Decode Matches (UINT32)
+    # Blob is match_row[0], rows is [1], cols is [2]
+    matches_blob = match_row[0]
+    matches = np.frombuffer(matches_blob, dtype=np.uint32).reshape(-1, 2)
+    
+    # If inputs were swapped relative to how COLMAP stores them, swap the columns
+    # so matches[:,0] refers to image_id1 and matches[:,1] refers to image_id2
+    if swapped:
+        matches = matches[:, [1, 0]]
+
+    # 3. Fetch Keypoints & Image Names
+    def get_image_data(img_id):
+        cursor.execute("SELECT name FROM images WHERE image_id = ?", (img_id,))
+        res = cursor.fetchone()
+        if not res:
+            raise ValueError(f"Image ID {img_id} not found in 'images' table.")
+        name = res[0]
+        
+        cursor.execute("SELECT data FROM keypoints WHERE image_id = ?", (img_id,))
+        kp_res = cursor.fetchone()
+        if not kp_res:
+             raise ValueError(f"No keypoints found for Image ID {img_id}.")
+        
+        # Decode Keypoints (FLOAT32)
+        kpts = np.frombuffer(kp_res[0], dtype=np.float32).reshape(-1, 2)
+        return name, kpts
+
+    name1, kpts1 = get_image_data(image_id1)
+    name2, kpts2 = get_image_data(image_id2)
+
+    # 4. Filter Keypoints using Match Indices
+    valid_kpts1 = kpts1[matches[:, 0]]
+    valid_kpts2 = kpts2[matches[:, 1]]
+
+    # 5. Visualization
+    path1 = os.path.join(image_dir, name1)
+    path2 = os.path.join(image_dir, name2)
+    
+    if not os.path.exists(path1) or not os.path.exists(path2):
+        print(f"Error: Could not find image files at \n{path1}\n{path2}")
+        return
+
+    img1 = cv2.imread(path1)
+    img2 = cv2.imread(path2)
+    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
+    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
+
+    # Create canvas
+    h1, w1 = img1.shape[:2]
+    h2, w2 = img2.shape[:2]
+    height = max(h1, h2)
+    width = w1 + w2
+    canvas = np.zeros((height, width, 3), dtype=np.uint8)
+    canvas[:h1, :w1] = img1
+    canvas[:h2, w1:w1+w2] = img2
+
+    plt.figure(figsize=(20, 10))
+    plt.imshow(canvas)
+    
+    # Plot matches
+    # x2 coordinates need to be shifted by w1
+    for (x1, y1), (x2, y2) in zip(valid_kpts1, valid_kpts2):
+        plt.plot([x1, x2 + w1], [y1, y2], 'g-', alpha=0.5, linewidth=1.5)
+        plt.plot(x1, y1, 'r.', markersize=4)
+        plt.plot(x2 + w1, y2, 'r.', markersize=4)
+
+    plt.title(f"{name1} <-> {name2} | Total Matches: {len(matches)}")
+    plt.axis('off')
+    plt.tight_layout()
+    plt.show()
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument("--database", type=Path, required=True)
+    parser.add_argument("--rgb_path", type=Path, required=True)
+    parser.add_argument("--feature", type=str, required=True)
+    parser.add_argument("--matcher", type=str, required=True)
+
+    args, _ = parser.parse_known_args()
+
+    DB_PATH = args.database
+    IMAGE_DIR = args.rgb_path
+    FEATURE_TYPE = args.feature
+    MATCHER_TYPE = args.matcher
+    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+    matches_file_path = os.path.join(os.path.dirname(DB_PATH), "matches.txt")
+
+    conn, cursor = load_colmap_db(DB_PATH)
+    cursor.execute("SELECT image_id, name FROM images")
+    images_info = {row[0]: row[1] for row in cursor.fetchall()}
+    image_ids = sorted(images_info.keys())
+
+    clean_database(cursor)
+    conn.commit() 
+
+    # Keypoint Extraction
+    fts = {}
+    for i in tqdm(range(len(image_ids)), desc="Feature Extraction"):
+        id = image_ids[i]
+        fname = images_info[id]
+        path = os.path.join(IMAGE_DIR, fname)
+
+        feats_dict, h, w = extract_keypoints(path, features=FEATURE_TYPE)
+        
+        fts[id] = feats_dict
+
+        kpts = feats_dict['keypoints'].squeeze(0).cpu().numpy().astype(np.float32)
+        descs = feats_dict['descriptors'].squeeze(0).cpu().numpy().astype(np.float32)
+
+        if FEATURE_TYPE == 'superpoint':
+            kpts_rot = unrotate_kps_W(kpts, feats_dict['rotations'].squeeze(0).cpu().numpy().astype(np.float32), h, w)
+        else:
+            kpts_rot = kpts
+        insert_keypoints(cursor, id, kpts_rot, descs)
+
+    conn.commit() 
+
+    # Feature Matching
+    if MATCHER_TYPE == 'lightglue':
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        matcher = LightGlue(features='superpoint', depth_confidence=-1, width_confidence=-1, flash=True).eval().to(device)
+    else:
+        matcher = None
+
+    with open(matches_file_path, "w") as f_match:
+        for i in tqdm(range(len(image_ids)), desc="Feature Extraction"):
+            id1 = image_ids[i]
+            fname1 = images_info[id1]
+            path1 = os.path.join(IMAGE_DIR, fname1)
+
+            for j in range(i + 1, len(image_ids)):
+                if j == i:
+                    continue
+                id2 = image_ids[j]
+
+                fname2 = images_info[id2]
+                path2 = os.path.join(IMAGE_DIR, fname2)
+                matches_tensor = feature_matching(fts[id1], fts[id2], matcher=matcher, features=FEATURE_TYPE, matcher_type=MATCHER_TYPE) 
+
+                if matches_tensor is not None and len(matches_tensor) > 0:
+                    matches_np = matches_tensor.cpu().numpy().astype(np.uint32)
+                    #insert_matches(cursor, id1, id2, matches_np)
+                    
+                    f_match.write(f"{fname1} {fname2}\n")
+                    np.savetxt(f_match, matches_np, fmt="%d")
+                    f_match.write("\n")
+                
+                    #verify_matches_visual(cursor, image_ids[i], image_ids[j], IMAGE_DIR)                
+            #plt.show()
+    
+    conn.commit()
+
+    #plot_matches_from_db(cursor, image_ids[0], image_ids[1], IMAGE_DIR)
+
+    conn.close()
+    print("Database overwrite complete.")    

feature_matcher_utilities.py

@@ -0,0 +1,140 @@
+import torch
+import numpy as np
+import cv2
+from lightglue import LightGlue
+from lightglue.utils import rbd
+from lightglue import  SuperPoint, SIFT
+from lightglue.utils import load_image
+
+
+def unrotate_kps_W(kps_rot, k, H, W):
+    # Ensure inputs are Numpy
+    if hasattr(kps_rot, 'cpu'): kps_rot = kps_rot.cpu().numpy()
+    if hasattr(k, 'cpu'): k = k.cpu().numpy()
+    
+    # Squeeze if necessary
+    if k.ndim > 1: k = k.squeeze()
+    if kps_rot.ndim > 2: kps_rot = kps_rot.squeeze()
+
+    x_r = kps_rot[:, 0]
+    y_r = kps_rot[:, 1]
+    
+    x = np.zeros_like(x_r)
+    y = np.zeros_like(y_r)
+    
+    mask0 = (k == 0)
+    x[mask0], y[mask0] = x_r[mask0], y_r[mask0]
+    
+    mask1 = (k == 1)
+    x[mask1], y[mask1] = (W - 1) - y_r[mask1], x_r[mask1]
+    
+    mask2 = (k == 2)
+    x[mask2], y[mask2] = (W - 1) - x_r[mask2], (H - 1) - y_r[mask2]
+    
+    mask3 = (k == 3)
+    x[mask3], y[mask3] = y_r[mask3], (H - 1) - x_r[mask3]
+    
+    return np.stack([x, y], axis=-1)
+
+def extract_keypoints(path_to_image0, features='superpoint', rotations = [0,1,2,3]):
+    # --- Models on GPU ---
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    # --- Load images as Torch tensors (3,H,W) in [0,1] ---
+    timg = load_image(path_to_image0).to(device)
+    _, h, w = timg.shape
+
+    if features == 'sift':
+        extractor = SIFT(max_num_keypoints=2048).eval().to(device)
+        feats = extractor.extract(timg)
+        return feats , h, w
+    
+    if features == 'superpoint':
+        extractor = SuperPoint(max_num_keypoints=2048).eval().to(device)
+
+    # --- Extract local features ---
+    feats = {}
+    for k in (rotations):
+        timg_rotated = torch.rot90(timg, k, dims=(1, 2))
+        feats[k] = extractor.extract(timg_rotated)
+        #print(f"Extracted {feats[k]['keypoints'].shape[1]} keypoints for rotation {k}")
+
+    # --- Merge features back to original coordinate system ---
+    all_keypoints = []
+    all_scores = []
+    all_descriptors = []
+    all_rotations = []
+    for k, feat in feats.items():
+        kpts = feat['keypoints']  # Shape (1, N, 2)
+        num_kpts = kpts.shape[1]
+        # if k == 0:
+        #     kpts_corrected = kpts
+        # elif k == 1:
+        #     kpts_corrected = torch.stack(
+        #         [w - 1 - kpts[..., 1], kpts[..., 0]], dim=-1
+        #     )
+        # elif k == 2:
+        #     kpts_corrected = torch.stack(
+        #         [w - 1 - kpts[..., 0], h - 1 - kpts[..., 1]], dim=-1
+        #     )
+        # elif k == 3:
+        #     kpts_corrected = torch.stack(
+        #         [kpts[..., 1], h - 1 - kpts[..., 0]], dim=-1
+        #     )
+
+        rot_indices = torch.full((1, num_kpts), k, dtype=torch.long, device=device)
+        all_keypoints.append(feat['keypoints'])
+        all_scores.append(feat['keypoint_scores'])
+        all_descriptors.append(feat['descriptors'])
+        all_rotations.append(rot_indices)
+
+    # Concatenate all features along the keypoint dimension (dim=1)
+    feats_merged = {
+        'keypoints': torch.cat(all_keypoints, dim=1),
+        'keypoint_scores': torch.cat(all_scores, dim=1),
+        'descriptors': torch.cat(all_descriptors, dim=1),
+        'rotations': torch.cat(all_rotations, dim=1)
+    }
+    
+    num_kpts = feats_merged['keypoints'].shape[1]
+    # perm = torch.randperm(num_kpts, device=device)
+
+    # feats_merged['keypoints'] = feats_merged['keypoints'][:, perm, :]
+    # feats_merged['keypoint_scores'] = feats_merged['keypoint_scores'][:, perm]
+    # feats_merged['descriptors'] = feats_merged['descriptors'][:, perm, :]
+
+    # Optional: If you want to retain other keys like 'shape' or 'image_size'
+    #feats_merged['image_size'] = torch.tensor([w, h], device=device).unsqueeze(0)
+    #feats_merged['scales'] = torch.tensor([w, h], device=device).unsqueeze(0)
+
+    # for f in feats_merged:
+    #     if 'scales' not in f:
+    #         f['scales'] = torch.ones(all_keypoints.shape[:-1], device=device)
+    #     if 'oris' not in f:
+    #         f['oris'] = torch.zeros(all_keypoints.shape[:-1], device=device)
+
+    return feats_merged , h, w
+
+def feature_matching(feats0, feats1, matcher = None, features='superpoint', matcher_type='lightglue'):
+    #print(f"Matching features using {matcher_type} matcher with {features} features.")
+    if matcher_type == "lightglue":
+        if matcher is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+            matcher = LightGlue(features=features).eval().to(device)
+    
+        out_k = matcher({'image0': feats0, 'image1': feats1})
+        _, _, out_k = [rbd(x) for x in [feats0, feats1, out_k]]   # remove batch dim
+        return out_k['matches'] 
+
+    if matcher_type == "brute_force":
+        desc0 = feats0['descriptors'].squeeze(0) 
+        desc1 = feats1['descriptors'].squeeze(0) 
+        dists = torch.cdist(desc0, desc1, p=2)
+        min_dist0, idx0 = dists.min(dim=1)
+        min_dist1, idx1 = dists.min(dim=0)
+        indices0 = torch.arange(len(idx0), device='cuda')
+        mutual_mask = (idx1[idx0] == indices0) 
+        matches0 = indices0[mutual_mask]
+        matches1 = idx0[mutual_mask]
+            
+        return torch.stack([matches0, matches1], dim=-1) # [N, 2]