File size: 5,195 Bytes
6a4d9dc 2d4593d 6a4d9dc 6ef5a6e 6a4d9dc 2d4593d 6a4d9dc 2d4593d 6a4d9dc 2d4593d 6a4d9dc 2d4593d 68ed7d4 2d4593d 6a4d9dc 2d4593d 6a4d9dc 2d4593d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import os
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
# Load the dataset
dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k")
# Take a subset of the dataset
subset_size = 50
dataset_subset = dataset['train'].shuffle(seed=42).select(range(subset_size))
# Directory to save images
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')
# Function to extract features
def extract_features(img_path, model):
img = image.load_img(img_path, target_size=(224, 224))
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)
return features.flatten()
# Extract features for a sample of images
features = []
image_paths = []
model_names = []
for sample in tqdm(dataset_subset):
img_url = sample['url'] # Adjust based on the correct column name
model_name = sample['Model'] # Adjust based on the correct column name
img_path = os.path.join(image_dir, os.path.basename(img_url))
# Download the image
try:
response = requests.get(img_url)
response.raise_for_status() # Check if the download was successful
if 'image' not in response.headers['Content-Type']:
raise ValueError("URL does not contain an image")
with open(img_path, 'wb') as f:
f.write(response.content)
# Extract features
try:
img_features = extract_features(img_path, model)
features.append(img_features)
image_paths.append(img_path)
model_names.append(model_name)
except UnidentifiedImageError:
print(f"UnidentifiedImageError: Skipping file {img_path}")
os.remove(img_path)
except requests.exceptions.RequestException as e:
print(f"RequestException: Failed to download {img_url} - {e}")
# Convert features to numpy array
features = np.array(features)
# Build the NearestNeighbors model
nbrs = NearestNeighbors(n_neighbors=5, algorithm='ball_tree').fit(features)
# Save the model and features
joblib.dump(nbrs, 'nearest_neighbors_model.pkl')
np.save('image_features.npy', features)
np.save('image_paths.npy', image_paths)
np.save('model_names.npy', model_names)
# Load the NearestNeighbors model and features
nbrs = joblib.load('nearest_neighbors_model.pkl')
features = np.load('image_features.npy')
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):
img_features = extract_features(img_path, model)
distances, indices = nbrs.kneighbors([img_features])
recommended_images = [image_paths[idx] for idx in indices.flatten()]
recommended_model_names = [model_names[idx] for idx in indices.flatten()]
recommended_distances = distances.flatten()
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)):
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"), # 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() |