Files
ieemoo-ai-searchv2/utils/monitor.py
2022-11-22 15:32:06 +08:00

62 lines
2.0 KiB
Python

import os
import json
import h5py
import numpy as np
from utils.config import cfg
from cirtorch.networks.imageretrievalnet import init_network, extract_vectors
from utils.tools import createNet,ManagingFeature
class Moniting:
def __init__(self, barcode = ''):
self.barcode = barcode
self.jsonpath = cfg.MONITORPATH
self.MF = ManagingFeature()
if not os.path.exists(self.jsonpath):
jsontext = {"monitor":[]}
jsondata = json.dumps(jsontext)
f = open(self.jsonpath, 'w+', encoding='utf-8')
f.write(jsondata)
f.close()
def add(self):
with open(self.jsonpath, 'r', encoding='utf-8') as add_f:
jsonfile = json.load(add_f)
add_f.close()
data = list(set(jsonfile['monitor']+self.barcode))
jsonfile['monitor'] = data
with open(self.jsonpath, 'w') as add_f:
json.dump(jsonfile, add_f)
add_f.close()
def search(self):
with open(self.jsonpath, 'r', encoding='utf-8') as f:
jsonfile = json.load(f)
f.close()
data = set(jsonfile['monitor'])
if self.barcode in data:
return 'success'
else:
return 'nomatch'
def update(self, net, transform, ms):#update monitor.json
Dict = {}
path = []
collectbarcode = set()
for name in os.listdir(cfg.IMG_DIR_TOTAL):
barcode = name.split('_')[-1].split('.')[0]
path.append(os.sep.join([cfg.IMG_DIR_TOTAL, name]))
collectbarcode.add(barcode)
vecs, img_paths = extract_vectors(net, path, cfg.RESIZE, transform, ms=ms)
data = list(vecs.detach().cpu().numpy().T)
for code, feature in zip(img_paths, data):
barcode = code.split('_')[-1].split('.')[0]
feature = feature.tolist()
self.MF.addfeature(barcode, feature)
Moniting(list(collectbarcode)).add()
if __name__ == '__main__':
barcode = '1'
mo = Moniting(barcode)
print(mo.search())
mo.add(barcode)