Welding-Parts-Anomaly-Detection/rorate_image_detection.py

#!/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:
            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,
        }
    )
    _, _, _, _, matches = superglue_matcher.match(ref_img, query_img)
    return matches





def registration(ref_image_dir, ref_selection_dir):
    # filename = os.listdir(ref_image_dir)
    # image_files = [f for f in filename if f.endswith(('.jpg', '.jpeg', '.png', '.JPG'))]
    # image_file = random.choice(image_files)
    # ref_image_path = os.path.join(ref_image_dir, image_file)
    # ref_image = cv2.imread(ref_image_path, 0)

    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)


        # positive_false = []
        save_path = './data/region_rotate/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)

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

            if len(matches1) > len(matches2):
                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)
        # print(positive_false)
        # print(file_name)

        # negative_false = []
        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)

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

            if len(matches1) > len(matches2):
                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)
        # print(negative_false)
        # print(file_name)

    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/E1'
    ref_selection_dir = './data/region_rotate/E'
    # 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} 秒")