Spaces:
Running
Running
| import cv2 | |
| import typing | |
| import matplotlib | |
| import numpy as np | |
| import seaborn as sns | |
| import matplotlib.pyplot as plt | |
| from pathlib import Path | |
| from typing import Dict, Any, Optional, Tuple, List, Union | |
| from hloc.utils.viz import add_text, plot_keypoints | |
| 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. 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): | |
| cmaps = [cmaps] * n | |
| figsize = (size * n, size * 6 / 5) if size is not None else None | |
| fig, ax = plt.subplots(1, n, figsize=figsize, dpi=dpi) | |
| if n == 1: | |
| ax = [ax] | |
| for i in range(n): | |
| ax[i].imshow(imgs[i], cmap=plt.get_cmap(cmaps[i])) | |
| ax[i].get_yaxis().set_ticks([]) | |
| ax[i].get_xaxis().set_ticks([]) | |
| ax[i].set_axis_off() | |
| for spine in ax[i].spines.values(): # remove frame | |
| spine.set_visible(False) | |
| if titles: | |
| ax[i].set_title(titles[i]) | |
| fig.tight_layout(pad=pad) | |
| return fig | |
| 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) 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 = 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: | |
| alphas[~np.array(correct_matches)] = 0.2 | |
| fig = plt.gcf() | |
| ax = typing.cast(List[matplotlib.axes.Axes], fig.axes) | |
| assert len(ax) > max(indices) | |
| axes = [ax[i] for i in indices] | |
| fig.canvas.draw() | |
| # Plot the lines | |
| for a, l in zip(axes, lines): | |
| # Transform the points into the figure coordinate system | |
| transFigure = fig.transFigure.inverted() | |
| endpoint0 = transFigure.transform(a.transData.transform(l[:, 0])) | |
| endpoint1 = transFigure.transform(a.transData.transform(l[:, 1])) | |
| fig.lines += [ | |
| matplotlib.lines.Line2D( | |
| (endpoint0[i, 0], endpoint1[i, 0]), | |
| (endpoint0[i, 1], endpoint1[i, 1]), | |
| zorder=1, | |
| transform=fig.transFigure, | |
| c=colors[i], | |
| alpha=alphas[i], | |
| linewidth=lw, | |
| ) | |
| for i in range(n_lines) | |
| ] | |
| return fig | |
| def make_matching_figure( | |
| 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 | |
| fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi) | |
| axes[0].imshow(img0) # , cmap='gray') | |
| axes[1].imshow(img1) # , cmap='gray') | |
| for i in range(2): # clear all frames | |
| axes[i].get_yaxis().set_ticks([]) | |
| axes[i].get_xaxis().set_ticks([]) | |
| for spine in axes[i].spines.values(): | |
| spine.set_visible(False) | |
| if titles is not None: | |
| axes[i].set_title(titles[i]) | |
| plt.tight_layout(pad=pad) | |
| if kpts0 is not None: | |
| assert kpts1 is not None | |
| axes[0].scatter(kpts0[:, 0], kpts0[:, 1], c="w", s=5) | |
| axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c="w", s=5) | |
| # draw matches | |
| if ( | |
| mkpts0.shape[0] != 0 | |
| and mkpts1.shape[0] != 0 | |
| and mkpts0.shape == mkpts1.shape | |
| ): | |
| fig.canvas.draw() | |
| transFigure = fig.transFigure.inverted() | |
| fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0)) | |
| fkpts1 = transFigure.transform(axes[1].transData.transform(mkpts1)) | |
| fig.lines = [ | |
| matplotlib.lines.Line2D( | |
| (fkpts0[i, 0], fkpts1[i, 0]), | |
| (fkpts0[i, 1], fkpts1[i, 1]), | |
| transform=fig.transFigure, | |
| c=color[i], | |
| linewidth=2, | |
| ) | |
| for i in range(len(mkpts0)) | |
| ] | |
| # freeze the axes to prevent the transform to change | |
| axes[0].autoscale(enable=False) | |
| axes[1].autoscale(enable=False) | |
| axes[0].scatter(mkpts0[:, 0], mkpts0[:, 1], c=color[..., :3], s=4) | |
| axes[1].scatter(mkpts1[:, 0], mkpts1[:, 1], c=color[..., :3], s=4) | |
| # put txts | |
| txt_color = "k" if img0[:100, :200].mean() > 200 else "w" | |
| fig.text( | |
| 0.01, | |
| 0.99, | |
| "\n".join(text), | |
| transform=fig.axes[0].transAxes, | |
| fontsize=15, | |
| va="top", | |
| ha="left", | |
| color=txt_color, | |
| ) | |
| # save or return figure | |
| if path: | |
| plt.savefig(str(path), bbox_inches="tight", pad_inches=0) | |
| plt.close() | |
| else: | |
| return fig | |
| 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( | |
| np.stack( | |
| [2 - x * 2, x * 2, np.zeros_like(x), np.ones_like(x) * alpha], -1 | |
| ), | |
| 0, | |
| 1, | |
| ) | |
| np.random.seed(1995) | |
| color_map = np.arange(100) | |
| np.random.shuffle(color_map) | |
| 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() | |
| (width, height) = fig.canvas.get_width_height() | |
| buf_ndarray = np.frombuffer(fig.canvas.tostring_rgb(), dtype="u1") | |
| return buf_ndarray.reshape(height, width, 3) | |
| def draw_matches_core( | |
| mkpts0: List[np.ndarray], | |
| mkpts1: List[np.ndarray], | |
| img0: np.ndarray, | |
| img1: np.ndarray, | |
| conf: np.ndarray, | |
| titles: Optional[List[str]] = None, | |
| texts: 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 = [ | |
| # "image name", | |
| f"#Matches: {len(mkpts0)}", | |
| ] | |
| if path: | |
| fig2im( | |
| make_matching_figure( | |
| img0, | |
| img1, | |
| mkpts0, | |
| mkpts1, | |
| color, | |
| titles=titles, | |
| text=text, | |
| path=path, | |
| dpi=dpi, | |
| pad=pad, | |
| ) | |
| ) | |
| else: | |
| return fig2im( | |
| make_matching_figure( | |
| img0, | |
| img1, | |
| mkpts0, | |
| mkpts1, | |
| color, | |
| titles=titles, | |
| text=text, | |
| pad=pad, | |
| dpi=dpi, | |
| ) | |
| ) | |
| 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') | |
| axes[1].imshow(img1) # , cmap='gray') | |
| for i in range(2): # clear all frames | |
| axes[i].get_yaxis().set_ticks([]) | |
| axes[i].get_xaxis().set_ticks([]) | |
| for spine in axes[i].spines.values(): | |
| spine.set_visible(False) | |
| plt.tight_layout(pad=pad) | |
| # put txts | |
| txt_color = "k" if img0[:100, :200].mean() > 200 else "w" | |
| fig.text( | |
| 0.01, | |
| 0.99, | |
| "\n".join(text), | |
| transform=fig.axes[0].transAxes, | |
| fontsize=15, | |
| va="top", | |
| ha="left", | |
| color=txt_color, | |
| ) | |
| # save or return figure | |
| if path: | |
| plt.savefig(str(path), bbox_inches="tight", pad_inches=0) | |
| plt.close() | |
| else: | |
| return fig2im(fig) | |
| def display_keypoints(pred: dict, titles: List[str] = []): | |
| img0 = pred["image0_orig"] | |
| img1 = pred["image1_orig"] | |
| output_keypoints = plot_images([img0, img1], titles=titles, dpi=300) | |
| if "keypoints0_orig" in pred.keys() and "keypoints1_orig" in pred.keys(): | |
| plot_keypoints([pred["keypoints0_orig"], pred["keypoints1_orig"]]) | |
| text = ( | |
| f"# keypoints0: {len(pred['keypoints0_orig'])} \n" | |
| + f"# keypoints1: {len(pred['keypoints1_orig'])}" | |
| ) | |
| add_text(0, text, fs=15) | |
| output_keypoints = fig2im(output_keypoints) | |
| return output_keypoints | |
| def display_matches( | |
| pred: Dict[str, np.ndarray], | |
| titles: List[str] = [], | |
| texts: List[str] = [], | |
| dpi: int = 300, | |
| tag: str = "KPTS_RAW", # KPTS_RAW, KPTS_RANSAC, LINES_RAW, LINES_RANSAC, | |
| ) -> 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 | |
| KPTS0_KEY = None | |
| KPTS1_KEY = None | |
| confid = None | |
| if tag == "KPTS_RAW": | |
| KPTS0_KEY = "mkeypoints0_orig" | |
| KPTS1_KEY = "mkeypoints1_orig" | |
| if "mconf" in pred: | |
| confid = pred["mconf"] | |
| elif tag == "KPTS_RANSAC": | |
| KPTS0_KEY = "mmkeypoints0_orig" | |
| KPTS1_KEY = "mmkeypoints1_orig" | |
| if "mmconf" in pred: | |
| confid = pred["mmconf"] | |
| else: | |
| # TODO: LINES_RAW, LINES_RANSAC | |
| raise ValueError(f"Unknown tag: {tag}") | |
| # draw raw matches | |
| if ( | |
| KPTS0_KEY in pred | |
| and KPTS1_KEY in pred | |
| and pred[KPTS0_KEY] is not None | |
| and pred[KPTS1_KEY] is not None | |
| ): # draw ransac matches | |
| mkpts0 = pred[KPTS0_KEY] | |
| mkpts1 = pred[KPTS1_KEY] | |
| num_inliers = len(mkpts0) | |
| if confid is None: | |
| confid = np.ones(len(mkpts0)) | |
| fig_mkpts = draw_matches_core( | |
| mkpts0, | |
| mkpts1, | |
| img0, | |
| img1, | |
| confid, | |
| dpi=dpi, | |
| titles=titles, | |
| texts=texts, | |
| ) | |
| fig = fig_mkpts | |
| # TODO: draw lines | |
| if ( | |
| "line0_orig" in pred | |
| and "line1_orig" in pred | |
| and pred["line0_orig"] is not None | |
| and pred["line1_orig"] is not None | |
| and (tag == "LINES_RAW" or tag == "LINES_RANSAC") | |
| ): | |
| # lines | |
| mtlines0 = pred["line0_orig"] | |
| mtlines1 = pred["line1_orig"] | |
| num_inliers = len(mtlines0) | |
| fig_lines = plot_images( | |
| [img0.squeeze(), img1.squeeze()], | |
| ["Image 0 - matched lines", "Image 1 - matched lines"], | |
| dpi=300, | |
| ) | |
| fig_lines = plot_color_line_matches([mtlines0, mtlines1], lw=2) | |
| fig_lines = fig2im(fig_lines) | |
| # keypoints | |
| mkpts0 = pred.get("line_keypoints0_orig") | |
| mkpts1 = pred.get("line_keypoints1_orig") | |
| fig = None | |
| breakpoint() | |
| if mkpts0 is not None and mkpts1 is not None: | |
| num_inliers = len(mkpts0) | |
| if "mconf" in pred: | |
| mconf = pred["mconf"] | |
| else: | |
| mconf = np.ones(len(mkpts0)) | |
| fig_mkpts = draw_matches_core( | |
| mkpts0, mkpts1, img0, img1, mconf, dpi=300 | |
| ) | |
| fig_lines = cv2.resize( | |
| fig_lines, (fig_mkpts.shape[1], fig_mkpts.shape[0]) | |
| ) | |
| fig = np.concatenate([fig_mkpts, fig_lines], axis=0) | |
| else: | |
| fig = fig_lines | |
| return fig, num_inliers | |