-
Notifications
You must be signed in to change notification settings - Fork 0
/
evaluate_dinov2_qkv.py
114 lines (90 loc) · 3.23 KB
/
evaluate_dinov2_qkv.py
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
import os
import einops as ein
import numpy as np
import torch
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
from src.model.dinov2extractor import DinoV2FeatureExtractor
from src.utils.dataset import configdataset
from src.utils.download import download_datasets
from src.utils.evaluate import compute_map
def process_image(img_path, extractor, device):
img = Image.open(img_path).convert("RGB")
transform = transforms.Compose(
[
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
img_tensor = transform(img)[None, ...].to(device)
features = extractor(img_tensor).cpu()
patches_h, patches_w = img_tensor.shape[-2] // 14, img_tensor.shape[-1] // 14
return (
ein.rearrange(features[0], "(h w) d -> h w d", h=patches_h, w=patches_w)
.flatten(1)
.mean(dim=0)
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
extractor = DinoV2FeatureExtractor(
model_type="dinov2_vits14", layer=11, facet="value", device=device
)
def extract_features_with_dino(image_paths):
features_list = []
for image_path in tqdm(image_paths, desc="Extracting Features"):
features = process_image(image_path, extractor, device=device).numpy()
features /= np.linalg.norm(features, keepdims=True)
features_list.append(features)
return np.vstack(features_list)
data_root = "data"
download_datasets(data_root)
test_dataset = "rparis6k"
cfg = configdataset(test_dataset, os.path.join(data_root, "datasets"))
query_images = [
os.path.join(data_root, "datasets", test_dataset, "jpg", f + ".jpg")
for f in cfg["qimlist"]
]
database_images = [
os.path.join(data_root, "datasets", test_dataset, "jpg", f + ".jpg")
for f in cfg["imlist"]
]
print(f">> {test_dataset}: Extracting query features...")
Q = extract_features_with_dino(query_images).T
print(f">> {test_dataset}: Extracting database features...")
X = extract_features_with_dino(database_images).T
print(f">> {test_dataset}: Retrieval...")
sim = np.dot(X.T, Q)
ranks = np.argsort(-sim, axis=0)
gnd = cfg["gnd"]
ks = [1, 5, 10]
gnd_t = []
for i in range(len(gnd)):
g = {
"ok": np.concatenate([gnd[i]["easy"]]),
"junk": np.concatenate([gnd[i]["junk"], gnd[i]["hard"]]),
}
gnd_t.append(g)
mapE, apsE, mprE, prsE = compute_map(ranks, gnd_t, ks)
gnd_t = []
for i in range(len(gnd)):
g = {
"ok": np.concatenate([gnd[i]["easy"], gnd[i]["hard"]]),
"junk": np.concatenate([gnd[i]["junk"]]),
}
gnd_t.append(g)
mapM, apsM, mprM, prsM = compute_map(ranks, gnd_t, ks)
gnd_t = []
for i in range(len(gnd)):
g = {
"ok": np.concatenate([gnd[i]["hard"]]),
"junk": np.concatenate([gnd[i]["junk"], gnd[i]["easy"]]),
}
gnd_t.append(g)
mapH, apsH, mprH, prsH = compute_map(ranks, gnd_t, ks)
print(
f">> {test_dataset}: mAP E: {np.around(mapE*100, decimals=2)}, M: {np.around(mapM*100, decimals=2)}, H: {np.around(mapH*100, decimals=2)}"
)
print(
f">> {test_dataset}: mP@k{ks} E: {np.around(mprE*100, decimals=2)}, M: {np.around(mprM*100, decimals=2)}, H: {np.around(mprH*100, decimals=2)}"
)