add: doc strings and type

app.py

@@ -1,5 +1,7 @@
 import argparse
 from pathlib import Path
+import numpy as np
+from typing import Dict, Any, Optional, Tuple, List, Union
 import gradio as gr
 from common.utils import (
     matcher_zoo,
@@ -56,36 +58,59 @@ def ui_change_imagebox(choice):
         }
 
 
-def ui_reset_state(*args):
+def ui_reset_state(
+    *args: Any,
+) -> Tuple[
+    Optional[np.ndarray],
+    Optional[np.ndarray],
+    float,
+    int,
+    float,
+    str,
+    Dict[str, Any],
+    Dict[str, Any],
+    str,
+    Optional[np.ndarray],
+    Optional[np.ndarray],
+    Optional[np.ndarray],
+    Dict[str, Any],
+    Dict[str, Any],
+    Optional[np.ndarray],
+    Dict[str, Any],
+    str,
+    int,
+    float,
+    int,
+]:
     """
     Reset the state of the UI.
 
     Returns:
         tuple: A tuple containing the initial values for the UI state.
     """
-    key = list(matcher_zoo.keys())[0]  # Get the first key from matcher_zoo
+    key: str = list(matcher_zoo.keys())[0]  # Get the first key from matcher_zoo
     return (
-        None,  # image0
-        None,  # image1
-        DEFAULT_MATCHING_THRESHOLD,  # matching_threshold
-        DEFAULT_SETTING_MAX_FEATURES,  # max_features
-        DEFAULT_DEFAULT_KEYPOINT_THRESHOLD,  # keypoint_threshold
-        key,  # matcher
-        ui_change_imagebox("upload"),  # input image0
-        ui_change_imagebox("upload"),  # input image1
-        "upload",  # match_image_src
-        None,  # keypoints
-        None,  # raw matches
-        None,  # ransac matches
-        {},  # matches result info
-        {},  # matcher config
-        None,  # warped imageInstance of 'Radio' has no 'change' member
-        {},  # geometry result
-        DEFAULT_RANSAC_METHOD,  # ransac_method
-        DEFAULT_RANSAC_REPROJ_THRESHOLD,  # ransac_reproj_threshold
-        DEFAULT_RANSAC_CONFIDENCE,  # ransac_confidence
-        DEFAULT_RANSAC_MAX_ITER,  # ransac_max_iter
-        DEFAULT_SETTING_GEOMETRY,  # geometry
+        None,  # image0: Optional[np.ndarray]
+        None,  # image1: Optional[np.ndarray]
+        DEFAULT_MATCHING_THRESHOLD,  # matching_threshold: float
+        DEFAULT_SETTING_MAX_FEATURES,  # max_features: int
+        DEFAULT_DEFAULT_KEYPOINT_THRESHOLD,  # keypoint_threshold: float
+        key,  # matcher: str
+        ui_change_imagebox("upload"),  # input image0: Dict[str, Any]
+        ui_change_imagebox("upload"),  # input image1: Dict[str, Any]
+        "upload",  # match_image_src: str
+        None,  # keypoints: Optional[np.ndarray]
+        None,  # raw matches: Optional[np.ndarray]
+        None,  # ransac matches: Optional[np.ndarray]
+        {},  # matches result info: Dict[str, Any]
+        {},  # matcher config: Dict[str, Any]
+        None,  # warped image: Optional[np.ndarray]
+        {},  # geometry result: Dict[str, Any]
+        DEFAULT_RANSAC_METHOD,  # ransac_method: str
+        DEFAULT_RANSAC_REPROJ_THRESHOLD,  # ransac_reproj_threshold: float
+        DEFAULT_RANSAC_CONFIDENCE,  # ransac_confidence: float
+        DEFAULT_RANSAC_MAX_ITER,  # ransac_max_iter: int
+        DEFAULT_SETTING_GEOMETRY,  # geometry: str
     )
 
 

common/utils.py

@@ -5,6 +5,7 @@ import torch
 import cv2
 import gradio as gr
 from pathlib import Path
+from typing import Dict, Any, Optional, Tuple, List, Union
 from itertools import combinations
 from hloc import matchers, extractors, logger
 from hloc.utils.base_model import dynamic_load
@@ -25,19 +26,39 @@ DEFAULT_RANSAC_MAX_ITER = 10000
 DEFAULT_MIN_NUM_MATCHES = 4
 DEFAULT_MATCHING_THRESHOLD = 0.2
 DEFAULT_SETTING_GEOMETRY = "Homography"
-GRADIO_VERSION = gr.__version__.split('.')[0]
+GRADIO_VERSION = gr.__version__.split(".")[0]
 
-def get_model(match_conf):
+
+def get_model(match_conf: Dict[str, Any]):
+    """
+    Load a matcher model from the provided configuration.
+
+    Args:
+        match_conf: A dictionary containing the model configuration.
+
+    Returns:
+        A matcher model instance.
+    """
     Model = dynamic_load(matchers, match_conf["model"]["name"])
     model = Model(match_conf["model"]).eval().to(device)
     return model
 
 
-def get_feature_model(conf):
+def get_feature_model(conf: Dict[str, Dict[str, Any]]):
+    """
+    Load a feature extraction model from the provided configuration.
+
+    Args:
+        conf: A dictionary containing the model configuration.
+
+    Returns:
+        A feature extraction model instance.
+    """
     Model = dynamic_load(extractors, conf["model"]["name"])
     model = Model(conf["model"]).eval().to(device)
     return model
 
+
 def gen_examples():
     random.seed(1)
     example_matchers = [
@@ -92,15 +113,30 @@ def gen_examples():
 
 
 def filter_matches(
-    pred,
-    ransac_method=DEFAULT_RANSAC_METHOD,
-    ransac_reproj_threshold=DEFAULT_RANSAC_REPROJ_THRESHOLD,
-    ransac_confidence=DEFAULT_RANSAC_CONFIDENCE,
-    ransac_max_iter=DEFAULT_RANSAC_MAX_ITER,
-):
-    mkpts0 = None
-    mkpts1 = None
-    feature_type = 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,
+) -> Dict[str, Any]:
+    """
+    Filter matches using RANSAC. If keypoints are available, filter by keypoints.
+    If lines are available, filter by lines. If both keypoints and lines are
+    available, filter by keypoints.
+
+    Args:
+        pred (Dict[str, Any]): dict of matches, including original keypoints.
+        ransac_method (str, optional): RANSAC method. Defaults to DEFAULT_RANSAC_METHOD.
+        ransac_reproj_threshold (float, optional): RANSAC reprojection threshold. Defaults to DEFAULT_RANSAC_REPROJ_THRESHOLD.
+        ransac_confidence (float, optional): RANSAC confidence. Defaults to DEFAULT_RANSAC_CONFIDENCE.
+        ransac_max_iter (int, optional): RANSAC maximum iterations. Defaults to DEFAULT_RANSAC_MAX_ITER.
+
+    Returns:
+        Dict[str, Any]: filtered matches.
+    """
+    mkpts0: Optional[np.ndarray] = None
+    mkpts1: Optional[np.ndarray] = None
+    feature_type: Optional[str] = None
     if "keypoints0_orig" in pred.keys() and "keypoints1_orig" in pred.keys():
         mkpts0 = pred["keypoints0_orig"]
         mkpts1 = pred["keypoints1_orig"]
@@ -142,20 +178,33 @@ def filter_matches(
 
 
 def compute_geom(
-    pred,
-    ransac_method=DEFAULT_RANSAC_METHOD,
-    ransac_reproj_threshold=DEFAULT_RANSAC_REPROJ_THRESHOLD,
-    ransac_confidence=DEFAULT_RANSAC_CONFIDENCE,
-    ransac_max_iter=DEFAULT_RANSAC_MAX_ITER,
-) -> dict:
-    mkpts0 = None
-    mkpts1 = 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,
+) -> Dict[str, List[float]]:
+    """
+    Compute geometric information of matches, including Fundamental matrix,
+    Homography matrix, and rectification matrices (if available).
+
+    Args:
+        pred (Dict[str, Any]): dict of matches, including original keypoints.
+        ransac_method (str, optional): RANSAC method. Defaults to DEFAULT_RANSAC_METHOD.
+        ransac_reproj_threshold (float, optional): RANSAC reprojection threshold. Defaults to DEFAULT_RANSAC_REPROJ_THRESHOLD.
+        ransac_confidence (float, optional): RANSAC confidence. Defaults to DEFAULT_RANSAC_CONFIDENCE.
+        ransac_max_iter (int, optional): RANSAC maximum iterations. Defaults to DEFAULT_RANSAC_MAX_ITER.
+
+    Returns:
+        Dict[str, List[float]]: geometric information in form of a dict.
+    """
+    mkpts0: Optional[np.ndarray] = None
+    mkpts1: Optional[np.ndarray] = None
 
     if "keypoints0_orig" in pred.keys() and "keypoints1_orig" in pred.keys():
         mkpts0 = pred["keypoints0_orig"]
         mkpts1 = pred["keypoints1_orig"]
-
-    if (
+    elif (
         "line_keypoints0_orig" in pred.keys()
         and "line_keypoints1_orig" in pred.keys()
     ):
@@ -166,7 +215,7 @@ def compute_geom(
         if len(mkpts0) < 2 * DEFAULT_MIN_NUM_MATCHES:
             return {}
         h1, w1, _ = pred["image0_orig"].shape
-        geo_info = {}
+        geo_info: Dict[str, List[float]] = {}
         F, inliers = cv2.findFundamentalMat(
             mkpts0,
             mkpts1,
@@ -197,22 +246,39 @@ def compute_geom(
                 geo_info["H1"] = H1.tolist()
                 geo_info["H2"] = H2.tolist()
             except cv2.error as e:
-                logger.error(f"e, skip")
+                logger.error(f"{e}, skip")
         return geo_info
     else:
         return {}
 
 
-def wrap_images(img0, img1, geo_info, geom_type):
+def wrap_images(
+    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.
+
+    Args:
+        img0: numpy array representing the first image.
+        img1: numpy array representing the second image.
+        geo_info: dictionary containing the geometric transformation information.
+        geom_type: type of geometric transformation used to align the images.
+
+    Returns:
+        A tuple containing a base64 encoded image string and a dictionary with the transformation matrix.
+    """
     h1, w1, _ = img0.shape
     h2, w2, _ = img1.shape
-    result_matrix = None
+    result_matrix: Optional[np.ndarray] = None
     if geo_info is not None and len(geo_info) != 0:
         rectified_image0 = img0
         rectified_image1 = None
         H = np.array(geo_info["Homography"])
         F = np.array(geo_info["Fundamental"])
-        title = []
+        title: List[str] = []
         if geom_type == "Homography":
             rectified_image1 = cv2.warpPerspective(
                 img1, H, (img0.shape[1], img0.shape[0])
@@ -242,15 +308,32 @@ def wrap_images(img0, img1, geo_info, geom_type):
         return None, None
 
 
-def change_estimate_geom(input_image0, input_image1, matches_info, choice):
+def change_estimate_geom(
+    input_image0: np.ndarray,
+    input_image1: np.ndarray,
+    matches_info: Dict[str, Any],
+    choice: str,
+) -> Tuple[Optional[np.ndarray], Optional[Dict[str, Any]]]:
+    """
+    Changes the estimate of the geometric transformation used to align the images.
+
+    Args:
+        input_image0: First input image.
+        input_image1: Second input image.
+        matches_info: Dictionary containing information about the matches.
+        choice: Type of geometric transformation to use ('Homography' or 'Fundamental') or 'No' to disable.
+
+    Returns:
+        A tuple containing the updated images and the updated matches info.
+    """
     if (
         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"]
-    wrapped_images = None
+    geom_info: Dict[str, Any] = matches_info["geom_info"]
+    wrapped_images: Optional[np.ndarray] = None
     if choice != "No":
         wrapped_images, _ = wrap_images(
             input_image0, input_image1, geom_info, choice
@@ -260,16 +343,34 @@ def change_estimate_geom(input_image0, input_image1, matches_info, choice):
         return None, None
 
 
-def display_matches(pred: dict, titles=[], dpi=300):
+def display_matches(
+    pred: Dict[str, np.ndarray], titles: List[str] = [], dpi: int = 300
+) -> Tuple[np.ndarray, int]:
+    """
+    Displays the matches between two images.
+
+    Args:
+        pred: Dictionary containing the original images and the matches.
+        titles: Optional titles for the plot.
+        dpi: Resolution of the plot.
+
+    Returns:
+        The resulting concatenated plot and the number of inliers.
+    """
     img0 = pred["image0_orig"]
     img1 = pred["image1_orig"]
 
     num_inliers = 0
-    if "keypoints0_orig" in pred.keys() and "keypoints1_orig" in pred.keys():
+    if (
+        "keypoints0_orig" in pred
+        and "keypoints1_orig" in pred
+        and pred["keypoints0_orig"] is not None
+        and pred["keypoints1_orig"] is not None
+    ):
         mkpts0 = pred["keypoints0_orig"]
         mkpts1 = pred["keypoints1_orig"]
         num_inliers = len(mkpts0)
-        if "mconf" in pred.keys():
+        if "mconf" in pred:
             mconf = pred["mconf"]
         else:
             mconf = np.ones(len(mkpts0))
@@ -283,7 +384,12 @@ def display_matches(pred: dict, titles=[], dpi=300):
             titles=titles,
         )
         fig = fig_mkpts
-    if "line0_orig" in pred.keys() and "line1_orig" in pred.keys():
+    if (
+        "line0_orig" in pred
+        and "line1_orig" in pred
+        and pred["line0_orig"] is not None
+        and pred["line1_orig"] is not None
+    ):
         # lines
         mtlines0 = pred["line0_orig"]
         mtlines1 = pred["line1_orig"]
@@ -297,12 +403,12 @@ def display_matches(pred: dict, titles=[], dpi=300):
         fig_lines = fig2im(fig_lines)
 
         # keypoints
-        mkpts0 = pred["line_keypoints0_orig"]
-        mkpts1 = pred["line_keypoints1_orig"]
+        mkpts0 = pred.get("line_keypoints0_orig")
+        mkpts1 = pred.get("line_keypoints1_orig")
 
         if mkpts0 is not None and mkpts1 is not None:
             num_inliers = len(mkpts0)
-            if "mconf" in pred.keys():
+            if "mconf" in pred:
                 mconf = pred["mconf"]
             else:
                 mconf = np.ones(len(mkpts0))
@@ -317,18 +423,51 @@ def display_matches(pred: dict, titles=[], dpi=300):
 
 
 def run_matching(
-    image0,
-    image1,
-    match_threshold,
-    extract_max_keypoints,
-    keypoint_threshold,
-    key,
-    ransac_method=DEFAULT_RANSAC_METHOD,
-    ransac_reproj_threshold=DEFAULT_RANSAC_REPROJ_THRESHOLD,
-    ransac_confidence=DEFAULT_RANSAC_CONFIDENCE,
-    ransac_max_iter=DEFAULT_RANSAC_MAX_ITER,
-    choice_estimate_geom=DEFAULT_SETTING_GEOMETRY,
-):
+    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_estimate_geom: str = DEFAULT_SETTING_GEOMETRY,
+) -> Tuple[
+    np.ndarray,
+    np.ndarray,
+    np.ndarray,
+    Dict[str, int],
+    Dict[str, Dict[str, Any]],
+    Dict[str, Dict[str, float]],
+    np.ndarray,
+]:
+    """Match two images using the given parameters.
+
+    Args:
+        image0 (np.ndarray): RGB image 0.
+        image1 (np.ndarray): RGB image 1.
+        match_threshold (float): match threshold.
+        extract_max_keypoints (int): number of keypoints to extract.
+        keypoint_threshold (float): keypoint threshold.
+        key (str): key of the model to use.
+        ransac_method (str, optional): RANSAC method to use.
+        ransac_reproj_threshold (int, optional): RANSAC reprojection threshold.
+        ransac_confidence (float, optional): RANSAC confidence level.
+        ransac_max_iter (int, optional): RANSAC maximum number of iterations.
+        choice_estimate_geom (str, optional): setting of geometry estimation.
+
+    Returns:
+        tuple:
+            - output_keypoints (np.ndarray): image with keypoints.
+            - output_matches_raw (np.ndarray): image with raw matches.
+            - output_matches_ransac (np.ndarray): image with RANSAC matches.
+            - num_matches (Dict[str, int]): number of raw and RANSAC matches.
+            - configs (Dict[str, Dict[str, Any]]): match and feature extraction configs.
+            - geom_info (Dict[str, Dict[str, float]]): geometry information.
+            - output_wrapped (np.ndarray): wrapped images.
+    """
     # image0 and image1 is RGB mode
     if image0 is None or image1 is None:
         raise gr.Error("Error: No images found! Please upload two images.")

common/viz.py

@@ -1,20 +1,35 @@
 import numpy as np
-import matplotlib.pyplot as plt
-import matplotlib
 import seaborn as sns
+import matplotlib
+import matplotlib.pyplot as plt
+from pathlib import Path
+from typing import Dict, Any, Optional, Tuple, List, Union
 
 
-def plot_images(imgs, titles=None, cmaps="gray", dpi=100, size=5, pad=0.5):
+def plot_images(
+    imgs: List[np.ndarray],
+    titles: Optional[List[str]] = None,
+    cmaps: Union[str, List[str]] = "gray",
+    dpi: int = 100,
+    size: Optional[int] = 5,
+    pad: float = 0.5,
+) -> plt.Figure:
     """Plot a set of images horizontally.
     Args:
         imgs: a list of NumPy or PyTorch images, RGB (H, W, 3) or mono (H, W).
         titles: a list of strings, as titles for each image.
-        cmaps: colormaps for monochrome images.
+        cmaps: colormaps for monochrome images. If a single string is given,
+            it is used for all images.
+        dpi: DPI of the figure.
+        size: figure size in inches (width). If not provided, the figure
+            size is determined automatically.
+        pad: padding between subplots, in inches.
+    Returns:
+        The created figure.
     """
     n = len(imgs)
-    if not isinstance(cmaps, (list, tuple)):
+    if not isinstance(cmaps, list):
         cmaps = [cmaps] * n
-    # figsize = (size*n, size*3/4) if size is not None else None
     figsize = (size * n, size * 6 / 5) if size is not None else None
     fig, ax = plt.subplots(1, n, figsize=figsize, dpi=dpi)
 
@@ -33,24 +48,33 @@ def plot_images(imgs, titles=None, cmaps="gray", dpi=100, size=5, pad=0.5):
     return fig
 
 
-def plot_color_line_matches(lines, correct_matches=None, lw=2, indices=(0, 1)):
+def plot_color_line_matches(
+    lines: List[np.ndarray],
+    correct_matches: Optional[np.ndarray] = None,
+    lw: float = 2.0,
+    indices: Tuple[int, int] = (0, 1),
+) -> matplotlib.figure.Figure:
     """Plot line matches for existing images with multiple colors.
+
     Args:
-        lines: list of ndarrays of size (N, 2, 2).
-        correct_matches: bool array of size (N,) indicating correct matches.
-        lw: line width as float pixels.
-        indices: indices of the images to draw the matches on.
+        lines: List of ndarrays of size (N, 2, 2) representing line segments.
+        correct_matches: Optional bool array of size (N,) indicating correct
+            matches. If not None, display wrong matches with a low alpha.
+        lw: Line width as float pixels.
+        indices: Indices of the images to draw the matches on.
+
+    Returns:
+        The modified matplotlib figure.
     """
-    n_lines = len(lines[0])
+    n_lines = lines[0].shape[0]
     colors = sns.color_palette("husl", n_colors=n_lines)
     np.random.shuffle(colors)
     alphas = np.ones(n_lines)
-    # If correct_matches is not None, display wrong matches with a low alpha
     if correct_matches is not None:
         alphas[~np.array(correct_matches)] = 0.2
 
     fig = plt.gcf()
-    ax = fig.axes
+    ax = typing.cast(List[matplotlib.axes.Axes], fig.axes)
     assert len(ax) > max(indices)
     axes = [ax[i] for i in indices]
     fig.canvas.draw()
@@ -78,21 +102,39 @@ def plot_color_line_matches(lines, correct_matches=None, lw=2, indices=(0, 1)):
 
 
 def make_matching_figure(
-    img0,
-    img1,
-    mkpts0,
-    mkpts1,
-    color,
-    titles=None,
-    kpts0=None,
-    kpts1=None,
-    text=[],
-    dpi=75,
-    path=None,
-    pad=0,
-):
+    img0: np.ndarray,
+    img1: np.ndarray,
+    mkpts0: np.ndarray,
+    mkpts1: np.ndarray,
+    color: np.ndarray,
+    titles: Optional[List[str]] = None,
+    kpts0: Optional[np.ndarray] = None,
+    kpts1: Optional[np.ndarray] = None,
+    text: List[str] = [],
+    dpi: int = 75,
+    path: Optional[Path] = None,
+    pad: float = 0.0,
+) -> Optional[plt.Figure]:
+    """Draw image pair with matches.
+
+    Args:
+        img0: image0 as HxWx3 numpy array.
+        img1: image1 as HxWx3 numpy array.
+        mkpts0: matched points in image0 as Nx2 numpy array.
+        mkpts1: matched points in image1 as Nx2 numpy array.
+        color: colors for the matches as Nx4 numpy array.
+        titles: titles for the two subplots.
+        kpts0: keypoints in image0 as Kx2 numpy array.
+        kpts1: keypoints in image1 as Kx2 numpy array.
+        text: list of strings to display in the top-left corner of the image.
+        dpi: dots per inch of the saved figure.
+        path: if not None, save the figure to this path.
+        pad: padding around the image as a fraction of the image size.
+
+    Returns:
+        The matplotlib Figure object if path is None.
+    """
     # draw image pair
-    # assert mkpts0.shape[0] == mkpts1.shape[0], f'mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}'
     fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi)
     axes[0].imshow(img0)  # , cmap='gray')
     axes[1].imshow(img1)  # , cmap='gray')
@@ -156,7 +198,20 @@ def make_matching_figure(
         return fig
 
 
-def error_colormap(err, thr, alpha=1.0):
+def error_colormap(
+    err: np.ndarray, thr: float, alpha: float = 1.0
+) -> np.ndarray:
+    """
+    Create a colormap based on the error values.
+
+    Args:
+        err: Error values as a numpy array of shape (N,).
+        thr: Threshold value for the error.
+        alpha: Alpha value for the colormap, between 0 and 1.
+
+    Returns:
+        Colormap as a numpy array of shape (N, 4) with values in [0, 1].
+    """
     assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
     x = 1 - np.clip(err / (thr * 2), 0, 1)
     return np.clip(
@@ -173,22 +228,57 @@ color_map = np.arange(100)
 np.random.shuffle(color_map)
 
 
-def fig2im(fig):
+def fig2im(fig: matplotlib.figure.Figure) -> np.ndarray:
+    """
+    Convert a matplotlib figure to a numpy array with RGB values.
+
+    Args:
+        fig: A matplotlib figure.
+
+    Returns:
+        A numpy array with shape (height, width, 3) and dtype uint8 containing
+        the RGB values of the figure.
+    """
     fig.canvas.draw()
-    w, h = fig.canvas.get_width_height()
+    (width, height) = fig.canvas.get_width_height()
     buf_ndarray = np.frombuffer(fig.canvas.tostring_rgb(), dtype="u1")
-    im = buf_ndarray.reshape(h, w, 3)
-    return im
+    return buf_ndarray.reshape(height, width, 3)
 
 
 def draw_matches(
-    mkpts0, mkpts1, img0, img1, conf, titles=None, dpi=150, path=None, pad=0.5
-):
+    mkpts0: List[np.ndarray],
+    mkpts1: List[np.ndarray],
+    img0: np.ndarray,
+    img1: np.ndarray,
+    conf: np.ndarray,
+    titles: Optional[List[str]] = None,
+    dpi: int = 150,
+    path: Optional[str] = None,
+    pad: float = 0.5,
+) -> np.ndarray:
+    """
+    Draw matches between two images.
+
+    Args:
+        mkpts0: List of matches from the first image, with shape (N, 2)
+        mkpts1: List of matches from the second image, with shape (N, 2)
+        img0: First image, with shape (H, W, 3)
+        img1: Second image, with shape (H, W, 3)
+        conf: Confidence values for the matches, with shape (N,)
+        titles: Optional list of title strings for the plot
+        dpi: DPI for the saved image
+        path: Optional path to save the image to. If None, the image is not saved.
+        pad: Padding between subplots
+
+    Returns:
+        The figure as a numpy array with shape (height, width, 3) and dtype uint8
+        containing the RGB values of the figure.
+    """
     thr = 5e-4
     thr = 0.5
     color = error_colormap(conf, thr, alpha=0.1)
     text = [
-        f"image name",
+        "image name",
         f"#Matches: {len(mkpts0)}",
     ]
     if path:
@@ -222,7 +312,31 @@ def draw_matches(
         )
 
 
-def draw_image_pairs(img0, img1, text=[], dpi=75, path=None, pad=0.5):
+def draw_image_pairs(
+    img0: np.ndarray,
+    img1: np.ndarray,
+    text: List[str] = [],
+    dpi: int = 75,
+    path: Optional[str] = None,
+    pad: float = 0.5,
+) -> np.ndarray:
+    """Draw image pair horizontally.
+
+    Args:
+        img0: First image, with shape (H, W, 3)
+        img1: Second image, with shape (H, W, 3)
+        text: List of strings to print. Each string is a new line.
+        dpi: DPI of the figure.
+        path: Path to save the image to. If None, the image is not saved and
+            the function returns the figure as a numpy array with shape
+            (height, width, 3) and dtype uint8 containing the RGB values of the
+            figure.
+        pad: Padding between subplots
+
+    Returns:
+        The figure as a numpy array with shape (height, width, 3) and dtype uint8
+        containing the RGB values of the figure, or None if path is not None.
+    """
     # draw image pair
     fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi)
     axes[0].imshow(img0)  # , cmap='gray')