third_party/ASpanFormer/demo/demo_utils.py

import cv2
import numpy as np


def resize(image, long_dim):
    h, w = image.shape[0], image.shape[1]
    image = cv2.resize(
        image, (int(w * long_dim / max(h, w)), int(h * long_dim / max(h, w)))
    )
    return image


def draw_points(img, points, color=(0, 255, 0), radius=3):
    dp = [(int(points[i, 0]), int(points[i, 1])) for i in range(points.shape[0])]
    for i in range(points.shape[0]):
        cv2.circle(img, dp[i], radius=radius, color=color)
    return img


def draw_match(
    img1,
    img2,
    corr1,
    corr2,
    inlier=[True],
    color=None,
    radius1=1,
    radius2=1,
    resize=None,
):
    if resize is not None:
        scale1, scale2 = [img1.shape[1] / resize[0], img1.shape[0] / resize[1]], [
            img2.shape[1] / resize[0],
            img2.shape[0] / resize[1],
        ]
        img1, img2 = cv2.resize(img1, resize, interpolation=cv2.INTER_AREA), cv2.resize(
            img2, resize, interpolation=cv2.INTER_AREA
        )
        corr1, corr2 = (
            corr1 / np.asarray(scale1)[np.newaxis],
            corr2 / np.asarray(scale2)[np.newaxis],
        )
    corr1_key = [
        cv2.KeyPoint(corr1[i, 0], corr1[i, 1], radius1) for i in range(corr1.shape[0])
    ]
    corr2_key = [
        cv2.KeyPoint(corr2[i, 0], corr2[i, 1], radius2) for i in range(corr2.shape[0])
    ]

    assert len(corr1) == len(corr2)

    draw_matches = [cv2.DMatch(i, i, 0) for i in range(len(corr1))]
    if color is None:
        color = [(0, 255, 0) if cur_inlier else (0, 0, 255) for cur_inlier in inlier]
    if len(color) == 1:
        display = cv2.drawMatches(
            img1,
            corr1_key,
            img2,
            corr2_key,
            draw_matches,
            None,
            matchColor=color[0],
            singlePointColor=color[0],
            flags=4,
        )
    else:
        height, width = max(img1.shape[0], img2.shape[0]), img1.shape[1] + img2.shape[1]
        display = np.zeros([height, width, 3], np.uint8)
        display[: img1.shape[0], : img1.shape[1]] = img1
        display[: img2.shape[0], img1.shape[1] :] = img2
        for i in range(len(corr1)):
            left_x, left_y, right_x, right_y = (
                int(corr1[i][0]),
                int(corr1[i][1]),
                int(corr2[i][0] + img1.shape[1]),
                int(corr2[i][1]),
            )
            cur_color = (int(color[i][0]), int(color[i][1]), int(color[i][2]))
            cv2.line(
                display,
                (left_x, left_y),
                (right_x, right_y),
                cur_color,
                1,
                lineType=cv2.LINE_AA,
            )
    return display