Spaces:
Configuration error
Configuration error
| # Copyright (C) 2024-present Naver Corporation. All rights reserved. | |
| # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). | |
| # | |
| # -------------------------------------------------------- | |
| # Main Retriever class | |
| # -------------------------------------------------------- | |
| import os | |
| import argparse | |
| import numpy as np | |
| import torch | |
| from mast3r.model import AsymmetricMASt3R | |
| from mast3r.retrieval.model import RetrievalModel, extract_local_features | |
| try: | |
| import faiss | |
| faiss.StandardGpuResources() # when loading the checkpoint, it will try to instanciate FaissGpuL2Index | |
| except AttributeError as e: | |
| import asmk.index | |
| class FaissCpuL2Index(asmk.index.FaissL2Index): | |
| def __init__(self, gpu_id): | |
| super().__init__() | |
| self.gpu_id = gpu_id | |
| def _faiss_index_flat(self, dim): | |
| """Return initialized faiss.IndexFlatL2""" | |
| return faiss.IndexFlatL2(dim) | |
| asmk.index.FaissGpuL2Index = FaissCpuL2Index | |
| from asmk import asmk_method # noqa | |
| def get_args_parser(): | |
| parser = argparse.ArgumentParser('Retrieval scores from a set of retrieval', add_help=False, allow_abbrev=False) | |
| parser.add_argument('--model', type=str, required=True, | |
| help="shortname of a retrieval model or path to the corresponding .pth") | |
| parser.add_argument('--input', type=str, required=True, | |
| help="directory containing images or a file containing a list of image paths") | |
| parser.add_argument('--outfile', type=str, required=True, help="numpy file where to store the matrix score") | |
| return parser | |
| def get_impaths(imlistfile): | |
| with open(imlistfile, 'r') as fid: | |
| impaths = [f for f in imlistfile.read().splitlines() if not f.startswith('#') | |
| and len(f) > 0] # ignore comments and empty lines | |
| return impaths | |
| def get_impaths_from_imdir(imdir, extensions=['png', 'jpg', 'PNG', 'JPG']): | |
| assert os.path.isdir(imdir) | |
| impaths = [os.path.join(imdir, f) for f in sorted(os.listdir(imdir)) if any(f.endswith(ext) for ext in extensions)] | |
| return impaths | |
| def get_impaths_from_imdir_or_imlistfile(input_imdir_or_imlistfile): | |
| if os.path.isfile(input_imdir_or_imlistfile): | |
| return get_impaths(input_imdir_or_imlistfile) | |
| else: | |
| return get_impaths_from_imdir(input_imdir_or_imlistfile) | |
| class Retriever(object): | |
| def __init__(self, modelname, backbone=None, device='cuda'): | |
| # load the model | |
| assert os.path.isfile(modelname), modelname | |
| print(f'Loading retrieval model from {modelname}') | |
| ckpt = torch.load(modelname, 'cpu') # TODO from pretrained to download it automatically | |
| ckpt_args = ckpt['args'] | |
| if backbone is None: | |
| backbone = AsymmetricMASt3R.from_pretrained(ckpt_args.pretrained) | |
| self.model = RetrievalModel( | |
| backbone, freeze_backbone=ckpt_args.freeze_backbone, prewhiten=ckpt_args.prewhiten, | |
| hdims=list(map(int, ckpt_args.hdims.split('_'))) if len(ckpt_args.hdims) > 0 else "", | |
| residual=getattr(ckpt_args, 'residual', False), postwhiten=ckpt_args.postwhiten, | |
| featweights=ckpt_args.featweights, nfeat=ckpt_args.nfeat | |
| ).to(device) | |
| self.device = device | |
| msg = self.model.load_state_dict(ckpt['model'], strict=False) | |
| assert all(k.startswith('backbone') for k in msg.missing_keys) | |
| assert len(msg.unexpected_keys) == 0 | |
| self.imsize = ckpt_args.imsize | |
| # load the asmk codebook | |
| dname, bname = os.path.split(modelname) # TODO they should both be in the same file ? | |
| bname_splits = bname.split('_') | |
| cache_codebook_fname = os.path.join(dname, '_'.join(bname_splits[:-1]) + '_codebook.pkl') | |
| assert os.path.isfile(cache_codebook_fname), cache_codebook_fname | |
| asmk_params = {'index': {'gpu_id': 0}, 'train_codebook': {'codebook': {'size': '64k'}}, | |
| 'build_ivf': {'kernel': {'binary': True}, 'ivf': {'use_idf': False}, | |
| 'quantize': {'multiple_assignment': 1}, 'aggregate': {}}, | |
| 'query_ivf': {'quantize': {'multiple_assignment': 5}, 'aggregate': {}, | |
| 'search': {'topk': None}, | |
| 'similarity': {'similarity_threshold': 0.0, 'alpha': 3.0}}} | |
| asmk_params['train_codebook']['codebook']['size'] = ckpt_args.nclusters | |
| self.asmk = asmk_method.ASMKMethod.initialize_untrained(asmk_params) | |
| self.asmk = self.asmk.train_codebook(None, cache_path=cache_codebook_fname) | |
| def __call__(self, input_imdir_or_imlistfile, outfile=None): | |
| # get impaths | |
| if isinstance(input_imdir_or_imlistfile, str): | |
| impaths = get_impaths_from_imdir_or_imlistfile(input_imdir_or_imlistfile) | |
| else: | |
| impaths = input_imdir_or_imlistfile # we're assuming a list has been passed | |
| print(f'Found {len(impaths)} images') | |
| # build the database | |
| feat, ids = extract_local_features(self.model, impaths, self.imsize, tocpu=True, device=self.device) | |
| feat = feat.cpu().numpy() | |
| ids = ids.cpu().numpy() | |
| asmk_dataset = self.asmk.build_ivf(feat, ids) | |
| # we actually retrieve the same set of images | |
| metadata, query_ids, ranks, ranked_scores = asmk_dataset.query_ivf(feat, ids) | |
| # well ... scores are actually reordered according to ranks ... | |
| # so we redo it the other way around... | |
| scores = np.empty_like(ranked_scores) | |
| scores[np.arange(ranked_scores.shape[0])[:, None], ranks] = ranked_scores | |
| # save | |
| if outfile is not None: | |
| if os.path.isdir(os.path.dirname(outfile)): | |
| os.makedirs(os.path.dirname(outfile), exist_ok=True) | |
| np.save(outfile, scores) | |
| print(f'Scores matrix saved in {outfile}') | |
| return scores | |