83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
import cv2
|
|
import numpy as np
|
|
import pickle as pickle
|
|
|
|
import LFUtilities
|
|
import ImageRecognitionSettings as settings
|
|
#from BEBLIDRescorerDB import BEBLIDRescorerDB
|
|
#from BEBLIDRescorerFAISS import BEBLIDRescorerFAISS
|
|
from BEBLIDRescorer import BEBLIDRescorer
|
|
import SearcherParameters
|
|
|
|
from FAISSSearchEngine2 import FAISSSearchEngine
|
|
import FeatureExtractor as fe
|
|
import BEBLIDExtractorQ as lfQ
|
|
import BEBLIDExtractorD as lfD
|
|
import logging
|
|
|
|
|
|
class Searcher:
|
|
|
|
def __init__(self):
|
|
# self.dataset = h5py.File(settings.dataset_file, 'r')['rmac'][...]
|
|
|
|
# np.save('/media/Data/data/beni_culturali/deploy/dataset', self.dataset)
|
|
self.search_engine = FAISSSearchEngine()
|
|
self.rescorer = BEBLIDRescorer()
|
|
|
|
def get_indexed_ids(self):
|
|
return self.search_engine.get_indexed_ids()
|
|
|
|
def get_id(self, idx):
|
|
return self.search_engine.get_id(idx)
|
|
|
|
def exists(self, doc_id):
|
|
return self.search_engine.exists(doc_id)
|
|
|
|
def add(self, img_file, doc_id):
|
|
desc = fe.extract(img_file)
|
|
self.search_engine.add(desc, doc_id)
|
|
|
|
kp, des = lfD.extract(img_file)
|
|
self.rescorer.add(doc_id, kp, des)
|
|
|
|
#orb = lf.extract(img_file)
|
|
|
|
self.save(True)
|
|
logging.info('added ' + doc_id)
|
|
|
|
def remove(self, doc_id):
|
|
self.search_engine.remove(doc_id)
|
|
self.rescorer.remove(doc_id)
|
|
self.save(True)
|
|
logging.info('removed ' + doc_id)
|
|
|
|
def search_by_id(self, query_id, k=10, search_threshold=0.25, search_deep_level=1):
|
|
kq = k
|
|
if search_deep_level > 0:
|
|
kq = SearcherParameters.SEARCH_DEEP_K[search_deep_level]
|
|
res = self.search_engine.search_by_id(query_id, kq)
|
|
if search_deep_level > 0:
|
|
res_lf = self.rescorer.rescore_by_id(query_id, res)
|
|
res = res_lf if res_lf else res[:k]
|
|
res = [result for result in res if result[1] >= search_threshold]
|
|
return res
|
|
|
|
def search_by_img(self, query_img, k=10, search_threshold=0.25, search_deep_level=1):
|
|
kq = k
|
|
if search_deep_level:
|
|
kq = SearcherParameters.SEARCH_DEEP_K[search_deep_level]
|
|
query_desc = fe.extract(query_img)
|
|
res = self.search_engine.search_by_img(query_desc, kq)
|
|
if search_deep_level > 0:
|
|
query_lf = lfQ.extract(query_img)
|
|
res_lf = self.rescorer.rescore_by_img(query_lf, res)
|
|
#res = res_lf if res_lf else res[:k]
|
|
res = res_lf if res_lf else res[:k]
|
|
res = [result for result in res if result[1] >= search_threshold]
|
|
return res
|
|
|
|
def save(self, is_backup=False):
|
|
self.search_engine.save(is_backup)
|
|
self.rescorer.save(is_backup)
|