72 lines
3.1 KiB
Python
72 lines
3.1 KiB
Python
from PIL import Image, ImageDraw, ImageFont
|
|
from tools.getHeatMap import cal_cam
|
|
import os
|
|
|
|
|
|
def merge_imgs(img1_path, img2_path, conf, similar=None, label=None, cam=None, save_path=None):
|
|
save = True
|
|
position = (50, 50) # 文字的左上角坐标
|
|
color = (255, 0, 0) # 红色文字，格式为 RGB
|
|
# if not os.path.exists(os.sep.join([save_path, str(label)])):
|
|
# os.makedirs(os.sep.join([save_path, str(label)]))
|
|
# save_path = os.sep.join([save_path, str(label)])
|
|
# img_name = os.path.basename(img1_path).split('.')[0] + '_' + os.path.basename(img2_path).split('.')[0] + '.png'
|
|
if save_path is None:
|
|
save_path = conf['data']['image_joint_pth']
|
|
if not conf['heatmap']['show_heatmap']:
|
|
img1 = Image.open(img1_path)
|
|
img2 = Image.open(img2_path)
|
|
img1 = img1.resize((224, 224))
|
|
img2 = img2.resize((224, 224))
|
|
new_img = Image.new('RGB', (img1.width + img2.width + 10, img1.height))
|
|
# save_path = conf['data']['image_joint_pth']
|
|
else:
|
|
assert cam is not None, 'cam is None'
|
|
img1 = cam.get_hot_map(img1_path)
|
|
img2 = cam.get_hot_map(img2_path)
|
|
img1_ori = Image.open(img1_path)
|
|
img2_ori = Image.open(img2_path)
|
|
img1_ori = img1_ori.resize((224, 224))
|
|
img2_ori = img2_ori.resize((224, 224))
|
|
new_img = Image.new('RGB',
|
|
(img1.width + img2.width + 10,
|
|
img1.height + img2.width + 10))
|
|
# save_path = conf['heatmap']['image_joint_pth']
|
|
# print('img1_path', img1)
|
|
# print('img2_path', img2)
|
|
if not os.path.exists(os.sep.join([save_path, str(label)])) and (label is not None):
|
|
os.makedirs(os.sep.join([save_path, str(label)]))
|
|
save_path = os.sep.join([save_path, str(label)])
|
|
if save_path is None:
|
|
# save_path = os.sep.join([save_path, str(label)])
|
|
pass
|
|
# img_name = os.path.basename(img1_path).split('.')[0] + '_' + os.path.basename(img2_path).split('.')[0] + '.png'
|
|
img_name = os.path.basename(img1_path).split('.')[0][:30] + '_' + os.path.basename(img2_path).split('.')[0][
|
|
:30] + '.png'
|
|
assert img1.height == img2.height
|
|
|
|
|
|
|
|
# print('new_img', new_img)
|
|
if not conf['heatmap']['show_heatmap']:
|
|
new_img.paste(img1, (0, 0))
|
|
new_img.paste(img2, (img1.width + 10, 0))
|
|
else:
|
|
new_img.paste(img1_ori, (10, 10))
|
|
new_img.paste(img2_ori, (img2_ori.width + 20, 10))
|
|
new_img.paste(img1, (10, img1.height+20))
|
|
new_img.paste(img2, (img2.width+20, img2.height+20))
|
|
|
|
if similar is not None:
|
|
if label == '1' and (similar > 0.5 or similar < 0.25):
|
|
save = False
|
|
elif label == '0' and similar > 0.25:
|
|
save = False
|
|
similar = str(similar) + '_' + str(label)
|
|
draw = ImageDraw.Draw(new_img)
|
|
draw.text(position, str(similar), color, font_size=36)
|
|
os.makedirs(save_path, exist_ok=True)
|
|
img_save = os.path.join(save_path, img_name)
|
|
if save:
|
|
new_img.save(img_save)
|