Update app.py

app.py

@@ -3,7 +3,6 @@ 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
@@ -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
@@ -122,15 +105,13 @@ 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))
     
-    # Display images with matplotlib
-    display_images(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>
+            <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>
@@ -138,16 +119,6 @@ def recommend(image):
     
     return html_output
 
-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"),