image retrieval on oxford and paris

facebookresearch · Jul 15, 2021 · 30aedce · 30aedce
1 parent 9085367
commit 30aedce
Show file tree

Hide file tree

Showing 4 changed files with 353 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -278,6 +278,21 @@ git clone https://github.com/davisvideochallenge/davis2017-evaluation $HOME/davi
 python $HOME/davis2017-evaluation/evaluation_method.py --task semi-supervised --results_path /path/to/saving_dir --davis_path $HOME/davis-2017/DAVIS/
 ```
 
+## Evaluation: Image Retrieval on revisited Oxford and Paris
+Step 1: Prepare revisited Oxford and Paris by following [this repo](https://github.com/filipradenovic/revisitop).
+
+Step 2: Image retrieval (if you do not specify weights with `--pretrained_weights` then by default [DINO weights pretrained on Google Landmark v2 dataset](https://dl.fbaipublicfiles.com/dino/dino_vitsmall16_googlelandmark_pretrain/dino_vitsmall16_googlelandmark_pretrain.pth) will be used).
+
+Paris:
+```
+python -m torch.distributed.launch --use_env --nproc_per_node=1 eval_image_retrieval.py --imsize 512 --multiscale 1 --data_path /path/to/revisited_paris_oxford/ --dataset rparis6k
+```
+
+Oxford:
+```
+python -m torch.distributed.launch --use_env --nproc_per_node=1 eval_image_retrieval.py --imsize 224 --multiscale 0 --data_path /path/to/revisited_paris_oxford/ --dataset roxford5k
+```
+
 ## License
 This repository is released under the Apache 2.0 license as found in the [LICENSE](LICENSE) file.
 

diff --git a/eval_image_retrieval.py b/eval_image_retrieval.py
@@ -0,0 +1,201 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# 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
+# 
+#     http://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.
+import os
+import sys
+import pickle
+import argparse
+
+import torch
+from torch import nn
+import torch.distributed as dist
+import torch.backends.cudnn as cudnn
+from torchvision import models as torchvision_models
+from torchvision import transforms as pth_transforms
+from PIL import Image, ImageFile
+import numpy as np
+
+import utils
+import vision_transformer as vits
+from eval_knn import extract_features
+
+
+class OxfordParisDataset(torch.utils.data.Dataset):
+    def __init__(self, dir_main, dataset, split, transform=None, imsize=None):
+        if dataset not in ['roxford5k', 'rparis6k']:
+            raise ValueError('Unknown dataset: {}!'.format(dataset))
+
+        # loading imlist, qimlist, and gnd, in cfg as a dict
+        gnd_fname = os.path.join(dir_main, dataset, 'gnd_{}.pkl'.format(dataset))
+        with open(gnd_fname, 'rb') as f:
+            cfg = pickle.load(f)
+        cfg['gnd_fname'] = gnd_fname
+        cfg['ext'] = '.jpg'
+        cfg['qext'] = '.jpg'
+        cfg['dir_data'] = os.path.join(dir_main, dataset)
+        cfg['dir_images'] = os.path.join(cfg['dir_data'], 'jpg')
+        cfg['n'] = len(cfg['imlist'])
+        cfg['nq'] = len(cfg['qimlist'])
+        cfg['im_fname'] = config_imname
+        cfg['qim_fname'] = config_qimname
+        cfg['dataset'] = dataset
+        self.cfg = cfg
+
+        self.samples = cfg["qimlist"] if split == "query" else cfg["imlist"]
+        self.transform = transform
+        self.imsize = imsize
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        path = os.path.join(self.cfg["dir_images"], self.samples[index] + ".jpg")
+        ImageFile.LOAD_TRUNCATED_IMAGES = True
+        with open(path, 'rb') as f:
+            img = Image.open(f)
+            img = img.convert('RGB')
+        if self.imsize is not None:
+            img.thumbnail((self.imsize, self.imsize), Image.ANTIALIAS)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, index
+
+
+def config_imname(cfg, i):
+    return os.path.join(cfg['dir_images'], cfg['imlist'][i] + cfg['ext'])
+
+
+def config_qimname(cfg, i):
+    return os.path.join(cfg['dir_images'], cfg['qimlist'][i] + cfg['qext'])
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('Image Retrieval on revisited Paris and Oxford')
+    parser.add_argument('--data_path', default='/path/to/revisited_paris_oxford/', type=str)
+    parser.add_argument('--dataset', default='roxford5k', type=str, choices=['roxford5k', 'rparis6k'])
+    parser.add_argument('--multiscale', default=False, type=utils.bool_flag)
+    parser.add_argument('--imsize', default=224, type=int, help='Image size (square)')
+    parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
+    parser.add_argument('--use_cuda', default=True, type=utils.bool_flag)
+    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
+    parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
+    parser.add_argument("--checkpoint_key", default="teacher", type=str,
+        help='Key to use in the checkpoint (example: "teacher")')
+    parser.add_argument('--num_workers', default=10, type=int, help='Number of data loading workers per GPU.')
+    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
+        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
+    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
+    args = parser.parse_args()
+
+    utils.init_distributed_mode(args)
+    print("git:\n  {}\n".format(utils.get_sha()))
+    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+    cudnn.benchmark = True
+
+    # ============ preparing data ... ============
+    transform = pth_transforms.Compose([
+        pth_transforms.ToTensor(),
+        pth_transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+    dataset_train = OxfordParisDataset(args.data_path, args.dataset, split="train", transform=transform, imsize=args.imsize)
+    dataset_query = OxfordParisDataset(args.data_path, args.dataset, split="query", transform=transform, imsize=args.imsize)
+    sampler = torch.utils.data.DistributedSampler(dataset_train, shuffle=False)
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train,
+        sampler=sampler,
+        batch_size=1,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+    data_loader_query = torch.utils.data.DataLoader(
+        dataset_query,
+        batch_size=1,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+    print(f"train: {len(dataset_train)} imgs / query: {len(dataset_query)} imgs")
+
+    # ============ building network ... ============
+    if "vit" in args.arch:
+        model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0)
+        print(f"Model {args.arch} {args.patch_size}x{args.patch_size} built.")
+    elif "xcit" in args.arch:
+        model = torch.hub.load('facebookresearch/xcit', args.arch, num_classes=0)
+    elif args.arch in torchvision_models.__dict__.keys():
+        model = torchvision_models.__dict__[args.arch](num_classes=0)
+    else:
+        print(f"Architecture {args.arch} non supported")
+        sys.exit(1)
+    if args.use_cuda:
+        model.cuda()
+    model.eval()
+
+    # load pretrained weights
+    if os.path.isfile(args.pretrained_weights):
+        state_dict = torch.load(args.pretrained_weights, map_location="cpu")
+        if args.checkpoint_key is not None and args.checkpoint_key in state_dict:
+            print(f"Take key {args.checkpoint_key} in provided checkpoint dict")
+            state_dict = state_dict[args.checkpoint_key]
+        # remove `module.` prefix
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+        # remove `backbone.` prefix induced by multicrop wrapper
+        state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+        msg = model.load_state_dict(state_dict, strict=False)
+        print('Pretrained weights found at {} and loaded with msg: {}'.format(args.pretrained_weights, msg))
+    elif args.arch == "vit_small" and args.patch_size == 16:
+        print("Since no pretrained weights have been provided, we load pretrained DINO weights on Google Landmark v2.")
+        model.load_state_dict(torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/dino_vitsmall16_googlelandmark_pretrain/dino_vitsmall16_googlelandmark_pretrain.pth"))
+    else:
+        print("Warning: We use random weights.")
+
+    ############################################################################
+    # Step 1: extract features
+    train_features = extract_features(model, data_loader_train, args.use_cuda, multiscale=args.multiscale)
+    query_features = extract_features(model, data_loader_query, args.use_cuda, multiscale=args.multiscale)
+
+    if utils.get_rank() == 0:  # only rank 0 will work from now on
+        # normalize features
+        train_features = nn.functional.normalize(train_features, dim=1, p=2)
+        query_features = nn.functional.normalize(query_features, dim=1, p=2)
+
+        ############################################################################
+        # Step 2: similarity
+        sim = torch.mm(train_features, query_features.T)
+        ranks = torch.argsort(-sim, dim=0).cpu().numpy()
+
+        ############################################################################
+        # Step 3: evaluate
+        gnd = dataset_train.cfg['gnd']
+        # evaluate ranks
+        ks = [1, 5, 10]
+        # search for easy & hard
+        gnd_t = []
+        for i in range(len(gnd)):
+            g = {}
+            g['ok'] = np.concatenate([gnd[i]['easy'], gnd[i]['hard']])
+            g['junk'] = np.concatenate([gnd[i]['junk']])
+            gnd_t.append(g)
+        mapM, apsM, mprM, prsM = utils.compute_map(ranks, gnd_t, ks)
+        # search for hard
+        gnd_t = []
+        for i in range(len(gnd)):
+            g = {}
+            g['ok'] = np.concatenate([gnd[i]['hard']])
+            g['junk'] = np.concatenate([gnd[i]['junk'], gnd[i]['easy']])
+            gnd_t.append(g)
+        mapH, apsH, mprH, prsH = utils.compute_map(ranks, gnd_t, ks)
+        print('>> {}: mAP M: {}, H: {}'.format(args.dataset, np.around(mapM*100, decimals=2), np.around(mapH*100, decimals=2)))
+        print('>> {}: mP@k{} M: {}, H: {}'.format(args.dataset, np.array(ks), np.around(mprM*100, decimals=2), np.around(mprH*100, decimals=2)))
+    dist.barrier()
diff --git a/eval_knn.py b/eval_knn.py
@@ -21,6 +21,7 @@
 import torch.backends.cudnn as cudnn
 from torchvision import datasets
 from torchvision import transforms as pth_transforms
+from torchvision import models as torchvision_models
 
 import utils
 import vision_transformer as vits
@@ -60,6 +61,8 @@ def extract_feature_pipeline(args):
         print(f"Model {args.arch} {args.patch_size}x{args.patch_size} built.")
     elif "xcit" in args.arch:
         model = torch.hub.load('facebookresearch/xcit', args.arch, num_classes=0)
+    elif args.arch in torchvision_models.__dict__.keys():
+        model = torchvision_models.__dict__[args.arch](num_classes=0)
     else:
         print(f"Architecture {args.arch} non supported")
         sys.exit(1)
@@ -69,9 +72,9 @@ def extract_feature_pipeline(args):
 
     # ============ extract features ... ============
     print("Extracting features for train set...")
-    train_features = extract_features(model, data_loader_train)
+    train_features = extract_features(model, data_loader_train, args.use_cuda)
     print("Extracting features for val set...")
-    test_features = extract_features(model, data_loader_val)
+    test_features = extract_features(model, data_loader_val, args.use_cuda)
 
     if utils.get_rank() == 0:
         train_features = nn.functional.normalize(train_features, dim=1, p=2)
@@ -89,18 +92,21 @@ def extract_feature_pipeline(args):
 
 
 @torch.no_grad()
-def extract_features(model, data_loader):
+def extract_features(model, data_loader, use_cuda=True, multiscale=False):
     metric_logger = utils.MetricLogger(delimiter="  ")
     features = None
     for samples, index in metric_logger.log_every(data_loader, 10):
         samples = samples.cuda(non_blocking=True)
         index = index.cuda(non_blocking=True)
-        feats = model(samples).clone()
+        if multiscale:
+            feats = utils.multi_scale(samples, model)
+        else:
+            feats = model(samples).clone()
 
         # init storage feature matrix
         if dist.get_rank() == 0 and features is None:
             features = torch.zeros(len(data_loader.dataset), feats.shape[-1])
-            if args.use_cuda:
+            if use_cuda:
                 features = features.cuda(non_blocking=True)
             print(f"Storing features into tensor of shape {features.shape}")
 
@@ -125,7 +131,7 @@ def extract_features(model, data_loader):
 
         # update storage feature matrix
         if dist.get_rank() == 0:
-            if args.use_cuda:
+            if use_cuda:
                 features.index_copy_(0, index_all, torch.cat(output_l))
             else:
                 features.index_copy_(0, index_all.cpu(), torch.cat(output_l).cpu())
@@ -191,7 +197,7 @@ def __getitem__(self, idx):
     parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
     parser.add_argument('--use_cuda', default=True, type=utils.bool_flag,
         help="Should we store the features on GPU? We recommend setting this to False if you encounter OOM")
-    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture (support only ViT and XCiT atm).')
+    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
     parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
     parser.add_argument("--checkpoint_key", default="teacher", type=str,
         help='Key to use in the checkpoint (example: "teacher")')