File size: 8,351 Bytes
cff8c58 |
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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 |
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wrapper for performing SuperPoint inference."""
import math
from typing import Optional, Tuple
import cv2
import numpy as np
from omniglue import utils
import tensorflow.compat.v1 as tf1
class SuperPointExtract:
"""Class to initialize SuperPoint model and extract features from an image.
To stay consistent with SuperPoint training and eval configurations, resize
images to (320x240) or (640x480).
Attributes
model_path: string, filepath to saved SuperPoint TF1 model weights.
"""
def __init__(self, model_path: str):
self.model_path = model_path
self._graph = tf1.Graph()
self._sess = tf1.Session(graph=self._graph)
tf1.saved_model.loader.load(
self._sess, [tf1.saved_model.tag_constants.SERVING], model_path
)
def __call__(
self,
image,
segmentation_mask=None,
num_features=1024,
pad_random_features=False,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
return self.compute(
image,
segmentation_mask=segmentation_mask,
num_features=num_features,
pad_random_features=pad_random_features,
)
def compute(
self,
image: np.ndarray,
segmentation_mask: Optional[np.ndarray] = None,
num_features: int = 1024,
pad_random_features: bool = False,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Feeds image through SuperPoint model to extract keypoints and features.
Args:
image: (H, W, 3) numpy array, decoded image bytes.
segmentation_mask: (H, W) binary numpy array or None. If not None,
extracted keypoints are restricted to being within the mask.
num_features: max number of features to extract (or 0 to indicate keeping
all extracted features).
pad_random_features: if True, adds randomly sampled keypoints to the
output such that there are exactly 'num_features' keypoints. Descriptors
for these sampled keypoints are taken from the network's descriptor map
output, and scores are set to 0. No action taken if num_features = 0.
Returns:
keypoints: (N, 2) numpy array, coordinates of keypoints as floats.
descriptors: (N, 256) numpy array, descriptors for keypoints as floats.
scores: (N, 1) numpy array, confidence values for keypoints as floats.
"""
# Resize image so both dimensions are divisible by 8.
image, keypoint_scale_factors = self._resize_input_image(image)
if segmentation_mask is not None:
segmentation_mask, _ = self._resize_input_image(
segmentation_mask, interpolation=cv2.INTER_NEAREST
)
assert (
segmentation_mask is None
or image.shape[:2] == segmentation_mask.shape[:2]
)
# Preprocess and feed-forward image.
image_preprocessed = self._preprocess_image(image)
input_image_tensor = self._graph.get_tensor_by_name('superpoint/image:0')
output_prob_nms_tensor = self._graph.get_tensor_by_name(
'superpoint/prob_nms:0'
)
output_desc_tensors = self._graph.get_tensor_by_name(
'superpoint/descriptors:0'
)
out = self._sess.run(
[output_prob_nms_tensor, output_desc_tensors],
feed_dict={input_image_tensor: np.expand_dims(image_preprocessed, 0)},
)
# Format output from network.
keypoint_map = np.squeeze(out[0])
descriptor_map = np.squeeze(out[1])
if segmentation_mask is not None:
keypoint_map = np.where(segmentation_mask, keypoint_map, 0.0)
keypoints, descriptors, scores = self._extract_superpoint_output(
keypoint_map, descriptor_map, num_features, pad_random_features
)
# Rescale keypoint locations to match original input image size, and return.
keypoints = keypoints / keypoint_scale_factors
return (keypoints, descriptors, scores)
def _resize_input_image(self, image, interpolation=cv2.INTER_LINEAR):
"""Resizes image such that both dimensions are divisble by 8."""
# Calculate new image dimensions and per-dimension resizing scale factor.
new_dim = [-1, -1]
keypoint_scale_factors = [1.0, 1.0]
for i in range(2):
dim_size = image.shape[i]
mod_eight = dim_size % 8
if mod_eight < 4:
# Round down to nearest multiple of 8.
new_dim[i] = dim_size - mod_eight
elif mod_eight >= 4:
# Round up to nearest multiple of 8.
new_dim[i] = dim_size + (8 - mod_eight)
keypoint_scale_factors[i] = (new_dim[i] - 1) / (dim_size - 1)
# Resize and return image + scale factors.
new_dim = new_dim[::-1] # Convert from (row, col) to (x,y).
keypoint_scale_factors = keypoint_scale_factors[::-1]
image = cv2.resize(image, tuple(new_dim), interpolation=interpolation)
return image, keypoint_scale_factors
def _preprocess_image(self, image):
"""Converts image to grayscale and normalizes values for model input."""
image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
image = np.expand_dims(image, 2)
image = image.astype(np.float32)
image = image / 255.0
return image
def _extract_superpoint_output(
self,
keypoint_map,
descriptor_map,
keep_k_points=512,
pad_random_features=False,
):
"""Converts from raw SuperPoint output (feature maps) into numpy arrays.
If keep_k_points is 0, then keep all detected keypoints. Otherwise, sort by
confidence and keep only the top k confidence keypoints.
Args:
keypoint_map: (H, W, 1) numpy array, raw output confidence values from
SuperPoint model.
descriptor_map: (H, W, 256) numpy array, raw output descriptors from
SuperPoint model.
keep_k_points: int, number of keypoints to keep (or 0 to indicate keeping
all detected keypoints).
pad_random_features: if True, adds randomly sampled keypoints to the
output such that there are exactly 'num_features' keypoints. Descriptors
for these sampled keypoints are taken from the network's descriptor map
output, and scores are set to 0. No action taken if keep_k_points = 0.
Returns:
keypoints: (N, 2) numpy array, image coordinates (x, y) of keypoints as
floats.
descriptors: (N, 256) numpy array, descriptors for keypoints as floats.
scores: (N, 1) numpy array, confidence values for keypoints as floats.
"""
def _select_k_best(points, k):
sorted_prob = points[points[:, 2].argsort(), :]
start = min(k, points.shape[0])
return sorted_prob[-start:, :2], sorted_prob[-start:, 2]
keypoints = np.where(keypoint_map > 0)
prob = keypoint_map[keypoints[0], keypoints[1]]
keypoints = np.stack([keypoints[0], keypoints[1], prob], axis=-1)
# Keep only top k points, or all points if keep_k_points param is 0.
if keep_k_points == 0:
keep_k_points = keypoints.shape[0]
keypoints, scores = _select_k_best(keypoints, keep_k_points)
# Optionally, pad with random features (and confidence scores of 0).
image_shape = np.array(keypoint_map.shape[:2])
if pad_random_features and (keep_k_points > keypoints.shape[0]):
num_pad = keep_k_points - keypoints.shape[0]
keypoints_pad = (image_shape - 1) * np.random.uniform(size=(num_pad, 2))
keypoints = np.concatenate((keypoints, keypoints_pad))
scores_pad = np.zeros((num_pad))
scores = np.concatenate((scores, scores_pad))
# Lookup descriptors via bilinear interpolation.
# TODO: batch descriptor lookup with bilinear interpolation.
keypoints[:, [0, 1]] = keypoints[:, [1, 0]] # Swap from (row,col) to (x,y).
descriptors = []
for kp in keypoints:
descriptors.append(utils.lookup_descriptor_bilinear(kp, descriptor_map))
descriptors = np.array(descriptors)
return keypoints, descriptors, scores
|