Image / app.py
omar87's picture
Update app.py
039c558
raw
history blame
1.91 kB
import numpy as np
from skimage.metrics import structural_similarity as ssim
import gradio as gr
import cv2
import os
# Function to calculate SSIM between two images
def calculate_similarity(img1, img2):
if len(img1.shape) == 2:
img1 = cv2.cvtColor(img1, cv2.COLOR_GRAY2RGB)
if len(img2.shape) == 2:
img2 = cv2.cvtColor(img2, cv2.COLOR_GRAY2RGB)
return ssim(img1, img2, win_size=3)
# Function to compute similarity scores for all images
def compute_similarity(target_image, image_list):
scores = []
target_image_resized = cv2.resize(target_image, (target_image.shape[1], target_image.shape[0]))
for image_path in image_list:
image = cv2.imread(image_path)
image_resized = cv2.resize(image, (target_image.shape[1], target_image.shape[0]))
similarity_score = calculate_similarity(target_image_resized, image_resized)
scores.append(similarity_score)
return scores
# Function to handle the Gradio interface
def image_similarity(target_image, image_folder):
target_image = target_image.astype(np.uint8)
image_paths = [os.path.join(image_folder, file) for file in os.listdir(image_folder) if os.path.isfile(os.path.join(image_folder, file))]
scores = compute_similarity(target_image, image_paths)
results = []
for image_path, score in zip(image_paths, scores):
result = f"Image: {image_path}\nScore: {score:.4f}\n"
results.append(result)
return "".join(results)
# Prepare Gradio interface
iface = gr.Interface(
fn=image_similarity,
inputs=[
gr.inputs.Image(type="numpy", label="Target Image"),
gr.inputs.Textbox(label="Image Folder", lines=1)
],
outputs="text",
title="Image Similarity Calculator",
description="Upload a target image and specify the path to the folder containing images. Get similarity scores."
)
# Launch the interface
iface.launch()