|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
"""Wrapper for performing OmniGlue inference, plus (optionally) SP/DINO.""" |
|
|
|
import numpy as np |
|
from omniglue import dino_extract |
|
from omniglue import superpoint_extract |
|
from omniglue import utils |
|
import tensorflow as tf |
|
|
|
DINO_FEATURE_DIM = 768 |
|
MATCH_THRESHOLD = 1e-3 |
|
|
|
|
|
class OmniGlue: |
|
|
|
|
|
def __init__( |
|
self, |
|
og_export: str, |
|
sp_export: str | None = None, |
|
dino_export: str | None = None, |
|
) -> None: |
|
self.matcher = tf.saved_model.load(og_export) |
|
if sp_export is not None: |
|
self.sp_extract = superpoint_extract.SuperPointExtract(sp_export) |
|
if dino_export is not None: |
|
self.dino_extract = dino_extract.DINOExtract(dino_export, feature_layer=1) |
|
|
|
def FindMatches(self, image0: np.ndarray, image1: np.ndarray): |
|
"""TODO(omniglue): docstring.""" |
|
height0, width0 = image0.shape[:2] |
|
height1, width1 = image1.shape[:2] |
|
sp_features0 = self.sp_extract(image0) |
|
sp_features1 = self.sp_extract(image1) |
|
dino_features0 = self.dino_extract(image0) |
|
dino_features1 = self.dino_extract(image1) |
|
dino_descriptors0 = dino_extract.get_dino_descriptors( |
|
dino_features0, |
|
tf.convert_to_tensor(sp_features0[0], dtype=tf.float32), |
|
tf.convert_to_tensor(height0, dtype=tf.int32), |
|
tf.convert_to_tensor(width0, dtype=tf.int32), |
|
DINO_FEATURE_DIM, |
|
) |
|
dino_descriptors1 = dino_extract.get_dino_descriptors( |
|
dino_features1, |
|
tf.convert_to_tensor(sp_features1[0], dtype=tf.float32), |
|
tf.convert_to_tensor(height1, dtype=tf.int32), |
|
tf.convert_to_tensor(width1, dtype=tf.int32), |
|
DINO_FEATURE_DIM, |
|
) |
|
|
|
inputs = self._construct_inputs( |
|
width0, |
|
height0, |
|
width1, |
|
height1, |
|
sp_features0, |
|
sp_features1, |
|
dino_descriptors0, |
|
dino_descriptors1, |
|
) |
|
|
|
og_outputs = self.matcher.signatures['serving_default'](**inputs) |
|
soft_assignment = og_outputs['soft_assignment'][:, :-1, :-1] |
|
|
|
match_matrix = ( |
|
utils.soft_assignment_to_match_matrix(soft_assignment, MATCH_THRESHOLD) |
|
.numpy() |
|
.squeeze() |
|
) |
|
|
|
|
|
match_indices = np.argwhere(match_matrix) |
|
keep = [] |
|
for i in range(match_indices.shape[0]): |
|
match = match_indices[i, :] |
|
if (sp_features0[2][match[0]] > 0.0) and ( |
|
sp_features1[2][match[1]] > 0.0 |
|
): |
|
keep.append(i) |
|
match_indices = match_indices[keep] |
|
|
|
|
|
match_kp0s = [] |
|
match_kp1s = [] |
|
match_confidences = [] |
|
for match in match_indices: |
|
match_kp0s.append(sp_features0[0][match[0], :]) |
|
match_kp1s.append(sp_features1[0][match[1], :]) |
|
match_confidences.append(soft_assignment[0, match[0], match[1]]) |
|
match_kp0s = np.array(match_kp0s) |
|
match_kp1s = np.array(match_kp1s) |
|
match_confidences = np.array(match_confidences) |
|
return match_kp0s, match_kp1s, match_confidences |
|
|
|
|
|
|
|
def _construct_inputs( |
|
self, |
|
width0, |
|
height0, |
|
width1, |
|
height1, |
|
sp_features0, |
|
sp_features1, |
|
dino_descriptors0, |
|
dino_descriptors1, |
|
): |
|
inputs = { |
|
'keypoints0': tf.convert_to_tensor( |
|
np.expand_dims(sp_features0[0], axis=0), |
|
dtype=tf.float32, |
|
), |
|
'keypoints1': tf.convert_to_tensor( |
|
np.expand_dims(sp_features1[0], axis=0), dtype=tf.float32 |
|
), |
|
'descriptors0': tf.convert_to_tensor( |
|
np.expand_dims(sp_features0[1], axis=0), dtype=tf.float32 |
|
), |
|
'descriptors1': tf.convert_to_tensor( |
|
np.expand_dims(sp_features1[1], axis=0), dtype=tf.float32 |
|
), |
|
'scores0': tf.convert_to_tensor( |
|
np.expand_dims(np.expand_dims(sp_features0[2], axis=0), axis=-1), |
|
dtype=tf.float32, |
|
), |
|
'scores1': tf.convert_to_tensor( |
|
np.expand_dims(np.expand_dims(sp_features1[2], axis=0), axis=-1), |
|
dtype=tf.float32, |
|
), |
|
'descriptors0_dino': tf.expand_dims(dino_descriptors0, axis=0), |
|
'descriptors1_dino': tf.expand_dims(dino_descriptors1, axis=0), |
|
'width0': tf.convert_to_tensor( |
|
np.expand_dims(width0, axis=0), dtype=tf.int32 |
|
), |
|
'width1': tf.convert_to_tensor( |
|
np.expand_dims(width1, axis=0), dtype=tf.int32 |
|
), |
|
'height0': tf.convert_to_tensor( |
|
np.expand_dims(height0, axis=0), dtype=tf.int32 |
|
), |
|
'height1': tf.convert_to_tensor( |
|
np.expand_dims(height1, axis=0), dtype=tf.int32 |
|
), |
|
} |
|
return inputs |
|
|