Welding-Parts-Anomaly-Detection/count_env.py

#!/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')