File size: 3,368 Bytes
a9f0f33 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
#############################################################################
#
# Source from:
# https://www.tensorflow.org/hub/tutorials/tf_hub_delf_module
# Forked from:
# https://www.tensorflow.org/hub/tutorials/tf_hub_delf_module
# Reimplemented by: Leonel Hernández
#
##############################################################################
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from PIL import Image, ImageOps
from huggingface_hub import snapshot_download
from scipy.spatial import cKDTree
from skimage.feature import plot_matches
from skimage.measure import ransac
from skimage.transform import AffineTransform
DELF_REPO_ID = "leonelhs/delf"
def match_images(image1, image2, result1, result2):
distance_threshold = 0.8
# Read features.
num_features_1 = result1['locations'].shape[0]
print("Loaded image 1's %d features" % num_features_1)
num_features_2 = result2['locations'].shape[0]
print("Loaded image 2's %d features" % num_features_2)
# Find nearest-neighbor matches using a KD tree.
d1_tree = cKDTree(result1['descriptors'])
_, indices = d1_tree.query(
result2['descriptors'],
distance_upper_bound=distance_threshold)
# Select feature locations for putative matches.
locations_2_to_use = np.array([
result2['locations'][i,]
for i in range(num_features_2)
if indices[i] != num_features_1
])
locations_1_to_use = np.array([
result1['locations'][indices[i],]
for i in range(num_features_2)
if indices[i] != num_features_1
])
# Perform geometric verification using RANSAC.
_, inliers = ransac(
(locations_1_to_use, locations_2_to_use),
AffineTransform,
min_samples=3,
residual_threshold=20,
max_trials=1000)
print('Found %d inliers' % sum(inliers))
# Visualize correspondences.
fig, ax = plt.subplots()
inlier_idxs = np.nonzero(inliers)[0]
stack = np.column_stack((inlier_idxs, inlier_idxs))
plot_matches(
ax,
image1,
image2,
locations_1_to_use,
locations_2_to_use,
stack,
matches_color='b')
ax.axis('off')
ax.set_title('DELF correspondences')
fig.canvas.draw()
image_array = np.array(fig.canvas.renderer.buffer_rgba())
image = Image.fromarray(image_array)
return image.convert("RGB")
def crop_image(image, width=256, height=256):
return ImageOps.fit(image, (width, height), Image.LANCZOS)
class DeepLocalFeatures:
def __init__(self):
model_path = snapshot_download(DELF_REPO_ID)
self.model = tf.saved_model.load(model_path).signatures['default']
def run_delf(self, image):
np_image = np.array(image)
float_image = tf.image.convert_image_dtype(np_image, tf.float32)
return self.model(
image=float_image,
score_threshold=tf.constant(100.0),
image_scales=tf.constant([0.25, 0.3536, 0.5, 0.7071, 1.0, 1.4142, 2.0]),
max_feature_num=tf.constant(1000))
def match(self, image_a, image_b):
image_a = crop_image(image_a)
image_b = crop_image(image_b)
result_a = self.run_delf(image_a)
result_b = self.run_delf(image_b)
return match_images(image_a, image_b, result_a, result_b)
|