class Pipeline(object):
def __init__(self, args):
'''
sum of step is: 4:0-4
:param n_step:
'''
self.n_step = args.n_step
self.pool_size = args.pool_size
self.args = args
self.hash = Hash(hash_size=self.args.hash_size)
self.schedule()
def multi(self, F, iterable):
pool = Pool(size=self.pool_size)
res = pool.imap(F, iterable)
return res
def schedule(self):
ori_p = self.args.ori_p
enc_p = self.args.enc_p
if os.path.isfile(enc_p) and self.n_step == 0:
enc_img = self.step_preprocess(enc_p)
enc_img = self.step_encode(enc_img)
hash_str = self.step_hash(enc_img)
# res_list = self.step_query(hash_str)
# return res_list
return hash_str
else:
if self.n_step >= 1:
try:
if self.args.ori_p != "":
self.step_preprocess_n(ori_p)
except:
raise print("请输入正确的原始图像路径！！！")
if self.args.train == 1:
self.step_train()
if self.n_step >= 2:
imgs = self.multi(cv.imread, [self.args.enc_p + m for m in os.listdir(self.args.enc_p)])
enc_imgs = [i for i in self.step_encode_n(imgs)]
# print(enc_imgs)
# enc_imgs = []
# for i in self.step_encode_n(imgs):
# enc_imgs.append(i)
else:
return None
if self.n_step >=3:
hash_strs = [i for i in self.step_hash_n(enc_imgs)]
# hash_strs = []
# for i in self.step_hash_n(enc_imgs):
# hash_strs.append(i)
else:
return enc_imgs
if self.n_step >= 4:
self.step_index(hash_strs)
else:
return hash_strs
def step_preprocess(self, ori_p):
img = img_cut(ori_p)
img = cv.resize(img, (self.args.img_size, self.args.img_size))
return img
def step_preprocess_n(self, ori_p):
img_migrate(ori_p=ori_p)
img_reshape(self.args)
def step_train(self):
train(self.args)
def step_encode(self, img):
encode = encode_transform(img, self.args)
return encode
def step_encode_n(self, imgs):
return self.multi(self.step_encode, imgs)
def step_hash(self, img):
h = self.hash.ahash(img)
return h
def step_hash_n(self, imgs):
return self.multi(self.step_hash, imgs)
def step_index(self, hash_strs):
self.hash.index()
def step_query(self, hash_str):
self.hash.query()
'''
import cv2 as cv
from torchvision.utils import save_image
from setting import TEST_DIR
from aae import encode_transform
from hash import Hash
from main import args
img_1 = cv.imread(TEST_DIR + "enc/1.jpg")
img_2 = cv.imread(TEST_DIR + "enc/2.jpg")
enc_1 = encode_transform(img_1, args)
enc_2 = encode_transform(img_2, args)
my_hash = Hash()
ah_1 = my_hash.ahash(enc_1)
ah_2 = my_hash.ahash(enc_2)
print("1 and 2:\t\n", ah_1, "\n", ah_2)
a1 = my_hash.compare(ah_1.__str__(), ah_2.__str__())
print("1 and 2 similarity:\t", a1)
img_3 = cv.imread(TEST_DIR + "enc/3.jpg")
img_4 = cv.imread(TEST_DIR + "enc/4.jpg")
enc_3 = encode_transform(img_3, args)
enc_4 = encode_transform(img_4, args)
my_hash = Hash()
ah_3 = my_hash.ahash(enc_3)
ah_4 = my_hash.ahash(enc_4)
print("3 and 4:\t\n", ah_3, "\n", ah_4)
a2 = my_hash.compare(ah_3.__str__(), ah_4.__str__())
