def compare(self, features1, features2):
features1 = json.loads(features1)
features2 = json.loads(features2)
sum = 0.0
if "sift" in self.features:
score = self.compare_sift(features1, features2)
if score > 0.2: # 相近
for alg, weight in self.sift_weights.items():
if alg == 'sift':
sum += self.get_sift_score(score) * weight
continue
method = getattr(self, "compare_" + alg)
scores = method(features1[alg], features2[alg])
logger.debug("{} provide score is {}".format(alg, scores))
if isinstance(scores, tuple):
sum += self.multiscore(scores, weight)
else:
sum += scores * weight
return sum
for alg, weight in self.weights.items():
method = getattr(self, "compare_" + alg)
scores = method(features1[alg], features2[alg])
logger.debug("{} provide score is {}".format(alg, scores))
if isinstance(scores, tuple):
sum += self.multiscore(scores, weight)
else:
sum += scores * weight
return sum
def detection_local_extrema2(self,
dog_pyr,
octaves,
intervals,
extrema_img=None):
img_border = self.img_border
thresh = 0.5 * self.dxthreshold / intervals
extrmums = {"unstable": [], "stable": [], "edge": []}
t = time.time()
for o in range(octaves):
layer = dog_pyr[o]
_, rows, cols = layer.shape
squares = self.stride(layer, (3, 3, 3), (1, 1, 1))
max_v = squares.max(axis=(3, 4, 5))
max_v[max_v < thresh] = 0.0
min_v = squares.min(axis=(3, 4, 5))
min_v[min_v > -thresh] = 0.0
idx_i, idx_x, idx_y = np.nonzero((layer[1:-1, 1:-1, 1:-1] == max_v)
| (layer[1:-1, 1:-1,
1:-1] == min_v))
for i, x, y in zip(idx_i, idx_x, idx_y):
if x + 1 < img_border or x + 1 >= rows - img_border: continue
if y + 1 < img_border or y + 1 >= cols - img_border: continue
# 修正极值点，删除不稳定点
extrmum = self.interploation_extremum(dog_pyr, o, i + 1, x + 1,
y + 1)
if extrmum:
# hessian矩阵，排除边缘点
if self.pass_edge_response(dog_pyr, o, i + 1, x + 1,
y + 1):
extrmums['stable'].append(extrmum)
# logger.debug("stable({},{},{},{})".format(o, i, extrmum['dx'], extrmum['dy']))
else:
extrmums['edge'].append(extrmum)
# logger.debug("edge({},{},{},{})".format(o, i, extrmum['dx'], extrmum['dy']))
else:
extrmum = self.create_keypoint(dog_pyr, o, i + 1, x + 1,
y + 1)
extrmums['unstable'].append(extrmum)
# logger.debug("unstable({},{},{},{})".format(o, i, extrmum['dx'], extrmum['dy']))
stables = extrmums['stable']
logger.debug("size={} cost={}".format(len(stables), (time.time() - t)))
if extrema_img:
self.draw_key_points(extrema_img,
stables,
color=(0, 0, 255),
adjust=1)
# draw_key_points(img, extremas['unstable'], color=(0, 255, 0), adjust=1)
# draw_key_points(img, extremas['edge'], color=(255, 0, 0), adjust=1)
cv.imwrite('sample/extremas.png', extrema_img,
[int(cv.IMWRITE_PNG_COMPRESSION), 3])
return extrmums
def detection_local_extrema(self,
dog_pyr,
octaves,
intervals,
extrema_img=None):
"""
检测当地极值点
:return:
"""
img_border = self.img_border
extrmums = {"unstable": [], "stable": [], "edge": []}
thresh = 0.5 * self.dxthreshold / intervals
t = time.time()
for o in range(octaves):
for i in range(1, (intervals + 2) - 1): # 第一层和最后一层极值忽略
img = dog_pyr[o][i]
rows, cols = img.shape
for x in range(img_border, rows - img_border):
for y in range(img_border, cols - img_border):
val = img[x, y]
if math.fabs(val) <= thresh: # 排除阈值过小的点
continue
if not self.is_extremum(dog_pyr, o, i, x, y):
continue
# 修正极值点，删除不稳定点
extrmum = self.interploation_extremum(
dog_pyr, o, i, x, y)
if extrmum:
# hessian矩阵，排除边缘点
if self.pass_edge_response(dog_pyr, o, i, x, y):
extrmums['stable'].append(extrmum)
# logger.debug("stable({},{},{},{})".format(o, i, extrmum['dx'], extrmum['dy']))
else:
extrmums['edge'].append(extrmum)
# logger.debug("edge({},{},{},{})".format(o, i, extrmum['dx'], extrmum['dy']))
else:
extrmum = self.create_keypoint(dog_pyr, o, i, x, y)
extrmums['unstable'].append(extrmum)
# logger.debug("unstable({},{},{},{})".format(o, i, extrmum['dx'], extrmum['dy']))
stables = extrmums['stable']
logger.debug("size={} cost={}".format(len(stables), (time.time() - t)))
if extrema_img:
self.draw_key_points(extrema_img,
stables,
color=(0, 0, 255),
adjust=1)
# draw_key_points(img, extremas['unstable'], color=(0, 255, 0), adjust=1)
# draw_key_points(img, extremas['edge'], color=(255, 0, 0), adjust=1)
cv.imwrite('sample/extremas.png', extrema_img,
[int(cv.IMWRITE_PNG_COMPRESSION), 3])
return extrmums
def gaussian_pyramid(self, img, octaves, intervals, sigma, archive=False):
"""
高斯金字塔
:return:
"""
k = math.pow(2.0, 1.0 / intervals)
invl = intervals + 3
sigmas = np.ones((invl, ))
sigmas[0] = sigma
for i in range(1, invl, 1):
sig_prev = math.pow(k, i - 1) * sigma
sig_total = sig_prev * k
sigmas[i] = math.sqrt(sig_total * sig_total - sig_prev * sig_prev)
gauss_pyr = [None] * octaves
# for i in range(sigmas.shape[0]):
# logger.debug(sigmas[i])
rows, cols = img.shape
for o in range(octaves):
gauss_pyr[o] = np.ones((invl, rows, cols))
for i in range(invl):
if o == 0 and i == 0:
gauss_pyr[o][i] = np.copy(img)
elif i == 0:
# 前一组高斯图像的倒数第三层
# 第一组第一张图(下标为6)的图像是0组下标为3的图像降采样得来
# gauss_pyr[o][i] = down_sample(gauss_pyr[o-1][intervals])
gauss_pyr[o][i] = cv.pyrDown(gauss_pyr[o - 1][intervals],
dstsize=(cols, rows))
else:
gauss_pyr[o][i] = cv.GaussianBlur(gauss_pyr[o][i - 1],
(0, 0), sigmas[i],
sigmas[i])
rows = rows // 2
cols = cols // 2
if archive:
for o in range(octaves):
for i in range(invl):
logger.debug(gauss_pyr[o][i].shape)
cv.imwrite('sample/gauss_{}_{}.png'.format(o, i),
gauss_pyr[o][i],
[int(cv.IMWRITE_PNG_COMPRESSION), 3])
return gauss_pyr
def query_img_to_wm(mysql_client, compartor, pool):
count, result = mysql_client.select_many(
f"""
SELECT
imgNumber,
originalPicX,
originalPicY
FROM
{DB_TABLE}
WHERE waterMark = 0 or waterMark is NULL
""")
logger.debug(f"查询{DB_TABLE}中未添加水印的图片，总共找到{count}条记录")
result_dict = {}
for line in result:
result_dict['imgNumber'] = line[0]
result_dict['originalPicX'] = line[1]
result_dict['originalPicY'] = line[2]
pool.submit(add_watermark,result_dict.copy(),mysql_client, compartor).add_done_callback(upload_filepath_to_cos)
logger.debug("当前线程池队列堆积长度：{}".format(getattr(pool, '_work_queue')._qsize()))
def add_watermark(task_dict,mysql_client, compartor):
imgNumber = task_dict['imgNumber']
response, image = cos.get_img_file(imgNumber)
x = task_dict['originalPicX']
y = task_dict['originalPicY']
assert x >= WATER_MARK_X * EMBEDER_DWT_WINDOW * (2 ** EMBEDER_DWT_TIMES)
assert y >= WATER_MARK_Y * EMBEDER_DWT_WINDOW * (2 ** EMBEDER_DWT_TIMES)
x_multiple = x // (WATER_MARK_X * EMBEDER_DWT_WINDOW)
y_multiple = y // (WATER_MARK_Y * EMBEDER_DWT_WINDOW)
min_multiple = min(x_multiple, y_multiple)
COMPUTE_DWT_TIMES = int(np.log2(min_multiple))
# 如果计算出来小波变换级数比配置文件中的多
# 那么最终的小波变换级数就用计算出来的
# 否则使用配置文件的小波变换级数
if COMPUTE_DWT_TIMES > EMBEDER_DWT_TIMES:
DWT_TIMES = COMPUTE_DWT_TIMES
else:
DWT_TIMES = EMBEDER_DWT_TIMES
seed_wm = int(imgNumber[:4])
seed_dct = int(imgNumber[7:11])
logger.debug(f'seed_wm:{seed_wm}')
logger.debug(f'seed_dct:{seed_dct}')
embedder = DWTWatermarkEmbeder(
seed_wm,seed_dct,EMBEDER_DWT_MOD_1,
block_shape=(EMBEDER_DWT_WINDOW,EMBEDER_DWT_WINDOW),
dwt_deep=DWT_TIMES,
)
# 用当前时间生成水印
now = arrow.now()
now_text = now.format()
today = now.format("YYYY-MM-DD")
wm,md5 = embedder.generate_wm(now_text)
if not os.path.exists(f'tmp/{today}'):
os.mkdir(f'tmp/{today}')
output_file_path = os.path.join('tmp',today,f'{imgNumber}.png')
output_wm_path = os.path.join('tmp',today,f'wm-{imgNumber}.png')
# 嵌入水印
embed_image = embedder.embed(image,wm_img=wm,filename=output_file_path)
# 生成水印hash
img_hash = generate_img_hash(embed_image)
# 水印写盘
cv.imwrite(output_wm_path,wm)
# 计算嵌入水印后图片的特征值
json_str = compartor.extract(embed_image)
# 更新数据库
mysql_client.exexcute_update(
f"""
UPDATE
`{DB_TABLE}`
SET
random_seed_1 = {seed_wm},
random_seed_2 = {seed_dct},
mod_1 = {EMBEDER_DWT_MOD_1},
dwt_times = {DWT_TIMES},
features = '{json_str}'
WHERE `imgNumber`='{imgNumber}'
"""
)
return output_file_path,mysql_client,md5,img_hash,imgNumber
def main():
if args.mode == 'run':
try:
logger.info("当前模式不是debug模式，启动调度器")
compartor = Compartor()
mysql_client = MysqlClient()
pool = ThreadPoolExecutor(max_workers=THREAD_POOL_MAX_WORKER, thread_name_prefix=THREAD_NAME_PREFIX)
scheduler = BlockingScheduler()
scheduler.add_job(query_img_to_wm, 'interval', minutes=SD_QUERY_INTERVAL,args=[mysql_client, compartor, pool,])
scheduler.start()
except (KeyboardInterrupt, SystemExit):
pool.shutdown(wait=True)
scheduler.shutdown(wait=True)
logger.debug("程序异常退出")
if args.mode == 'debug':
# -------------------------------------------------------------------------------------------------------------- #
if args.gw:
# TODO ： 调整水印大小
if args.name and args.id:
# python eye.py debug -g --name 差不多 --id 312391432483842
logger.info(f"开始测试生成水印,要生成水印的姓名:[{args.name}] ,身份证号:[{args.id}]")
wm = DWTWatermarkEmbeder.generate_wm(args.name,args.id)
if args.output_path:
if os.path.isdir(args.output_path):
output_path = os.path.join(args.output_path,'debug-generate-watermark.png')
if os.path.exists(output_path):
os.remove(output_path)
cv.imwrite(output_path,wm)
else:
cv.imwrite(output_path, wm)
else:
raise Exception("参数--output_path必须是一个存在的文件夹!")
else:
output_path = './debug-output'
if not os.path.exists(output_path):
os.mkdir(output_path)
output_path = os.path.join(output_path,'debug-generate-watermark.png')
if os.path.exists(output_path):
os.remove(output_path)
cv.imwrite(output_path, wm)
cv.imshow('x',wm)
cv.waitKey(0)
cv.destroyAllWindows()
logger.info("水印生成完毕。")
# -------------------------------------------------------------------------------------------------------------- #
if args.embed:
logger.debug("开始测试对图片嵌入水印..")
if not args.original_file_path:
raise Exception("请在--original_file_path参数中提供要嵌入水印的原图路径！")
else:
if not os.path.isfile(args.original_file_path):
raise Exception("--original_file_path参数必须指向某个文件!")
else:
if not args.original_file_path.endswith(".jpg") and not args.original_file_path.endswith(".png"):
logger.warn("--original_file_path参数指向的文件不是.jpg或者.png文件！")
if not args.watermark_file_path:
raise Exception("请在--watermark_file_path参数中提供要嵌入的水印图片路径！")
else:
if not os.path.isfile(args.watermark_file_path):
raise Exception("--watermark_file_path参数必须指向某个文件!")
else:
if not args.watermark_file_path.endswith(".jpg") and not args.watermark_file_path.endswith(".png"):
logger.warn("--watermark_file_path参数指向的文件不是.jpg或者.png文件！")
original_image = cv.imread(args.original_file_path)
water_mark = cv.imread(args.watermark_file_path)
assert original_image.shape[1] >= water_mark.shape[1] * EMBEDER_DWT_WINDOW * (2 ** EMBEDER_DWT_TIMES)
assert original_image.shape[0] >= water_mark.shape[0] * EMBEDER_DWT_WINDOW * (2 ** EMBEDER_DWT_TIMES)
x_multiple = original_image.shape[1] // (water_mark.shape[1] * EMBEDER_DWT_WINDOW)
y_multiple = original_image.shape[0] // (water_mark.shape[0] * EMBEDER_DWT_WINDOW)
min_multiple = min(x_multiple,y_multiple)
COMPUTE_DWT_TIMES = int(np.log2(min_multiple))
# 如果计算出来小波变换级数比配置文件中的多
# 那么最终的小波变换级数就用计算出来的
# 否则使用配置文件的小波变换级数
if COMPUTE_DWT_TIMES > EMBEDER_DWT_TIMES:
DWT_TIMES = COMPUTE_DWT_TIMES
else:
DWT_TIMES = EMBEDER_DWT_TIMES
embedder = DWTWatermarkEmbeder(
4399, 2333, 64,
dwt_deep=DWT_TIMES
)
embedded_image = embedder.embed(ori_img=original_image, wm_img=water_mark)
if args.output_path:
if os.path.isdir(args.output_path):
output_path = os.path.join(args.output_path, 'debug-embed-watermark.png')
if os.path.exists(output_path):
os.remove(output_path)
cv.imwrite(output_path, embedded_image)
else:
cv.imwrite(output_path, embedded_image)
else:
raise Exception("参数--output_path必须是一个存在的文件夹!")
else:
output_path = './debug-output'
if not os.path.exists(output_path):
os.mkdir(output_path)
output_path = os.path.join(output_path, 'debug-embed-watermark.png')
if os.path.exists(output_path):
os.remove(output_path)
cv.imwrite(output_path, embedded_image)
# cv.imshow('x',wm)
# cv.waitKey(0)
# cv.destroyAllWindows()
# -------------------------------------------------------------------------------------------------------------- #
if args.extract:
logger.debug("开始测试提取图片中的水印..")
if not args.extract_file_path:
raise Exception("请在--extract_file_path参数中提供要嵌入水印的原图路径！")
else:
if not os.path.isfile(args.extract_file_path):
raise Exception("--extract_file_path参数必须指向某个文件!")
else:
if not args.extract_file_path.endswith(".jpg") and not args.extract_file_path.endswith(".png"):
logger.warn("--extract_file_path参数指向的文件不是.jpg或者.png文件！")
extract_file = cv.imread(args.extract_file_path)
assert extract_file.shape[1] >= WATER_MARK_X * EMBEDER_DWT_WINDOW * (2 ** EMBEDER_DWT_TIMES)
assert extract_file.shape[0] >= WATER_MARK_Y * EMBEDER_DWT_WINDOW * (2 ** EMBEDER_DWT_TIMES)
x_multiple = extract_file.shape[1] // (WATER_MARK_X * EMBEDER_DWT_WINDOW)
y_multiple = extract_file.shape[0] // (WATER_MARK_Y * EMBEDER_DWT_WINDOW)
min_multiple = min(x_multiple, y_multiple)
COMPUTE_DWT_TIMES = int(np.log2(min_multiple))
# 如果计算出来小波变换级数比配置文件中的多
# 那么最终的小波变换级数就用计算出来的
# 否则使用配置文件的小波变换级数
if COMPUTE_DWT_TIMES > EMBEDER_DWT_TIMES:
DWT_TIMES = COMPUTE_DWT_TIMES
else:
DWT_TIMES = EMBEDER_DWT_TIMES
extractor = DwtWatermarkExractor(
4399,2333,64,
wm_shape=(WATER_MARK_Y,WATER_MARK_X),
dwt_deep=DWT_TIMES
)
wm = extractor.extract(extract_file)
if args.output_path:
if os.path.isdir(args.output_path):
output_path = os.path.join(args.output_path, 'debug-extract-watermark.png')
if os.path.exists(output_path):
os.remove(output_path)
cv.imwrite(output_path, wm)
else:
cv.imwrite(output_path, wm)
else:
raise Exception("参数--output_path必须是一个存在的文件夹!")
else:
output_path = './debug-output'
if not os.path.exists(output_path):
os.mkdir(output_path)
output_path = os.path.join(output_path, 'debug-extract-watermark.png')
if os.path.exists(output_path):
os.remove(output_path)
cv.imwrite(output_path, wm)
def extract(self, filename, output_path=None):
if not isinstance(filename, str):
embed_img = filename.astype(np.float32)
else:
embed_img = cv2.imread(filename).astype(np.float32)
if self.color_mod == 'RGB':
embed_img_YUV = embed_img
elif self.color_mod == 'YUV':
embed_img_YUV = cv2.cvtColor(embed_img, cv2.COLOR_BGR2YUV)
if not embed_img_YUV.shape[0] % (2**self.dwt_deep) == 0:
temp = (
2**self.dwt_deep) - embed_img_YUV.shape[0] % (2**self.dwt_deep)
embed_img_YUV = np.concatenate(
(embed_img_YUV, np.zeros((temp, embed_img_YUV.shape[1], 3))),
axis=0)
if not embed_img_YUV.shape[1] % (2**self.dwt_deep) == 0:
temp = (
2**self.dwt_deep) - embed_img_YUV.shape[1] % (2**self.dwt_deep)
embed_img_YUV = np.concatenate(
(embed_img_YUV, np.zeros((embed_img_YUV.shape[0], temp, 3))),
axis=1)
assert embed_img_YUV.shape[0] % (2**self.dwt_deep) == 0
assert embed_img_YUV.shape[1] % (2**self.dwt_deep) == 0
logger.info("读取图片成功！要抽取水印的图片大小为 ： {0}".format(embed_img_YUV.shape))
logger.info("开始进行{}次小波变换".format(self.dwt_deep))
embed_img_Y = embed_img_YUV[:, :, 0]
embed_img_U = embed_img_YUV[:, :, 1]
embed_img_V = embed_img_YUV[:, :, 2]
coeffs_Y = dwt2(embed_img_Y, 'haar')
coeffs_U = dwt2(embed_img_U, 'haar')
coeffs_V = dwt2(embed_img_V, 'haar')
ha_Y = coeffs_Y[0]
ha_U = coeffs_U[0]
ha_V = coeffs_V[0]
# 对ha进一步进行小波变换,并把下一级ha保存到ha中
for i in range(self.dwt_deep - 1):
coeffs_Y = dwt2(ha_Y, 'haar')
ha_Y = coeffs_Y[0]
coeffs_U = dwt2(ha_U, 'haar')
ha_U = coeffs_U[0]
coeffs_V = dwt2(ha_V, 'haar')
ha_V = coeffs_V[0]
# 初始化块索引数组
try:
if self.ha_Y.shape == ha_Y.shape:
self.init_block_add_index(ha_Y.shape)
else:
logger.warn('你现在要解水印的图片与之前读取的原图的形状不同,这是不被允许的')
except:
self.init_block_add_index(ha_Y.shape)
ha_block_shape = (int(ha_Y.shape[0] / self.block_shape[0]),
int(ha_Y.shape[1] / self.block_shape[1]),
self.block_shape[0], self.block_shape[1])
strides = ha_Y.itemsize * (np.array([
ha_Y.shape[1] * self.block_shape[0], self.block_shape[1],
ha_Y.shape[1], 1
]))
ha_Y_block = np.lib.stride_tricks.as_strided(ha_Y.copy(),
ha_block_shape, strides)
ha_U_block = np.lib.stride_tricks.as_strided(ha_U.copy(),
ha_block_shape, strides)
ha_V_block = np.lib.stride_tricks.as_strided(ha_V.copy(),
ha_block_shape, strides)
extract_wm = np.array([])
extract_wm_Y = np.array([])
extract_wm_U = np.array([])
extract_wm_V = np.array([])
self.random_dct = np.random.RandomState(self.random_seed_dct)
start = time.time()
logger.info("小波变换完成，开始提取水印")
index = np.arange(self.block_shape[0] * self.block_shape[1])
for i in range(self.length):
if i % 5000 == 0:
logger.debug(f"正在提取第{i+1}个窗口的像素点")
self.random_dct.shuffle(index)
wm_Y = self.block_get_wm(
ha_Y_block[self.block_add_index0[i], self.block_add_index1[i]],
index)
wm_U = self.block_get_wm(
ha_U_block[self.block_add_index0[i], self.block_add_index1[i]],
index)
wm_V = self.block_get_wm(
ha_V_block[self.block_add_index0[i], self.block_add_index1[i]],
index)
wm = round((wm_Y + wm_U + wm_V) / 3)
# else情况是对循环嵌入的水印的提取
if i < self.wm_shape[0] * self.wm_shape[1]:
extract_wm = np.append(extract_wm, wm)
extract_wm_Y = np.append(extract_wm_Y, wm_Y)
extract_wm_U = np.append(extract_wm_U, wm_U)
extract_wm_V = np.append(extract_wm_V, wm_V)
else:
times = int(i / (self.wm_shape[0] * self.wm_shape[1]))
ii = i % (self.wm_shape[0] * self.wm_shape[1])
extract_wm[ii] = (extract_wm[ii] * times + wm) / (times + 1)
extract_wm_Y[ii] = (extract_wm_Y[ii] * times + wm_Y) / (times +
1)
extract_wm_U[ii] = (extract_wm_U[ii] * times + wm_U) / (times +
1)
extract_wm_V[ii] = (extract_wm_V[ii] * times + wm_V) / (times +
1)
wm_index = np.arange(extract_wm.size)
self.random_wm = np.random.RandomState(self.random_seed_wm)
self.random_wm.shuffle(wm_index)
extract_wm[wm_index] = extract_wm.copy()
extract_wm_Y[wm_index] = extract_wm_Y.copy()
extract_wm_U[wm_index] = extract_wm_U.copy()
extract_wm_V[wm_index] = extract_wm_V.copy()
end = time.time()
logger.info(f"水印提取完成，总共耗时 ： {end-start}")
if output_path:
cv2.imwrite(output_path,
extract_wm.reshape(self.wm_shape[0], self.wm_shape[1]))
return extract_wm.reshape(self.wm_shape[0], self.wm_shape[1])
show("img2", img2)
eyer = Compartor()
# hist1 = eyer.extract_grad_hist(img1)
# hist2 = eyer.extract_grad_hist(img2)
#
# hist1 = drawHist(hist1, 180)
# hist2 = drawHist(hist2, 180)
#
# plt_show((1, 2), ["original", "rotate"], [hist1, hist2])
features1 = eyer.extract(img1)
features2 = eyer.extract(img2)
score = eyer.compare(features1, features2)
logger.debug("score=", score)
# features1 = eyer.extract_sift(img1)
# features2 = eyer.extract_sift(img2)
# logger.debug("sift score=", eyer.compare_sift(features1, features2))
# hsv1 = eyer.extract_hsv(img1)
# hsv2 = eyer.extract_hsv(img2)
# logger.debug("hsv score=", eyer.compare_hsv(hsv1, hsv2))
# img1 = eyer.extract_contours(img1)
# img2 = eyer.extract_contours(img2)
# score = eyer.compare_contours(img1, img2)
# logger.debug("contours score=", score)
# show("img1", img1)
# show("img2", img2)
# show("img3", img3)
def extract(self, img, extrema_img=None):
'''
sift 算法
:return:
'''
sigma = self.sigma
intervals = self.intervals
bias = self.bias
# 创建初始灰度图像
# 初始图像先将原图像灰度化，再扩大一倍后，使用高斯模糊平滑
init_gray = self.create_init_smooth_gray(img, sigma)
# 计算组数
rows, cols = init_gray.shape
octaves = math.floor(math.log2(min(rows, cols)) - bias)
logger.debug(f"Target Info: size=({rows},{cols}), octaves={octaves}")
# 高斯金字塔
t = time.time()
gauss_pyr = self.gaussian_pyramid(init_gray,
octaves,
intervals,
sigma,
archive=self.archive)
logger.debug("生成高斯金字塔耗时: {}".format(time.time() - t))
# 差分金字塔
t = time.time()
dog_pyr = self.dog_pyramid(gauss_pyr,
octaves,
intervals,
archive=self.archive)
logger.debug("生成差分金字塔耗时: {}".format(time.time() - t))
# 检测当地极值点
t = time.time()
extremas = self.detection_local_extrema2(dog_pyr,
octaves,
intervals,
extrema_img=extrema_img)
stables = extremas['stable']
logger.debug("检测当地极值点耗时: {}".format(time.time() - t))
# 计算尺度
self.calculate_scale(stables, sigma, intervals)
# 关键点方向分配
t = time.time()
features = self.orientation_assignment(gauss_pyr, stables)
logger.debug("关键点方向分配耗时: {}".format(time.time() - t))
# 计算描述符
t = time.time()
self.descriptor_representation(gauss_pyr, features, DESCR_HIST_BINS,
DESCR_WINDOW_WIDTH)
logger.debug("计算描述符分配耗时: {}".format(time.time() - t))
# 剔除无用信息
w, h = img.shape[0:2]
labels = self.get_labels(features, w, h, octaves, intervals, sigma)
return labels
#
# datas, labels = get_labels(features1)
# kd1 = kd_tree.KDTree(datas, labels)
# logger.debug(kd1.length)
img1 = cv.imread("img/1.jpg", cv.IMREAD_COLOR)
sift = Sift()
features1 = sift.extract(img1)
sift.draw_sift_features(img1, features1)
cv.imwrite('sample/features1.png', img1,
[int(cv.IMWRITE_PNG_COMPRESSION), 3])
img2 = cv.imread("img/2.jpg", cv.IMREAD_COLOR)
features2 = sift.extract(img2)
sift.draw_sift_features(img2, features2)
cv.imwrite('sample/features2.png', img2,
[int(cv.IMWRITE_PNG_COMPRESSION), 3])
logger.debug("Match===================>")
score1 = sift.match(features1, features2, k=2)
score2 = sift.match(features2, features1, k=2)
score3 = sift.bimatch(features1, features2, k=2)
logger.debug("({}, {}, {})".format(score1, score2, score3))
score1 = sift.match(features1, features2, k=5)
score2 = sift.match(features2, features1, k=5)
score3 = sift.bimatch(features1, features2, k=5)
logger.debug("({}, {}, {})".format(score1, score2, score3))
def embed(self,ori_img,wm_img,filename=None):
self.read_ori_img(ori_img)
self.read_wm(wm_img)
embed_ha_Y_block = self.ha_Y_block.copy()
embed_ha_U_block = self.ha_U_block.copy()
embed_ha_V_block = self.ha_V_block.copy()
self.random_dct = np.random.RandomState(self.random_seed_dct)
index = np.arange(self.block_shape[0] * self.block_shape[1])
start = time.time()
logger.info("开始遍历坐标矩阵中的每个点，一共有{}个window".format(self.length))
for i in range(self.length):
self.random_dct.shuffle(index)
embed_ha_Y_block[self.block_add_index0[i], self.block_add_index1[i]] = self.block_add_wm(
embed_ha_Y_block[self.block_add_index0[i], self.block_add_index1[i]], index, i)
embed_ha_U_block[self.block_add_index0[i], self.block_add_index1[i]] = self.block_add_wm(
embed_ha_U_block[self.block_add_index0[i], self.block_add_index1[i]], index, i)
embed_ha_V_block[self.block_add_index0[i], self.block_add_index1[i]] = self.block_add_wm(
embed_ha_V_block[self.block_add_index0[i], self.block_add_index1[i]], index, i)
if i % 5000 == 0:
logger.debug(f"完成对第{(i+1)}个window,YUV三个分量的嵌入")
end = time.time()
logger.info(f"水印嵌入完成。总共耗时 ： {(end-start)}")
start = time.time()
logger.info("开始合并YUV通道,准备反向小波变换")
embed_ha_Y_part = np.concatenate(embed_ha_Y_block, 1)
embed_ha_Y_part = np.concatenate(embed_ha_Y_part, 1)
embed_ha_U_part = np.concatenate(embed_ha_U_block, 1)
embed_ha_U_part = np.concatenate(embed_ha_U_part, 1)
embed_ha_V_part = np.concatenate(embed_ha_V_block, 1)
embed_ha_V_part = np.concatenate(embed_ha_V_part, 1)
embed_ha_Y = self.ha_Y.copy()
embed_ha_Y[:self.part_shape[0], :self.part_shape[1]] = embed_ha_Y_part
embed_ha_U = self.ha_U.copy()
embed_ha_U[:self.part_shape[0], :self.part_shape[1]] = embed_ha_U_part
embed_ha_V = self.ha_V.copy()
embed_ha_V[:self.part_shape[0], :self.part_shape[1]] = embed_ha_V_part
for i in range(self.dwt_deep):
(cH, cV, cD) = self.coeffs_Y[-1 * (i + 1)]
embed_ha_Y = idwt2((embed_ha_Y.copy(), (cH, cV, cD)), "haar") # 其idwt得到父级的ha
(cH, cV, cD) = self.coeffs_U[-1 * (i + 1)]
embed_ha_U = idwt2((embed_ha_U.copy(), (cH, cV, cD)), "haar") # 其idwt得到父级的ha
(cH, cV, cD) = self.coeffs_V[-1 * (i + 1)]
embed_ha_V = idwt2((embed_ha_V.copy(), (cH, cV, cD)), "haar") # 其idwt得到父级的ha
# 最上级的ha就是嵌入水印的图,即for运行完的ha
embed_img_YUV = np.zeros(self.ori_img_YUV.shape, dtype=np.float32)
embed_img_YUV[:, :, 0] = embed_ha_Y
embed_img_YUV[:, :, 1] = embed_ha_U
embed_img_YUV[:, :, 2] = embed_ha_V
embed_img_YUV = embed_img_YUV[:self.ori_img_shape[0], :self.ori_img_shape[1]]
if self.color_mod == 'RGB':
embed_img = embed_img_YUV
elif self.color_mod == 'YUV':
embed_img = cv2.cvtColor(embed_img_YUV, cv2.COLOR_YUV2BGR)
end = time.time()
logger.info(f"反向小波变换完成,耗时 ： {(end-start)}")
embed_img[embed_img > 255] = 255
embed_img[embed_img < 0] = 0
if filename:
cv2.imwrite(filename, embed_img)
return embed_img