Spaces:
Sleeping
Sleeping
| 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="file") | |
| 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) | |