File size: 4,070 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
8b973ee
404d2af
8b973ee
 
404d2af
 
8b973ee
404d2af
 
8b973ee
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.")