Update app.py

app.py

@@ -3,14 +3,13 @@ import requests
 from tqdm import tqdm
 from datasets import load_dataset
 import numpy as np
-import tensorflow as tf
 from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
 from tensorflow.keras.preprocessing import image
 from sklearn.neighbors import NearestNeighbors
 import joblib
 from PIL import UnidentifiedImageError, Image
-import gradio as gr
 import matplotlib.pyplot as plt
+import gradio as gr
 
 # Load the dataset
 dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k")
@@ -23,23 +22,8 @@ dataset_subset = dataset['train'].shuffle(seed=42).select(range(subset_size))
 image_dir = 'civitai_images'
 os.makedirs(image_dir, exist_ok=True)
 
-# Try to use GPU, fall back to CPU if not available
-try:
-    gpus = tf.config.list_physical_devices('GPU')
-    if gpus:
-        tf.config.experimental.set_memory_growth(gpus[0], True)
-        device = '/GPU:0'
-        print("Using GPU")
-    else:
-        raise RuntimeError("No GPU found")
-except RuntimeError as e:
-    print(e)
-    device = '/CPU:0'
-    print("Using CPU")
-
 # Load the ResNet50 model pretrained on ImageNet
-with tf.device(device):
-    model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
+model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
 
 # Function to extract features
 def extract_features(img_path, model):
@@ -47,8 +31,7 @@ def extract_features(img_path, model):
     img_array = image.img_to_array(img)
     img_array = np.expand_dims(img_array, axis=0)
     img_array = preprocess_input(img_array)
-    with tf.device(device):
-        features = model.predict(img_array)
+    features = model.predict(img_array)
     return features.flatten()
 
 # Extract features for a sample of images
@@ -114,19 +97,6 @@ def get_recommendations(img_path, model, nbrs, image_paths, model_names, n_neigh
 
     return recommended_images, recommended_model_names, recommended_distances
 
-def recommend(image):
-    # Save uploaded image to a path
-    image_path = "uploaded_image.jpg"
-    image.save(image_path)
-    
-    recommended_images, recommended_model_names, recommended_distances = get_recommendations(image_path, model, nbrs, image_paths, model_names)
-    result = list(zip(recommended_images, recommended_model_names, recommended_distances))
-    
-    # Display images with matplotlib
-    display_images(recommended_images, recommended_model_names, recommended_distances)
-    
-    return result
-
 def display_images(image_paths, model_names, distances):
     plt.figure(figsize=(20, 10))
     for i, (img_path, model_name, distance) in enumerate(zip(image_paths, model_names, distances)):
@@ -137,12 +107,19 @@ def display_images(image_paths, model_names, distances):
         plt.axis('off')
     plt.show()
 
-interface = gr.Interface(
-    fn=recommend,
-    inputs=gr.Image(type="pil"),  # Updated input component
-    outputs=gr.Text(),  # Updated output component
+def recommend_images(img):
+  recommended_images, recommended_model_names, recommended_distances = get_recommendations(img, model, nbrs, image_paths, model_names)
+  return [Image.open(img_path) for img_path in recommended_images], [model_name for model_name in recommended_model_names], [distance for distance in recommended_distances]
+
+iface = gr.Interface(
+    fn=recommend_images,
+    inputs=gr.Image(label="Upload an image"),
+    outputs=[
+        gr.Gallery(label="Recommended Images", show_label=False),
+        gr.Textbox(label="Model Names", lines=5),
+        gr.Textbox(label="Distances", lines=5)
+    ],
     title="Image Recommendation System",
-    description="Upload an image and get 5 recommended similar images with model names and distances."
+    description="Upload an image and get recommendations based on similarity."
 )
-
-interface.launch()
+iface.launch()