Update app.py

app.py

@@ -8,13 +8,14 @@ from tensorflow.keras.preprocessing import image
 from sklearn.neighbors import NearestNeighbors
 import joblib
 from PIL import UnidentifiedImageError, Image
+import matplotlib.pyplot as plt
 import gradio as gr
 
 # Load the dataset
 dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k")
 
 # Take a subset of the dataset
-subset_size = 50
+subset_size = 100  # Reduce the subset size for faster execution in Spaces
 dataset_subset = dataset['train'].shuffle(seed=42).select(range(subset_size))
 
 # Directory to save images
@@ -86,7 +87,10 @@ image_paths = np.load('image_paths.npy', allow_pickle=True)
 model_names = np.load('model_names.npy', allow_pickle=True)
 
 # Function to get recommendations
-def get_recommendations(img_path, model, nbrs, image_paths, model_names, n_neighbors=5):
+def get_recommendations(img, n_neighbors=5):
+    img_path = "temp_input_image.png"
+    img.save(img_path)
+    
     img_features = extract_features(img_path, model)
     distances, indices = nbrs.kneighbors([img_features])
 
@@ -94,40 +98,20 @@ def get_recommendations(img_path, model, nbrs, image_paths, model_names, n_neigh
     recommended_model_names = [model_names[idx] for idx in indices.flatten()]
     recommended_distances = distances.flatten()
 
-    return recommended_images, recommended_model_names, recommended_distances
+    return [Image.open(img_path) for img_path in recommended_images], recommended_model_names, recommended_distances
 
-def get_recommendations_and_display(img_path):
-    recommended_images, recommended_model_names, recommended_distances = get_recommendations(img_path, model, nbrs, image_paths, model_names)
-    
-    results = []
-    for i in range(len(recommended_images)):
-        result = {
-            "Image": Image.open(recommended_images[i]),
-            "Model Name": recommended_model_names[i],
-            "Distance": recommended_distances[i]
-        }
-        results.append(result)
-    return results
-
-# Define Gradio interface
-def gradio_interface(input_image):
-    input_image.save("input_image.jpg")  # Save the input image
-    recommendations = get_recommendations_and_display("input_image.jpg")
-    outputs = []
-    for i, rec in enumerate(recommendations):
-        outputs.append((rec["Image"], f"{rec['Model Name']} (Distance: {rec['Distance']:.2f})"))
-    return outputs
-
-# Create the Gradio app
-iface = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(label=f"Recommendation {i+1} Image") for i in range(5)] +
-            [gr.Textbox(label=f"Recommendation {i+1} Details") for i in range(5)],
+# Gradio interface
+def display_images(input_image):
+    recommended_images, recommended_model_names, recommended_distances = get_recommendations(input_image)
+    return [(img, f'{name}, Distance: {dist:.2f}') for img, name, dist in zip(recommended_images, recommended_model_names, recommended_distances)]
+
+interface = gr.Interface(
+    fn=display_images,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs=gr.outputs.Carousel(label="Recommended Images", item_shape=(256, 256)),
     title="Image Recommendation System",
-    description="Upload an image to get recommendations based on the image"
+    description="Upload an image and get similar images with their model names and distances."
 )
 
-# Launch the Gradio app
 if __name__ == "__main__":
-    iface.launch()
+    interface.launch()