Spaces:
Running
Running
File size: 4,070 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 |
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 |
import argparse
import numpy as np
import imageio
import torch
from tqdm import tqdm
import scipy
import scipy.io
import scipy.misc
from lib.model_test import D2Net
from lib.utils import preprocess_image
from lib.pyramid import process_multiscale
# CUDA
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
# Argument parsing
parser = argparse.ArgumentParser(description="Feature extraction script")
parser.add_argument(
"--image_list_file",
type=str,
required=True,
help="path to a file containing a list of images to process",
)
parser.add_argument(
"--preprocessing",
type=str,
default="caffe",
help="image preprocessing (caffe or torch)",
)
parser.add_argument(
"--model_file", type=str, default="models/d2_tf.pth", help="path to the full model"
)
parser.add_argument(
"--max_edge", type=int, default=1600, help="maximum image size at network input"
)
parser.add_argument(
"--max_sum_edges",
type=int,
default=2800,
help="maximum sum of image sizes at network input",
)
parser.add_argument(
"--output_extension", type=str, default=".d2-net", help="extension for the output"
)
parser.add_argument(
"--output_type", type=str, default="npz", help="output file type (npz or mat)"
)
parser.add_argument(
"--multiscale",
dest="multiscale",
action="store_true",
help="extract multiscale features",
)
parser.set_defaults(multiscale=False)
parser.add_argument(
"--no-relu",
dest="use_relu",
action="store_false",
help="remove ReLU after the dense feature extraction module",
)
parser.set_defaults(use_relu=True)
args = parser.parse_args()
print(args)
# Creating CNN model
model = D2Net(model_file=args.model_file, use_relu=args.use_relu, use_cuda=use_cuda)
# Process the file
with open(args.image_list_file, "r") as f:
lines = f.readlines()
for line in tqdm(lines, total=len(lines)):
path = line.strip()
image = imageio.imread(path)
if len(image.shape) == 2:
image = image[:, :, np.newaxis]
image = np.repeat(image, 3, -1)
# TODO: switch to PIL.Image due to deprecation of scipy.misc.imresize.
resized_image = image
if max(resized_image.shape) > args.max_edge:
resized_image = scipy.misc.imresize(
resized_image, args.max_edge / max(resized_image.shape)
).astype("float")
if sum(resized_image.shape[:2]) > args.max_sum_edges:
resized_image = scipy.misc.imresize(
resized_image, args.max_sum_edges / sum(resized_image.shape[:2])
).astype("float")
fact_i = image.shape[0] / resized_image.shape[0]
fact_j = image.shape[1] / resized_image.shape[1]
input_image = preprocess_image(resized_image, preprocessing=args.preprocessing)
with torch.no_grad():
if args.multiscale:
keypoints, scores, descriptors = process_multiscale(
torch.tensor(
input_image[np.newaxis, :, :, :].astype(np.float32), device=device
),
model,
)
else:
keypoints, scores, descriptors = process_multiscale(
torch.tensor(
input_image[np.newaxis, :, :, :].astype(np.float32), device=device
),
model,
scales=[1],
)
# Input image coordinates
keypoints[:, 0] *= fact_i
keypoints[:, 1] *= fact_j
# i, j -> u, v
keypoints = keypoints[:, [1, 0, 2]]
if args.output_type == "npz":
with open(path + args.output_extension, "wb") as output_file:
np.savez(
output_file, keypoints=keypoints, scores=scores, descriptors=descriptors
)
elif args.output_type == "mat":
with open(path + args.output_extension, "wb") as output_file:
scipy.io.savemat(
output_file,
{"keypoints": keypoints, "scores": scores, "descriptors": descriptors},
)
else:
raise ValueError("Unknown output type.")
|