File size: 5,466 Bytes
404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af |
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 |
import os
import sys
import cv2
from pathlib import Path
import numpy as np
import torch
import torch.utils.data as data
from tqdm import tqdm
from copy import deepcopy
from torchvision.transforms import ToTensor
sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from alike import ALike, configs
dataset_root = "hseq/hpatches-sequences-release"
use_cuda = torch.cuda.is_available()
device = "cuda" if use_cuda else "cpu"
methods = ["alike-n", "alike-l", "alike-n-ms", "alike-l-ms"]
class HPatchesDataset(data.Dataset):
def __init__(self, root: str = dataset_root, alteration: str = "all"):
"""
Args:
root: dataset root path
alteration: # 'all', 'i' for illumination or 'v' for viewpoint
"""
assert Path(root).exists(), f"Dataset root path {root} dose not exist!"
self.root = root
# get all image file name
self.image0_list = []
self.image1_list = []
self.homographies = []
folders = [x for x in Path(self.root).iterdir() if x.is_dir()]
self.seqs = []
for folder in folders:
if alteration == "i" and folder.stem[0] != "i":
continue
if alteration == "v" and folder.stem[0] != "v":
continue
self.seqs.append(folder)
self.len = len(self.seqs)
assert self.len > 0, f"Can not find PatchDataset in path {self.root}"
def __getitem__(self, item):
folder = self.seqs[item]
imgs = []
homos = []
for i in range(1, 7):
img = cv2.imread(str(folder / f"{i}.ppm"), cv2.IMREAD_COLOR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # HxWxC
imgs.append(img)
if i != 1:
homo = np.loadtxt(str(folder / f"H_1_{i}")).astype("float32")
homos.append(homo)
return imgs, homos, folder.stem
def __len__(self):
return self.len
def name(self):
return self.__class__
def extract_multiscale(
model,
img,
scale_f=2**0.5,
min_scale=1.0,
max_scale=1.0,
min_size=0.0,
max_size=99999.0,
image_size_max=99999,
n_k=0,
sort=False,
):
H_, W_, three = img.shape
assert three == 3, "input image shape should be [HxWx3]"
old_bm = torch.backends.cudnn.benchmark
torch.backends.cudnn.benchmark = False # speedup
# ==================== image size constraint
image = deepcopy(img)
max_hw = max(H_, W_)
if max_hw > image_size_max:
ratio = float(image_size_max / max_hw)
image = cv2.resize(image, dsize=None, fx=ratio, fy=ratio)
# ==================== convert image to tensor
H, W, three = image.shape
image = ToTensor()(image).unsqueeze(0)
image = image.to(device)
s = 1.0 # current scale factor
keypoints, descriptors, scores, scores_maps, descriptor_maps = [], [], [], [], []
while s + 0.001 >= max(min_scale, min_size / max(H, W)):
if s - 0.001 <= min(max_scale, max_size / max(H, W)):
nh, nw = image.shape[2:]
# extract descriptors
with torch.no_grad():
descriptor_map, scores_map = model.extract_dense_map(image)
keypoints_, descriptors_, scores_, _ = model.dkd(
scores_map, descriptor_map
)
keypoints.append(keypoints_[0])
descriptors.append(descriptors_[0])
scores.append(scores_[0])
s /= scale_f
# down-scale the image for next iteration
nh, nw = round(H * s), round(W * s)
image = torch.nn.functional.interpolate(
image, (nh, nw), mode="bilinear", align_corners=False
)
# restore value
torch.backends.cudnn.benchmark = old_bm
keypoints = torch.cat(keypoints)
descriptors = torch.cat(descriptors)
scores = torch.cat(scores)
keypoints = (keypoints + 1) / 2 * keypoints.new_tensor([[W_ - 1, H_ - 1]])
if sort or 0 < n_k < len(keypoints):
indices = torch.argsort(scores, descending=True)
keypoints = keypoints[indices]
descriptors = descriptors[indices]
scores = scores[indices]
if 0 < n_k < len(keypoints):
keypoints = keypoints[0:n_k]
descriptors = descriptors[0:n_k]
scores = scores[0:n_k]
return {"keypoints": keypoints, "descriptors": descriptors, "scores": scores}
def extract_method(m):
hpatches = HPatchesDataset(root=dataset_root, alteration="all")
model = m[:7]
min_scale = 0.3 if m[8:] == "ms" else 1.0
model = ALike(**configs[model], device=device, top_k=0, scores_th=0.2, n_limit=5000)
progbar = tqdm(hpatches, desc="Extracting for {}".format(m))
for imgs, homos, seq_name in progbar:
for i in range(1, 7):
img = imgs[i - 1]
pred = extract_multiscale(
model, img, min_scale=min_scale, max_scale=1, sort=False, n_k=5000
)
kpts, descs, scores = pred["keypoints"], pred["descriptors"], pred["scores"]
with open(os.path.join(dataset_root, seq_name, f"{i}.ppm.{m}"), "wb") as f:
np.savez(
f,
keypoints=kpts.cpu().numpy(),
scores=scores.cpu().numpy(),
descriptors=descs.cpu().numpy(),
)
if __name__ == "__main__":
for method in methods:
extract_method(method)
|