98 lines
2.8 KiB
Python
Executable File
98 lines
2.8 KiB
Python
Executable File
import json
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pickle as pickle
|
|
|
|
import LFUtilities
|
|
import WebAppSettings as settings
|
|
#from ORBRescorer import ORBRescorer
|
|
from BEBLIDRescorer import BEBLIDRescorer
|
|
#from LATCHRescorer import LATCHRescorer
|
|
from GEMSearcher import GEMSearcher
|
|
import GEMExtractor as fe
|
|
import ORBExtractor as orb_lf
|
|
import BEBLIDExtractor as beblid_lf
|
|
import LATCHExtractor as latch_lf
|
|
|
|
|
|
class Searcher:
|
|
|
|
GEM_THRESHOLD = 0.3
|
|
|
|
def __init__(self):
|
|
# self.dataset = h5py.File(settings.dataset_file, 'r')['rmac'][...]
|
|
|
|
# np.save('/media/Data/data/beni_culturali/deploy/dataset', self.dataset)
|
|
self.ids = np.loadtxt(settings.DATASET_IDS, dtype=str).tolist()
|
|
|
|
self.search_engine = GEMSearcher()
|
|
|
|
# self.orb_rescorer = ORBRescorer()
|
|
self.beblid_rescorer = BEBLIDRescorer()
|
|
#self.latch_rescorer = LATCHRescorer()
|
|
|
|
def get_id(self, idx):
|
|
return self.search_engine.get_id(idx)
|
|
|
|
def add(self, img_file, id):
|
|
# self.save(True)
|
|
|
|
#desc = fe.extract(img_file)
|
|
#orb = lf.extract(img_file)
|
|
#self.search_engine.add(desc, id)
|
|
#self.rescorer.add(orb)
|
|
|
|
#self.save()
|
|
print('added ' + id)
|
|
|
|
def remove(self, id):
|
|
#self.save(True)
|
|
#self.search_engine.remove(id)
|
|
#self.rescorer.remove(idx)
|
|
#self.save()
|
|
print('removed ' + id)
|
|
|
|
def search_by_id(self, query_id, k=10, rescorer=False, lf_impl='BEBLID'):
|
|
kq = k
|
|
if rescorer:
|
|
kq = settings.k_reorder
|
|
res = self.search_engine.search_by_id(query_id, kq)
|
|
if rescorer:
|
|
if lf_impl == 'BEBLID':
|
|
res_lf = self.beblid_rescorer.rescore_by_id(query_id, res)
|
|
#res = res_lf if res_lf else res[:k]
|
|
res = res_lf
|
|
return res
|
|
|
|
def search_by_img(self, query_img, k=10, rescorer=False, lf_impl='BEBLID'):
|
|
if rescorer:
|
|
print(lf_impl)
|
|
kq = k
|
|
if rescorer:
|
|
kq = settings.k_reorder
|
|
query_desc = fe.extract(query_img)
|
|
res = self.search_engine.search_by_img(query_desc, kq)
|
|
print(res[0])
|
|
if rescorer:
|
|
if lf_impl == 'BEBLID':
|
|
query_lf = beblid_lf.extract(query_img)
|
|
res_lf = self.beblid_rescorer.rescore_by_img(query_lf, res)
|
|
if res_lf:
|
|
label = res_lf[0][0].split('/')[0]
|
|
res = [label, res_lf[0][1]]
|
|
else:
|
|
res = None
|
|
else:
|
|
if res[0][1] > self.GEM_THRESHOLD:
|
|
label = res[0][0].split('/')[0]
|
|
res = [label, res[0][1]]
|
|
else:
|
|
res = None
|
|
|
|
return res
|
|
|
|
def save(self, is_backup=False):
|
|
self.search_engine.save(is_backup)
|
|
#self.rescorer.save(is_backup)
|