fix on gpu

ui/app_class.py

@@ -83,7 +83,6 @@ a:hover {
 }
 """
 
-@spaces.GPU
 class ImageMatchingApp:
     def __init__(self, server_name="0.0.0.0", server_port=7860, **kwargs):
         self.server_name = server_name

ui/utils.py

@@ -1,22 +1,21 @@
+import cv2
+import gradio as gr
+import matplotlib.pyplot as plt
+import numpy as np
 import os
 import pickle
+import poselib
+import psutil
 import random
 import shutil
 import sys
 import time
 import warnings
+from PIL import Image
 from itertools import combinations
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
-import cv2
-import gradio as gr
-import matplotlib.pyplot as plt
-import numpy as np
-import poselib
-import psutil
-from PIL import Image
-
 sys.path.append(str(Path(__file__).parents[1]))
 
 from hloc import (
@@ -30,6 +29,7 @@ from hloc import (
 )
 from hloc.utils.base_model import dynamic_load
 from ui.viz import display_keypoints, display_matches, fig2im, plot_images
+import spaces
 
 warnings.simplefilter("ignore")
 
@@ -130,7 +130,7 @@ def load_config(config_name: str) -> Dict[str, Any]:
 
 
 def get_matcher_zoo(
-    matcher_zoo: Dict[str, Dict[str, Union[str, bool]]]
+        matcher_zoo: Dict[str, Dict[str, Union[str, bool]]]
 ) -> Dict[str, Dict[str, Union[Callable, bool]]]:
     """
     Restore matcher configurations from a dictionary.
@@ -220,7 +220,7 @@ def gen_examples():
                 img1 = os.path.join(path, lines[i].strip())
                 img2 = os.path.join(path, lines[i + 1].strip())
                 image_pairs.append((img1, img2))
-        count=len(image_pairs)
+        count = len(image_pairs)
 
         if len(image_pairs) < count:
             count = len(image_pairs)
@@ -276,13 +276,13 @@ def set_null_pred(feature_type: str, pred: dict):
 
 
 def _filter_matches_opencv(
-    kp0: np.ndarray,
-    kp1: np.ndarray,
-    method: int = cv2.RANSAC,
-    reproj_threshold: float = 3.0,
-    confidence: float = 0.99,
-    max_iter: int = 2000,
-    geometry_type: str = "Homography",
+        kp0: np.ndarray,
+        kp1: np.ndarray,
+        method: int = cv2.RANSAC,
+        reproj_threshold: float = 3.0,
+        confidence: float = 0.99,
+        max_iter: int = 2000,
+        geometry_type: str = "Homography",
 ) -> Tuple[np.ndarray, np.ndarray]:
     """
     Filters matches between two sets of keypoints using OpenCV's findHomography.
@@ -322,13 +322,13 @@ def _filter_matches_opencv(
 
 
 def _filter_matches_poselib(
-    kp0: np.ndarray,
-    kp1: np.ndarray,
-    method: int = None,  # not used
-    reproj_threshold: float = 3,
-    confidence: float = 0.99,
-    max_iter: int = 2000,
-    geometry_type: str = "Homography",
+        kp0: np.ndarray,
+        kp1: np.ndarray,
+        method: int = None,  # not used
+        reproj_threshold: float = 3,
+        confidence: float = 0.99,
+        max_iter: int = 2000,
+        geometry_type: str = "Homography",
 ) -> dict:
     """
     Filters matches between two sets of keypoints using the poselib library.
@@ -364,13 +364,13 @@ def _filter_matches_poselib(
 
 
 def proc_ransac_matches(
-    mkpts0: np.ndarray,
-    mkpts1: np.ndarray,
-    ransac_method: str = DEFAULT_RANSAC_METHOD,
-    ransac_reproj_threshold: float = 3.0,
-    ransac_confidence: float = 0.99,
-    ransac_max_iter: int = 2000,
-    geometry_type: str = "Homography",
+        mkpts0: np.ndarray,
+        mkpts1: np.ndarray,
+        ransac_method: str = DEFAULT_RANSAC_METHOD,
+        ransac_reproj_threshold: float = 3.0,
+        ransac_confidence: float = 0.99,
+        ransac_max_iter: int = 2000,
+        geometry_type: str = "Homography",
 ):
     if ransac_method.startswith("CV2"):
         logger.info(
@@ -403,12 +403,12 @@ def proc_ransac_matches(
 
 
 def filter_matches(
-    pred: Dict[str, Any],
-    ransac_method: str = DEFAULT_RANSAC_METHOD,
-    ransac_reproj_threshold: float = DEFAULT_RANSAC_REPROJ_THRESHOLD,
-    ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
-    ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
-    ransac_estimator: str = None,
+        pred: Dict[str, Any],
+        ransac_method: str = DEFAULT_RANSAC_METHOD,
+        ransac_reproj_threshold: float = DEFAULT_RANSAC_REPROJ_THRESHOLD,
+        ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
+        ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
+        ransac_estimator: str = None,
 ):
     """
     Filter matches using RANSAC. If keypoints are available, filter by keypoints.
@@ -433,8 +433,8 @@ def filter_matches(
         mkpts1 = pred["mkeypoints1_orig"]
         feature_type = "KEYPOINT"
     elif (
-        "line_keypoints0_orig" in pred.keys()
-        and "line_keypoints1_orig" in pred.keys()
+            "line_keypoints0_orig" in pred.keys()
+            and "line_keypoints1_orig" in pred.keys()
     ):
         mkpts0 = pred["line_keypoints0_orig"]
         mkpts1 = pred["line_keypoints1_orig"]
@@ -477,11 +477,11 @@ def filter_matches(
 
 
 def compute_geometry(
-    pred: Dict[str, Any],
-    ransac_method: str = DEFAULT_RANSAC_METHOD,
-    ransac_reproj_threshold: float = DEFAULT_RANSAC_REPROJ_THRESHOLD,
-    ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
-    ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
+        pred: Dict[str, Any],
+        ransac_method: str = DEFAULT_RANSAC_METHOD,
+        ransac_reproj_threshold: float = DEFAULT_RANSAC_REPROJ_THRESHOLD,
+        ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
+        ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
 ) -> Dict[str, List[float]]:
     """
     Compute geometric information of matches, including Fundamental matrix,
@@ -504,8 +504,8 @@ def compute_geometry(
         mkpts0 = pred["mkeypoints0_orig"]
         mkpts1 = pred["mkeypoints1_orig"]
     elif (
-        "line_keypoints0_orig" in pred.keys()
-        and "line_keypoints1_orig" in pred.keys()
+            "line_keypoints0_orig" in pred.keys()
+            and "line_keypoints1_orig" in pred.keys()
     ):
         mkpts0 = pred["line_keypoints0_orig"]
         mkpts1 = pred["line_keypoints1_orig"]
@@ -561,10 +561,10 @@ def compute_geometry(
 
 
 def wrap_images(
-    img0: np.ndarray,
-    img1: np.ndarray,
-    geo_info: Optional[Dict[str, List[float]]],
-    geom_type: str,
+        img0: np.ndarray,
+        img1: np.ndarray,
+        geo_info: Optional[Dict[str, List[float]]],
+        geom_type: str,
 ) -> Tuple[Optional[str], Optional[Dict[str, List[float]]]]:
     """
     Wraps the images based on the geometric transformation used to align them.
@@ -617,10 +617,10 @@ def wrap_images(
 
 
 def generate_warp_images(
-    input_image0: np.ndarray,
-    input_image1: np.ndarray,
-    matches_info: Dict[str, Any],
-    choice: str,
+        input_image0: np.ndarray,
+        input_image1: np.ndarray,
+        matches_info: Dict[str, Any],
+        choice: str,
 ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
     """
     Changes the estimate of the geometric transformation used to align the images.
@@ -635,9 +635,9 @@ def generate_warp_images(
         A tuple containing the updated images and the warpped images.
     """
     if (
-        matches_info is None
-        or len(matches_info) < 1
-        or "geom_info" not in matches_info.keys()
+            matches_info is None
+            or len(matches_info) < 1
+            or "geom_info" not in matches_info.keys()
     ):
         return None, None
     geom_info = matches_info["geom_info"]
@@ -671,12 +671,12 @@ def send_to_match(state_cache: Dict[str, Any]):
 
 
 def run_ransac(
-    state_cache: Dict[str, Any],
-    choice_geometry_type: str,
-    ransac_method: str = DEFAULT_RANSAC_METHOD,
-    ransac_reproj_threshold: int = DEFAULT_RANSAC_REPROJ_THRESHOLD,
-    ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
-    ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
+        state_cache: Dict[str, Any],
+        choice_geometry_type: str,
+        ransac_method: str = DEFAULT_RANSAC_METHOD,
+        ransac_reproj_threshold: int = DEFAULT_RANSAC_REPROJ_THRESHOLD,
+        ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
+        ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
 ) -> Tuple[Optional[np.ndarray], Optional[Dict[str, int]]]:
     """
     Run RANSAC matches and return the output images and the number of matches.
@@ -710,7 +710,7 @@ def run_ransac(
         ransac_confidence=ransac_confidence,
         ransac_max_iter=ransac_max_iter,
     )
-    logger.info(f"RANSAC matches done using: {time.time()-t1:.3f}s")
+    logger.info(f"RANSAC matches done using: {time.time() - t1:.3f}s")
     t1 = time.time()
 
     # plot images with ransac matches
@@ -721,7 +721,7 @@ def run_ransac(
     output_matches_ransac, num_matches_ransac = display_matches(
         state_cache, titles=titles, tag="KPTS_RANSAC"
     )
-    logger.info(f"Display matches done using: {time.time()-t1:.3f}s")
+    logger.info(f"Display matches done using: {time.time() - t1:.3f}s")
     t1 = time.time()
 
     # compute warp images
@@ -753,24 +753,24 @@ def run_ransac(
         tmp_state_cache,
     )
 
-
+@spaces.GPU
 def run_matching(
-    image0: np.ndarray,
-    image1: np.ndarray,
-    match_threshold: float,
-    extract_max_keypoints: int,
-    keypoint_threshold: float,
-    key: str,
-    ransac_method: str = DEFAULT_RANSAC_METHOD,
-    ransac_reproj_threshold: int = DEFAULT_RANSAC_REPROJ_THRESHOLD,
-    ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
-    ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
-    choice_geometry_type: str = DEFAULT_SETTING_GEOMETRY,
-    matcher_zoo: Dict[str, Any] = None,
-    force_resize: bool = False,
-    image_width: int = 640,
-    image_height: int = 480,
-    use_cached_model: bool = False,
+        image0: np.ndarray,
+        image1: np.ndarray,
+        match_threshold: float,
+        extract_max_keypoints: int,
+        keypoint_threshold: float,
+        key: str,
+        ransac_method: str = DEFAULT_RANSAC_METHOD,
+        ransac_reproj_threshold: int = DEFAULT_RANSAC_REPROJ_THRESHOLD,
+        ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
+        ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
+        choice_geometry_type: str = DEFAULT_SETTING_GEOMETRY,
+        matcher_zoo: Dict[str, Any] = None,
+        force_resize: bool = False,
+        image_width: int = 640,
+        image_height: int = 480,
+        use_cached_model: bool = False,
 ) -> Tuple[
     np.ndarray,
     np.ndarray,
@@ -846,7 +846,7 @@ def run_matching(
     else:
         matcher = get_model(match_conf)
         print('match_conf2', match_conf)
-    logger.info(f"Loading model using: {time.time()-t0:.3f}s")
+    logger.info(f"Loading model using: {time.time() - t0:.3f}s")
     t1 = time.time()
 
     if model["dense"]:
@@ -899,13 +899,13 @@ def run_matching(
         )
         pred = match_features.match_images(matcher, pred0, pred1)
         # print('pred', pred)
-        mconf= pred["mconf"]
+        mconf = pred["mconf"]
         print('mconf', mconf.min(), mconf.max())
         del extractor
     # gr.Info(
     #     f"Matching images done using: {time.time()-t1:.3f}s",
     # )
-    logger.info(f"Matching images done using: {time.time()-t1:.3f}s")
+    logger.info(f"Matching images done using: {time.time() - t1:.3f}s")
     t1 = time.time()
 
     # plot images with keypoints
@@ -932,7 +932,7 @@ def run_matching(
     )
 
     # gr.Info(f"RANSAC matches done using: {time.time()-t1:.3f}s")
-    logger.info(f"RANSAC matches done using: {time.time()-t1:.3f}s")
+    logger.info(f"RANSAC matches done using: {time.time() - t1:.3f}s")
     t1 = time.time()
 
     # plot images with ransac matches
@@ -944,7 +944,7 @@ def run_matching(
         pred, titles=titles, tag="KPTS_RANSAC"
     )
     # gr.Info(f"Display matches done using: {time.time()-t1:.3f}s")
-    logger.info(f"Display matches done using: {time.time()-t1:.3f}s")
+    logger.info(f"Display matches done using: {time.time() - t1:.3f}s")
 
     t1 = time.time()
     # plot wrapped images
@@ -956,7 +956,7 @@ def run_matching(
     )
     plt.close("all")
     # gr.Info(f"In summary, total time: {time.time()-t0:.3f}s")
-    logger.info(f"TOTAL time: {time.time()-t0:.3f}s")
+    logger.info(f"TOTAL time: {time.time() - t0:.3f}s")
 
     state_cache = pred
     state_cache["num_matches_raw"] = num_matches_raw