97 lines
4.5 KiB
Python
97 lines
4.5 KiB
Python
import os
|
|
import shutil
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from skimage.metrics import structural_similarity as ssim
|
|
|
|
def load_images(folder):
|
|
"""Load all PNG images from a folder, categorized by prefix."""
|
|
true_images = {}
|
|
false_images = {}
|
|
for filename in os.listdir(folder):
|
|
if filename.endswith('.png'):
|
|
filepath = os.path.join(folder, filename)
|
|
image = cv2.imread(filepath, cv2.IMREAD_UNCHANGED) # Load image with alpha channel
|
|
if filename.startswith("True_"):
|
|
true_images[filename] = image
|
|
elif filename.startswith("False_"):
|
|
false_images[filename] = image
|
|
return true_images, false_images
|
|
|
|
def resize_image(image, size=(125, 50)):
|
|
"""Resize an image to a standard size."""
|
|
return cv2.resize(image, size, interpolation=cv2.INTER_AREA)
|
|
|
|
def calculate_ssim(image1, image2):
|
|
"""Calculate SSIM between two images."""
|
|
if image1.shape[2] == 4: # Handle alpha channel
|
|
image1 = cv2.cvtColor(image1, cv2.COLOR_BGRA2BGR)
|
|
if image2.shape[2] == 4: # Handle alpha channel
|
|
image2 = cv2.cvtColor(image2, cv2.COLOR_BGRA2BGR)
|
|
gray1 = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
|
|
gray2 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
|
|
score, _ = ssim(gray1, gray2, full=True)
|
|
return score
|
|
|
|
def calculate_color_similarity(image1, image2):
|
|
"""Compare color histograms of two images."""
|
|
if image1.shape[2] == 4: # Handle alpha channel
|
|
image1 = cv2.cvtColor(image1, cv2.COLOR_BGRA2BGR)
|
|
if image2.shape[2] == 4: # Handle alpha channel
|
|
image2 = cv2.cvtColor(image2, cv2.COLOR_BGRA2BGR)
|
|
hist1 = cv2.calcHist([image1], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
|
|
hist2 = cv2.calcHist([image2], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
|
|
hist1 = cv2.normalize(hist1, hist1).flatten()
|
|
hist2 = cv2.normalize(hist2, hist2).flatten()
|
|
return cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)
|
|
|
|
def calculate_combined_similarity(image1, image2):
|
|
"""Combine SSIM and color histogram similarity."""
|
|
# ssim_score = calculate_ssim(image1, image2)
|
|
color_score = calculate_color_similarity(image1, image2)
|
|
# Weight SSIM and color similarity equally (adjust weights as needed)
|
|
return 1 * color_score
|
|
|
|
def find_most_similar_images(true_images, false_images):
|
|
"""Find the most similar image for each 'True_' image among 'False_' images."""
|
|
results = {}
|
|
for true_name, true_image in true_images.items():
|
|
max_similarity = -1
|
|
most_similar_image = None
|
|
for false_name, false_image in false_images.items():
|
|
# resized_true = resize_image(true_image)
|
|
# resized_false = resize_image(false_image)
|
|
|
|
if true_name.split("_")[2] == false_name.split("_")[2]:
|
|
continue
|
|
|
|
similarity = calculate_combined_similarity(true_image, false_image)
|
|
if similarity > max_similarity:
|
|
max_similarity = similarity
|
|
most_similar_image = false_name
|
|
results[true_name] = (most_similar_image, max_similarity)
|
|
os.makedirs("output14\\"+true_name+"_"+most_similar_image,exist_ok=True)
|
|
shutil.copyfile('output13\\'+true_name, 'output14\\'+true_name+"_"+most_similar_image+"\\"+true_name)
|
|
shutil.copyfile('output13\\' + most_similar_image, 'output14\\' +true_name+"_"+most_similar_image+"\\" + most_similar_image)
|
|
|
|
# shutil.copyfile('output12\\'+ "segment_"+true_name.split("_")[2]+".png", 'output14\\'+true_name+"_"+most_similar_image+"\\segment_"+true_name.split("_")[2]+".png")
|
|
# shutil.copyfile('output12\\' +"segment_"+most_similar_image.split("_")[2]+".png", 'output14\\'+true_name+"_"+most_similar_image+"\\segment_"+most_similar_image.split("_")[2]+".png")
|
|
|
|
shutil.copyfile('output10\\'+ "segment_"+true_name.split("_")[2]+".png", 'output14\\'+true_name+"_"+most_similar_image+"\\segment_"+true_name.split("_")[2]+".png")
|
|
shutil.copyfile('output10\\' +"segment_"+most_similar_image.split("_")[2]+".png", 'output14\\'+true_name+"_"+most_similar_image+"\\segment_"+most_similar_image.split("_")[2]+".png")
|
|
|
|
|
|
return results
|
|
|
|
def main():
|
|
folder = 'output13' # Replace with your folder path
|
|
true_images, false_images = load_images(folder)
|
|
results = find_most_similar_images(true_images, false_images)
|
|
for true_image, (most_similar, similarity) in results.items():
|
|
print(f"True Image: {true_image}, Most Similar False Image: {most_similar}, Similarity: {similarity:.4f}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|