Extension_image_recognition/LFSearcher.py

55 lines
1.4 KiB
Python
Executable File

import cv2
import numpy as np
import pickle as pickle
import LFUtilities
import WebAppSettings as settings
from BEBLIDSearcher import BEBLIDSearcher
import GEMExtractor as fe
import ORBExtractor as lf
import BEBLIDExtractor as beblid_lf
import LATCHExtractor as latch_lf
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 = BEBLIDSearcher()
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.save()
print('added ' + id)
def remove(self, id):
self.save(True)
self.search_engine.remove(id)
self.save()
print('removed ' + id)
def search_by_id(self, query_id, k=10):
kq = k
res = self.search_engine.search_by_id(query_id)
return res
def search_by_img(self, query_img, k=10):
kq = k
query_desc = beblid_lf.extract(query_img)
res = self.search_engine.search_by_img(query_desc)
return res
def save(self, is_backup=False):
self.search_engine.save(is_backup)
#self.rescorer.save(is_backup)