Spaces:
Running
Running
File size: 4,445 Bytes
657e67e 7c0f80b 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 cb35992 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 fe8cd4e 0414ba1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import os
import tempfile
import numpy as np
import cv2
import gradio as gr
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.preprocessing import image
from skimage.metrics import structural_similarity as ssim
from PIL import Image
# Disable GPU for TensorFlow
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
class ImageCharacterClassifier:
def __init__(self, similarity_threshold=0.5):
self.model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
self.similarity_threshold = similarity_threshold
def load_and_preprocess_image(self, img):
# Convert image to array and preprocess it
img = img.convert('RGB')
img_array = np.array(img)
img_array = cv2.resize(img_array, (224, 224)) # Ensure correct size
img_array = np.expand_dims(img_array, axis=0)
img_array = preprocess_input(img_array)
return img_array
def extract_features(self, img):
preprocessed_img = self.load_and_preprocess_image(img)
features = self.model.predict(preprocessed_img)
return features
def calculate_ssim(self, img1, img2):
img1_gray = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)
img2_gray = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)
img2_gray = cv2.resize(img2_gray, (img1_gray.shape[1], img1_gray.shape[0]))
return ssim(img1_gray, img2_gray)
def process_images(reference_image, comparison_images, similarity_threshold):
try:
if reference_image is None:
return "Please upload a reference image.", []
if not comparison_images:
return "Please upload comparison images.", []
classifier = ImageCharacterClassifier(similarity_threshold)
# Convert reference image to NumPy array
ref_image = Image.fromarray(reference_image)
ref_features = classifier.extract_features(ref_image)
results = []
html_output = "<h3>Comparison Results:</h3>"
for comp_image in comparison_images:
try:
# Read image file as PIL Image
comp_pil = Image.open(comp_image)
comp_pil = comp_pil.convert("RGB")
# Convert to NumPy format for SSIM
comp_array = np.array(comp_pil)
# Calculate SSIM score
ssim_score = classifier.calculate_ssim(reference_image, comp_array)
# Extract features
comp_features = classifier.extract_features(comp_pil)
max_feature_diff = np.max(np.abs(ref_features - comp_features))
is_similar = max_feature_diff < 6.0
status_text = "SIMILAR" if is_similar else "NOT SIMILAR"
status_color = "green" if is_similar else "red"
html_output += f"<p style='color:{status_color};'>{comp_image.name}: {status_text}</p>"
results.append(comp_array)
except Exception as e:
html_output += f"<p style='color:red;'>Error processing {comp_image.name}: {str(e)}</p>"
return html_output, results
except Exception as e:
return f"<p style='color:red;'>Error: {str(e)}</p>", []
def create_interface():
with gr.Blocks() as interface:
gr.Markdown("# Image Similarity Classifier")
gr.Markdown("Upload a reference image and multiple comparison images.")
with gr.Row():
with gr.Column():
reference_input = gr.Image(label="Reference Image", type="numpy")
comparison_input = gr.Files(label="Comparison Images", type="filepath")
threshold_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, step=0.05, label="Similarity Threshold")
submit_button = gr.Button("Compare Images")
with gr.Column():
output_html = gr.HTML(label="Results")
output_gallery = gr.Gallery(label="Processed Images", columns=3)
submit_button.click(
fn=process_images,
inputs=[reference_input, comparison_input, threshold_slider],
outputs=[output_html, output_gallery]
)
return interface
if __name__ == "__main__":
interface = create_interface()
interface.launch(share=True)
|