# -*- coding: utf-8 -*-
# @Author : xuelun
import cv2
import torch
import argparse
import warnings
import numpy as np
import matplotlib.pyplot as plt
import torchvision.transforms.functional as F
from os.path import join
from dkm.models.model_zoo.DKMv3 import DKMv3
from gluefactory.superpoint import SuperPoint
from gluefactory.models.matchers.lightglue import LightGlue
DEFAULT_MIN_NUM_MATCHES = 4
DEFAULT_RANSAC_MAX_ITER = 10000
DEFAULT_RANSAC_CONFIDENCE = 0.999
DEFAULT_RANSAC_REPROJ_THRESHOLD = 8
DEFAULT_RANSAC_METHOD = "USAC_MAGSAC"
RANSAC_ZOO = {
"RANSAC": cv2.RANSAC,
"USAC_FAST": cv2.USAC_FAST,
"USAC_MAGSAC": cv2.USAC_MAGSAC,
"USAC_PROSAC": cv2.USAC_PROSAC,
"USAC_DEFAULT": cv2.USAC_DEFAULT,
"USAC_FM_8PTS": cv2.USAC_FM_8PTS,
"USAC_ACCURATE": cv2.USAC_ACCURATE,
"USAC_PARALLEL": cv2.USAC_PARALLEL,
}
def read_image(path, grayscale=False):
if grayscale:
mode = cv2.IMREAD_GRAYSCALE
else:
mode = cv2.IMREAD_COLOR
image = cv2.imread(str(path), mode)
if image is None:
raise ValueError(f'Cannot read image {path}.')
if not grayscale and len(image.shape) == 3:
image = image[:, :, ::-1] # BGR to RGB
return image
def resize_image(image, size, interp):
assert interp.startswith('cv2_')
if interp.startswith('cv2_'):
interp = getattr(cv2, 'INTER_'+interp[len('cv2_'):].upper())
h, w = image.shape[:2]
if interp == cv2.INTER_AREA and (w < size[0] or h < size[1]):
interp = cv2.INTER_LINEAR
resized = cv2.resize(image, size, interpolation=interp)
# elif interp.startswith('pil_'):
# interp = getattr(PIL.Image, interp[len('pil_'):].upper())
# resized = PIL.Image.fromarray(image.astype(np.uint8))
# resized = resized.resize(size, resample=interp)
# resized = np.asarray(resized, dtype=image.dtype)
else:
raise ValueError(
f'Unknown interpolation {interp}.')
return resized
def fast_make_matching_figure(data, b_id):
color0 = (data['color0'][b_id].permute(1, 2, 0).cpu().detach().numpy() * 255).round().astype(np.uint8) # (rH, rW, 3)
color1 = (data['color1'][b_id].permute(1, 2, 0).cpu().detach().numpy() * 255).round().astype(np.uint8) # (rH, rW, 3)
gray0 = cv2.cvtColor(color0, cv2.COLOR_RGB2GRAY)
gray1 = cv2.cvtColor(color1, cv2.COLOR_RGB2GRAY)
kpts0 = data['mkpts0_f'].cpu().detach().numpy()
kpts1 = data['mkpts1_f'].cpu().detach().numpy()
mconf = data['mconf'].cpu().detach().numpy()
inliers = data['inliers']
rows = 2
margin = 2
(h0, w0), (h1, w1) = data['hw0_i'], data['hw1_i']
h = max(h0, h1)
H, W = margin * (rows + 1) + h * rows, margin * 3 + w0 + w1
# canvas
out = 255 * np.ones((H, W), np.uint8)
wx = [margin, margin + w0, margin + w0 + margin, margin + w0 + margin + w1]
hx = lambda row: margin * row + h * (row-1)
out = np.stack([out] * 3, -1)
sh = hx(row=1)
out[sh: sh + h0, wx[0]: wx[1]] = color0
out[sh: sh + h1, wx[2]: wx[3]] = color1
sh = hx(row=2)
out[sh: sh + h0, wx[0]: wx[1]] = color0
out[sh: sh + h1, wx[2]: wx[3]] = color1
mkpts0, mkpts1 = np.round(kpts0).astype(int), np.round(kpts1).astype(int)
for (x0, y0), (x1, y1) in zip(mkpts0[inliers], mkpts1[inliers]):
c = (0, 255, 0)
cv2.circle(out, (x0, y0 + sh), 3, c, -1, lineType=cv2.LINE_AA)
cv2.circle(out, (x1 + margin + w0, y1 + sh), 3, c, -1, lineType=cv2.LINE_AA)
return out
def fast_make_matching_overlay(data, b_id):
color0 = (data['color0'][b_id].permute(1, 2, 0).cpu().detach().numpy() * 255).round().astype(np.uint8) # (rH, rW, 3)
color1 = (data['color1'][b_id].permute(1, 2, 0).cpu().detach().numpy() * 255).round().astype(np.uint8) # (rH, rW, 3)
gray0 = cv2.cvtColor(color0, cv2.COLOR_RGB2GRAY)
gray1 = cv2.cvtColor(color1, cv2.COLOR_RGB2GRAY)
kpts0 = data['mkpts0_f'].cpu().detach().numpy()
kpts1 = data['mkpts1_f'].cpu().detach().numpy()
mconf = data['mconf'].cpu().detach().numpy()
inliers = data['inliers']
rows = 2
margin = 2
(h0, w0), (h1, w1) = data['hw0_i'], data['hw1_i']
h = max(h0, h1)
H, W = margin * (rows + 1) + h * rows, margin * 3 + w0 + w1
# canvas
out = 255 * np.ones((H, W), np.uint8)
wx = [margin, margin + w0, margin + w0 + margin, margin + w0 + margin + w1]
hx = lambda row: margin * row + h * (row-1)
out = np.stack([out] * 3, -1)
sh = hx(row=1)
out[sh: sh + h0, wx[0]: wx[1]] = color0
out[sh: sh + h1, wx[2]: wx[3]] = color1
sh = hx(row=2)
out[sh: sh + h0, wx[0]: wx[1]] = color0
out[sh: sh + h1, wx[2]: wx[3]] = color1
mkpts0, mkpts1 = np.round(kpts0).astype(int), np.round(kpts1).astype(int)
for (x0, y0), (x1, y1) in zip(mkpts0[inliers], mkpts1[inliers]):
c = (0, 255, 0)
cv2.line(out, (x0, y0 + sh), (x1 + margin + w0, y1 + sh), color=c, thickness=1, lineType=cv2.LINE_AA)
cv2.circle(out, (x0, y0 + sh), 3, c, -1, lineType=cv2.LINE_AA)
cv2.circle(out, (x1 + margin + w0, y1 + sh), 3, c, -1, lineType=cv2.LINE_AA)
return out
def preprocess(image: np.ndarray, grayscale: bool = False, resize_max: int = None,
dfactor: int = 8):
image = image.astype(np.float32, copy=False)
size = image.shape[:2][::-1]
scale = np.array([1.0, 1.0])
if resize_max:
scale = resize_max / max(size)
if scale < 1.0:
size_new = tuple(int(round(x*scale)) for x in size)
image = resize_image(image, size_new, 'cv2_area')
scale = np.array(size) / np.array(size_new)
if grayscale:
assert image.ndim == 2, image.shape
image = image[None]
else:
image = image.transpose((2, 0, 1)) # HxWxC to CxHxW
image = torch.from_numpy(image / 255.0).float()
# assure that the size is divisible by dfactor
size_new = tuple(map(
lambda x: int(x // dfactor * dfactor),
image.shape[-2:]))
image = F.resize(image, size=size_new)
scale = np.array(size) / np.array(size_new)[::-1]
return image, scale
def compute_geom(data,
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 = data["mkpts0_f"].cpu().detach().numpy()
mkpts1 = data["mkpts1_f"].cpu().detach().numpy()
if len(mkpts0) < 2 * DEFAULT_MIN_NUM_MATCHES:
return {}
h1, w1 = data["hw0_i"]
geo_info = {}
F, inliers = cv2.findFundamentalMat(
mkpts0,
mkpts1,
method=RANSAC_ZOO[ransac_method],
ransacReprojThreshold=ransac_reproj_threshold,
confidence=ransac_confidence,
maxIters=ransac_max_iter,
)
if F is not None:
geo_info["Fundamental"] = F.tolist()
H, _ = cv2.findHomography(
mkpts1,
mkpts0,
method=RANSAC_ZOO[ransac_method],
ransacReprojThreshold=ransac_reproj_threshold,
confidence=ransac_confidence,
maxIters=ransac_max_iter,
)
if H is not None:
geo_info["Homography"] = H.tolist()
_, H1, H2 = cv2.stereoRectifyUncalibrated(
mkpts0.reshape(-1, 2),
mkpts1.reshape(-1, 2),
F,
imgSize=(w1, h1),
)
geo_info["H1"] = H1.tolist()
geo_info["H2"] = H2.tolist()
return geo_info
def wrap_images(img0, img1, geo_info, geom_type):
img0 = img0[0].permute((1, 2, 0)).cpu().detach().numpy()[..., ::-1]
img1 = img1[0].permute((1, 2, 0)).cpu().detach().numpy()[..., ::-1]
h1, w1, _ = img0.shape
h2, w2, _ = img1.shape
rectified_image0 = img0
rectified_image1 = None
H = np.array(geo_info["Homography"])
F = np.array(geo_info["Fundamental"])
title = []
if geom_type == "Homography":
rectified_image1 = cv2.warpPerspective(
img1, H, (img0.shape[1], img0.shape[0])
)
title = ["Image 0", "Image 1 - warped"]
elif geom_type == "Fundamental":
H1, H2 = np.array(geo_info["H1"]), np.array(geo_info["H2"])
rectified_image0 = cv2.warpPerspective(img0, H1, (w1, h1))
rectified_image1 = cv2.warpPerspective(img1, H2, (w2, h2))
title = ["Image 0 - warped", "Image 1 - warped"]
else:
print("Error: Unknown geometry type")
fig = plot_images(
[rectified_image0.squeeze(), rectified_image1.squeeze()],
title,
dpi=300,
)
img = fig2im(fig)
plt.close(fig)
return img
def plot_images(imgs, titles=None, cmaps="gray", dpi=100, size=5, pad=0.5):
"""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.
dpi:
size:
pad:
"""
n = len(imgs)
if not isinstance(cmaps, (list, tuple)):
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 fig2im(fig):
fig.canvas.draw()
w, h = fig.canvas.get_width_height()
buf_ndarray = np.frombuffer(fig.canvas.tostring_rgb(), dtype="u1")
im = buf_ndarray.reshape(h, w, 3)
return im
if __name__ == '__main__':
model_zoo = ['gim_dkm', 'gim_lightglue']
# model
parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, default='gim_dkm', choices=model_zoo)
args = parser.parse_args()
# device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# load model
ckpt = None
model = None
detector = None
if args.model == 'gim_dkm':
ckpt = 'gim_dkm_100h.ckpt'
model = DKMv3(weights=None, h=672, w=896)
elif args.model == 'gim_lightglue':
ckpt = 'gim_lightglue_100h.ckpt'
detector = SuperPoint({
'max_num_keypoints': 2048,
'force_num_keypoints': True,
'detection_threshold': 0.0,
'nms_radius': 3,
'trainable': False,
})
model = LightGlue({
'filter_threshold': 0.1,
'flash': False,
'checkpointed': True,
})
# weights path
checkpoints_path = join('weights', ckpt)
# load state dict
if args.model == 'gim_dkm':
state_dict = torch.load(checkpoints_path, map_location='cpu')
if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict']
for k in list(state_dict.keys()):
if k.startswith('model.'):
state_dict[k.replace('model.', '', 1)] = state_dict.pop(k)
if 'encoder.net.fc' in k:
state_dict.pop(k)
model.load_state_dict(state_dict)
elif args.model == 'gim_lightglue':
state_dict = torch.load(checkpoints_path, map_location='cpu')
if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict']
for k in list(state_dict.keys()):
if k.startswith('model.'):
state_dict.pop(k)
if k.startswith('superpoint.'):
state_dict[k.replace('superpoint.', '', 1)] = state_dict.pop(k)
detector.load_state_dict(state_dict)
state_dict = torch.load(checkpoints_path, map_location='cpu')
if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict']
for k in list(state_dict.keys()):
if k.startswith('superpoint.'):
state_dict.pop(k)
if k.startswith('model.'):
state_dict[k.replace('model.', '', 1)] = state_dict.pop(k)
model.load_state_dict(state_dict)
# eval mode
if detector is not None:
detector = detector.eval().to(device)
model = model.eval().to(device)
name0 = 'a1'
name1 = 'a2'
postfix = '.png'
image_dir = join('assets', 'demo')
img_path0 = join(image_dir, name0 + postfix)
img_path1 = join(image_dir, name1 + postfix)
image0 = read_image(img_path0)
image1 = read_image(img_path1)
image0, scale0 = preprocess(image0)
image1, scale1 = preprocess(image1)
image0 = image0.to(device)[None]
image1 = image1.to(device)[None]
data = dict(color0=image0, color1=image1, image0=image0, image1=image1)
if args.model == 'gim_dkm':
with warnings.catch_warnings():
warnings.simplefilter("ignore")
dense_matches, dense_certainty = model.match(image0, image1)
sparse_matches, mconf = model.sample(dense_matches, dense_certainty, 5000)
height0, width0 = image0.shape[-2:]
height1, width1 = image1.shape[-2:]
kpts0 = sparse_matches[:, :2]
kpts0 = torch.stack((
width0 * (kpts0[:, 0] + 1) / 2, height0 * (kpts0[:, 1] + 1) / 2), dim=-1,)
kpts1 = sparse_matches[:, 2:]
kpts1 = torch.stack((
width1 * (kpts1[:, 0] + 1) / 2, height1 * (kpts1[:, 1] + 1) / 2), dim=-1,)
b_ids = torch.where(mconf[None])[0]
elif args.model == 'gim_lightglue':
gray0 = read_image(img_path0, grayscale=True)
gray1 = read_image(img_path1, grayscale=True)
gray0 = preprocess(gray0, grayscale=True)[0]
gray1 = preprocess(gray1, grayscale=True)[0]
gray0 = gray0.to(device)[None]
gray1 = gray1.to(device)[None]
scale0 = torch.tensor(scale0).to(device)[None]
scale1 = torch.tensor(scale1).to(device)[None]
data.update(dict(gray0=gray0, gray1=gray1))
size0 = torch.tensor(data["gray0"].shape[-2:][::-1])[None]
size1 = torch.tensor(data["gray1"].shape[-2:][::-1])[None]
data.update(dict(size0=size0, size1=size1))
data.update(dict(scale0=scale0, scale1=scale1))
pred = {}
pred.update({k + '0': v for k, v in detector({
"image": data["gray0"],
"image_size": data["size0"],
}).items()})
pred.update({k + '1': v for k, v in detector({
"image": data["gray1"],
"image_size": data["size1"],
}).items()})
pred.update(model({**pred, **data,
**{'resize0': data['size0'], 'resize1': data['size1']}}))
kpts0 = torch.cat([kp * s for kp, s in zip(pred['keypoints0'], data['scale0'][:, None])])
kpts1 = torch.cat([kp * s for kp, s in zip(pred['keypoints1'], data['scale1'][:, None])])
m_bids = torch.nonzero(pred['keypoints0'].sum(dim=2) > -1)[:, 0]
matches = pred['matches']
bs = data['image0'].size(0)
kpts0 = torch.cat([kpts0[m_bids == b_id][matches[b_id][..., 0]] for b_id in range(bs)])
kpts1 = torch.cat([kpts1[m_bids == b_id][matches[b_id][..., 1]] for b_id in range(bs)])
b_ids = torch.cat([m_bids[m_bids == b_id][matches[b_id][..., 0]] for b_id in range(bs)])
mconf = torch.cat(pred['scores'])
# robust fitting
_, mask = cv2.findFundamentalMat(kpts0.cpu().detach().numpy(),
kpts1.cpu().detach().numpy(),
cv2.USAC_MAGSAC, ransacReprojThreshold=1.0,
confidence=0.999999, maxIters=10000)
mask = mask.ravel() > 0
data.update({
'hw0_i': image0.shape[-2:],
'hw1_i': image1.shape[-2:],
'mkpts0_f': kpts0,
'mkpts1_f': kpts1,
'm_bids': b_ids,
'mconf': mconf,
'inliers': mask,
})
# save visualization
alpha = 0.5
out = fast_make_matching_figure(data, b_id=0)
overlay = fast_make_matching_overlay(data, b_id=0)
out = cv2.addWeighted(out, 1 - alpha, overlay, alpha, 0)
cv2.imwrite(join(image_dir, f'{name0}_{name1}_{args.model}_match.png'), out[..., ::-1])
geom_info = compute_geom(data)
wrapped_images = wrap_images(image0, image1, geom_info,
"Homography")
cv2.imwrite(join(image_dir, f'{name0}_{name1}_{args.model}_warp.png'), wrapped_images)