|
@@ -0,0 +1,135 @@
|
|
|
+#!/usr/bin/env python
|
|
|
+# -*- coding: utf-8 -*-
|
|
|
+# @Time : 2024/6/13 0013 下午 12:03
|
|
|
+# @Author : liudan
|
|
|
+# @File : count_env.py
|
|
|
+# @Software: pycharm
|
|
|
+import cv2
|
|
|
+import os
|
|
|
+from PIL import Image, ImageDraw
|
|
|
+import numpy as np
|
|
|
+
|
|
|
+from skimage.metrics import structural_similarity as compare_ssim
|
|
|
+import json
|
|
|
+import random
|
|
|
+import yaml
|
|
|
+import demo_env
|
|
|
+from demo_env import registration_demo
|
|
|
+
|
|
|
+
|
|
|
+def compare_boxes_similarity(image1_path, image2_path, json_file_path, similarity_threshold=0.4):
|
|
|
+ try:
|
|
|
+ if not os.path.exists(image1_path):
|
|
|
+ raise FileNotFoundError(f"Image file {image1_path} not found")
|
|
|
+
|
|
|
+
|
|
|
+ image1 = Image.open(image1_path)# 原图尺寸,未resize
|
|
|
+ image2 = wrap_image
|
|
|
+
|
|
|
+ draw1 = ImageDraw.Draw(image1)
|
|
|
+
|
|
|
+ # 存储相似度结果和是否相同的判断
|
|
|
+ similarity_results = []
|
|
|
+ same_content_boxes = []
|
|
|
+
|
|
|
+
|
|
|
+ with open(json_file_path, 'r') as f:
|
|
|
+ data = json.load(f)
|
|
|
+
|
|
|
+
|
|
|
+ for shape in data['shapes']:
|
|
|
+ if 'points' in shape:
|
|
|
+ shape['points'] = [[int(round(x)), int(round(y))] for x, y in shape['points']]
|
|
|
+ x1, y1 = shape['points'][0]
|
|
|
+ x2, y2 = shape['points'][1]
|
|
|
+
|
|
|
+ # 从两幅图像中截取对应区域
|
|
|
+ region1 = image1.crop((x1, y1, x2, y2))
|
|
|
+ # region1 = image1.crop((x1, y1, x2, y2)).convert('L')
|
|
|
+ draw1.rectangle([x1, y1, x2, y2], outline='red', width=2)
|
|
|
+ image1.save(os.path.join(params['save_dir'], f'save_annotated1_{i}.jpg'))
|
|
|
+ region1.save(os.path.join(params['save_dir'], f'111111{i}.jpg'))
|
|
|
+
|
|
|
+
|
|
|
+ # region2 = image2.crop((left-80, top, right-80, bottom))
|
|
|
+ region2 = image2[y1:y2, x1:x2]
|
|
|
+ # region2 = cv2.cvtColor(region2, cv2.COLOR_BGR2GRAY)
|
|
|
+ # region2= region2.transpose(Image.FLIP_TOP_BOTTOM) #旋转180°针对pillowImage对象
|
|
|
+ # region2 = cv2.rotate(region2, cv2.ROTATE_180)
|
|
|
+
|
|
|
+ filename = f'json_image_{shape["label"]}_{i}.jpg'
|
|
|
+ cv2.imwrite(os.path.join(params['save_dir'], filename), region2)
|
|
|
+ cv2.rectangle(image2, (x1, y1), (x2, y2),(0,255,0), 2)
|
|
|
+ cv2.imwrite(os.path.join(params['save_dir'], f'save_annotated2_{i}.jpg'), image2)
|
|
|
+
|
|
|
+ # 将PIL图像转换为numpy数组,以便进行计算
|
|
|
+ arr1 = np.array(region1)
|
|
|
+ arr2 = region2 # region2一直是numpy数组,所以上述image1和image2处理方式不同
|
|
|
+
|
|
|
+ # 确保两个数组的形状是相同的
|
|
|
+ assert arr1.shape == arr2.shape, "Images do not have the same size for the given box"
|
|
|
+
|
|
|
+ # 使用SSIM计算相似度(范围在-1到1之间,1表示完全相似)
|
|
|
+ # ssim = compare_ssim(arr1, arr2, multichannel=False) # 这是旧版,可以计算灰度图相似度,对于计算彩色图像即使设置multichannel=True也错
|
|
|
+ ssim = compare_ssim(arr1, arr2, channel_axis=2)
|
|
|
+
|
|
|
+
|
|
|
+ similarity_results.append(ssim)
|
|
|
+
|
|
|
+ if ssim > similarity_threshold:
|
|
|
+ same_content_boxes.append(shape)
|
|
|
+ cv2.rectangle(image2, (x1, y1), (x2, y2),(0,255,0), 2)
|
|
|
+
|
|
|
+ text = "Similarity: " + str(round(ssim, 3))
|
|
|
+ text_pos = (x1, y1 - 5)
|
|
|
+ # 参数:图像, 文本, 文本位置, 字体类型, 字体大小, 字体颜色, 字体粗细
|
|
|
+ cv2.putText(image2, text, text_pos, cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 0), 2)
|
|
|
+ cv2.imwrite(os.path.join(params['visualization_dir'],f'{wrap_images_name[:-8]}_{i}.jpg'), image2)
|
|
|
+
|
|
|
+ else:
|
|
|
+ cv2.rectangle(image2, (x1, y1), (x2, y2), (0, 0, 255), 2)
|
|
|
+ text = "score: " + str(round(ssim, 3))
|
|
|
+ text_pos = (x1, y1 - 5)
|
|
|
+ # 参数:图像, 文本, 文本位置, 字体类型, 字体大小, 字体颜色, 字体粗细
|
|
|
+ cv2.putText(image2, text, text_pos, cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 0, 255), 2)
|
|
|
+ cv2.imwrite(os.path.join(params['visualization_dir'], f'{wrap_images_name[:-8]}_{i}.jpg'), image2)
|
|
|
+
|
|
|
+ return similarity_results, same_content_boxes
|
|
|
+
|
|
|
+ except FileNotFoundError as e:
|
|
|
+ print(f"An error occurred: {e}")
|
|
|
+ except Exception as e:
|
|
|
+ print(f"An unexpected error occurred: {e}")
|
|
|
+ return None, None
|
|
|
+
|
|
|
+
|
|
|
+def read_params_from_yml(yml_file_path):
|
|
|
+ with open(yml_file_path, 'r') as file:
|
|
|
+ params = yaml.safe_load(file)
|
|
|
+ return params
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ yml_file_path = 'params.yml'
|
|
|
+ params = read_params_from_yml(yml_file_path)
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+ wrap_images_all = registration_demo(params['image_dir'],params['demo_image_path'], params['json_ref_path'], params['ref_image_path'])
|
|
|
+ for i, item in enumerate(wrap_images_all):
|
|
|
+
|
|
|
+ wrap_image,wrap_images_name = item
|
|
|
+ similarity_results, same_content_boxes = compare_boxes_similarity(params['path_to_image1'], wrap_image, params['json_file_path'],
|
|
|
+ params['similarity_threshold'])
|
|
|
+ # 打印所有坐标框的相似度结果
|
|
|
+ print(f"{wrap_images_name}\n")
|
|
|
+ for idx, score in enumerate(similarity_results, 1):
|
|
|
+ print(f"Similarity Score for Box {idx}: {score}")
|
|
|
+
|
|
|
+ # 打印被认为是相同内容的坐标框
|
|
|
+ print("Boxes with the same content:")
|
|
|
+ for shape in same_content_boxes:
|
|
|
+ print(shape['label'] + ' object is same as template')
|