Welding-Parts-Anomaly-Detection/var_detection.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/6/19 0019 下午 12:24
# @Author  : liudan
# @File    : var_detection.py
# @Software: pycharm


import os

import cv2
from PIL import Image
import numpy as np
import random


# 计算两张图片的方差
def calculate_variance(image):
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    variance = np.var(gray_image)
    mean = np.mean(image)
    variance = variance / mean
    return variance


# 计算两张图的方差差异，ref_image为模板图，query_image为查询图
def calculate_variance_difference(ref_image, query_image):
    variance1= calculate_variance(ref_image)
    variance2 = calculate_variance(query_image)
    variance_diff = abs(variance1 - variance2)
    max_variance = max(variance1, variance2)
    normalized_diff = variance_diff / max_variance if max_variance != 0 else 0
    # diff = np.abs(image1.astype(float) - image2.astype(float)).var()
    return normalized_diff



# 计算Recall等，positive_folder正样本文件夹，negative_folder负样本文件夹，ref_image_dir模板图文件夹
def detection(positive_folder, negative_folder, ref_image_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_path = 'H:\\dataset\\region_samples\\positive_test_img\\A1\\image165246-002_A.jpg'
    ref_image = cv2.imread(ref_image_path) # 读取模板图
    print(f"{image_file}\n")

    true_positives = 0
    false_positives = 0
    false_negatives = 0
    true_negatives = 0
    threshold_score = 0.9

    # 在正样本中计算Recall等
    positive_false = []
    for file in os.listdir(positive_folder):
        query_image_path = os.path.join(positive_folder, file) # 查询图，来自正样本
        query_image = cv2.imread(query_image_path)
        variance_score = calculate_variance_difference(ref_image,query_image) # 计算模板图与查询图方差差异

        if variance_score < threshold_score:
            true_positives += 1
        else:
            false_negatives += 1
            positive_false.append(file)

    print(positive_false)

    # 在负样本中计算Recall等
    negative_false =[]
    for file in os.listdir(negative_folder):
        query_image_path = os.path.join(negative_folder, file) # 查询图，来自负样本
        query_image = cv2.imread(query_image_path)
        variance_score = calculate_variance_difference(ref_image,query_image)

        if variance_score > threshold_score:
            true_negatives += 1
        else:
            false_positives += 1
            negative_false.append((file,variance_score))
    print(negative_false)

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

    print(f"Accurary:{Accurary: .2f}")
    print(f"Precision: {Precision:.2f}")
    print(f"Recall: {Recall:.2f}")
    print(f"F1 Score: {F1_score:.2f}")

if __name__ == "__main__":
    positive_folder = './data/624_test/positive/Y'
    negative_folder = './data/624_test/negative/Y'
    ref_image_dir = './data/624_test/positive/Y'

    detection(positive_folder, negative_folder, ref_image_dir)