src/superpoint_pruning/distillation/utils.py

import numpy as np
import matplotlib.pyplot as plt
import cv2


def plot_images(images, titles=None, figsize=None):
    n = len(images)
    if figsize is None:
        figsize = (5 * n, 5)

    fig, axes = plt.subplots(1, n, figsize=figsize)
    axes = np.atleast_1d(axes)

    for i, (ax, image) in enumerate(zip(axes, images)):
        if image.ndim == 2:
            ax.imshow(image, cmap="gray")
        else:
            ax.imshow(image)
        if titles is not None:
            ax.set_title(titles[i])
        ax.axis("off")

    fig.tight_layout()
    return fig, axes


def plot_images_with_keypoints(
    images,
    keypoints,
    titles=None,
    keypoint_size=30,
    figsize=None,
):
    fig, axes = plot_images(
        images,
        titles=titles,
        figsize=figsize,
    )
    for ax, points in zip(axes, keypoints):
        points = np.asarray(points)
        if len(points) == 0:
            continue
        x = points[:, 0]
        y = points[:, 1]

        ax.scatter(x, y, s=keypoint_size, marker=".", color="lime")
    fig.tight_layout()


def load_grayscale_image(path):
    image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)

    if image is None:
        raise FileNotFoundError(f"Could not read image: {path}")

    image = image.astype(np.float32) / 255.0

    return image


def rescale_image(image, new_size=None):
    height, width = image.shape[:2]
    new_width, new_height = new_size

    resized = cv2.resize(
        image,
        (new_width, new_height),
        interpolation=cv2.INTER_LINEAR,
    )

    scale_x = new_width / width
    scale_y = new_height / height

    return resized, (scale_x, scale_y)