Welding-Parts-Anomaly-Detection/rotate_env.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/6/24 0024 上午 10:38
# @Author  : liudan
# @File    : rotate_env.py
# @Software: pycharm
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/6/19 0019 下午 4:35
# @Author  : liudan
# @File    : rorate_image_detection.py
# @Software: pycharm
import time

import cv2 as cv
import cv2
import numpy as np
from loguru import logger
import os

from superpoint_superglue_deployment import Matcher
from datetime import datetime
import random

import cv2
import numpy as np
import random


# 在模板图的基础上创造负样本（翻转图片）
def rotate_image(image_path):
    rotate_result_images = []
    for path in os.listdir(image_path):
        ref_image_path = os.path.join(image_path,path)
        ref_image = cv2.imread(ref_image_path, 0)
        if random.random() < 0.5:  # 以0.5的概率决定是否翻转
            rotated_image = cv2.rotate(ref_image, cv2.ROTATE_180)
            rotate_result_images.append((rotated_image, path))
        else:
            print(f"Skipped")
    print(len(rotate_result_images))

    return rotate_result_images


# 配准操作，计算模板图与查询图配准点的距离差异（翻转后的距离大于未翻转的）
def match_image(ref_img, query_img):
    superglue_matcher = Matcher(
        {
            "superpoint": {
                "input_shape": (-1, -1),
                "keypoint_threshold": 0.003,
            },
            "superglue": {
                "match_threshold": 0.5,
            },
            "use_gpu": True,
        }
    )
    query_kpts, ref_kpts, _, _, matches = superglue_matcher.match(ref_img, query_img)
    matched_query_kpts = [query_kpts[m.queryIdx].pt for m in matches]
    matched_ref_kpts = [ref_kpts[m.trainIdx].pt for m in matches]
    diff_query_ref = np.array(matched_ref_kpts) - np.array(matched_query_kpts)
    print(diff_query_ref)
    if len(matches)!=0:
        diff_query_ref = np.linalg.norm(diff_query_ref, axis=1, ord=2)
        diff_query_ref = np.sqrt(diff_query_ref)
        diff_query_ref = np.mean(diff_query_ref)
    else:
        diff_query_ref = np.inf
    return diff_query_ref, matches  # 返回差异值，用于后续做比较





def registration(ref_image_dir, ref_selection_dir):

    # 计算recall等需要的的参数
    true_positives = 0
    false_positives = 0
    false_negatives = 0
    true_negatives = 0

    # 返回预测错误的样本
    positive_false = []
    negative_false = []

    # 循环遍历模板图
    for file_name in os.listdir(ref_selection_dir):
        ref_image_path = os.path.join(ref_selection_dir,file_name)
        ref_image = cv2.imread(ref_image_path, 0)

        save_path = './data/region_rotate/vague/A' # 保存预测错误的原图片，可不用

        # 循环遍历正样本（这里正样本是查询图，因为查询图都是没有翻转的，所以可当作正样本）
        # 这是在正样本中预测
        for file in os.listdir(ref_image_dir):
            query_image_path = os.path.join(ref_image_dir, file)
            query_image = cv2.imread(query_image_path, 0)
            query_image_rotate = cv2.rotate(query_image, cv2.ROTATE_180)
            # cv2.imwrite('1111111111.jpg', query_image_rotate)

            diff1, matches1 = match_image(ref_image, query_image) # 计算模板图与查询图的配准点差异
            diff2, matches2 = match_image(ref_image,query_image_rotate) # 计算模板图与翻转180度图的配准点差异

            # if (len(matches1) > len(matches2)) or (diff1<diff2) :
            if diff1 < diff2:
                flag = True
            else:
                flag = False
            print(flag)

            if flag == True:
                true_positives += 1
            else:
                false_negatives += 1
                positive_false.append((file,file_name))
                cv2.imwrite(os.path.join(save_path, f'{file[:-4]}.jpg'),query_image)


        # 这是在负样本中预测
        # 因为缺少负样本，所以用rotate_image方法创造负样本
        nagetive_image = rotate_image(ref_image_dir)
        for i, item in enumerate(nagetive_image):
            file, path = item
            # query_image = cv2.cvtColor(file,cv2.COLOR_BGR2GRAY)
            query_image = file
            query_image_rotate = cv2.rotate(query_image, cv2.ROTATE_180)

            diff1, matches1 = match_image(ref_image, query_image)
            diff2, matches2 = match_image(ref_image, query_image_rotate)

            # if (len(matches1) > len(matches2)) or (diff1<diff2):
            if diff1 < diff2:
                flag = True
            else:
                flag = False
            print(flag)

            if flag == False:
                true_negatives += 1
            else:
                false_positives += 1
                negative_false.append((path, file_name))
                cv2.imwrite(os.path.join(save_path, f'{path[:-4]}.jpg'),query_image)

    Accurary = (true_positives + true_negatives) / (true_positives + true_negatives + false_positives + false_negatives)
    Precision = true_negatives / (true_negatives + false_negatives)
    Recall = true_negatives / (true_negatives + false_positives)
    F1_score = 2 * (Precision * Recall) / (Precision + Recall)

    print(positive_false)
    print(negative_false)
    # print(file_name)
    print(f"Accurary:{Accurary: .4f}")
    print(f"Precision: {Precision:.4f}")
    print(f"Recall: {Recall:.4f}")
    print(f"F1 Score: {F1_score:.4f}")





if __name__ == "__main__":
    ref_image_path = './data/region_rotate/vague/E2'
    ref_selection_dir = './data/region_rotate/vague/E22'
    # start_time = time.time()
    registration(ref_image_path, ref_selection_dir)
    # end_time = time.time()
    # elapsed_time = end_time - start_time
    # print(f"程序用时: {elapsed_time:.2f} 秒")