Update app.py

app.py

@@ -3,6 +3,7 @@ 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
@@ -22,8 +23,23 @@ 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
-model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
+with tf.device(device):
+    model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
 
 # Function to extract features
 def extract_features(img_path, model):
@@ -31,7 +47,8 @@ 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)
-    features = model.predict(img_array)
+    with tf.device(device):
+        features = model.predict(img_array)
     return features.flatten()
 
 # Extract features for a sample of images
@@ -105,26 +122,27 @@ def recommend(image):
     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))
     
-    # Prepare HTML output for Gradio
-    html_output = ""
-    for img_path, model_name, distance in zip(recommended_images, recommended_model_names, recommended_distances):
-        img_path = img_path.replace('\\', '/')
-        html_output += f"""
-        <div style='display:inline-block; text-align:center; margin:10px;'>
-            <img src='file/{img_path}' style='width:200px; height:200px;'><br>
-            <b>Model Name:</b> {model_name}<br>
-            <b>Distance:</b> {distance:.2f}<br>
-        </div>
-        """
+    # Display images with matplotlib
+    display_images(recommended_images, recommended_model_names, recommended_distances)
     
-    return html_output
+    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)):
+        img = Image.open(img_path)
+        plt.subplot(1, len(image_paths), i+1)
+        plt.imshow(img)
+        plt.title(f'{model_name}\nDistance: {distance:.2f}', fontsize=12)
+        plt.axis('off')
+    plt.show()
 
 interface = gr.Interface(
     fn=recommend,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.HTML(),  # Use HTML output for better formatting
+    inputs=gr.Image(type="pil"),  # Updated input component
+    outputs=gr.Text(),  # Updated output component
     title="Image Recommendation System",
     description="Upload an image and get 5 recommended similar images with model names and distances."
 )
 
-interface.launch()
+interface.launch()