Amazon-Lens / app.py
rishikaboinapally's picture
Rename main.py to app.py
025ab78 verified
import os
import numpy as np
import pickle
import tensorflow as tf
from PIL import Image
from tensorflow.keras.preprocessing import image
from tensorflow.keras.layers import GlobalMaxPooling2D
from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
from sklearn.neighbors import NearestNeighbors
from numpy.linalg import norm
import gradio as gr
# Function to load embeddings and filenames
def load_data():
# Load embeddings and filenames from the same directory as main.py
embeddings_path = os.path.join(os.path.dirname(__file__), 'embeddings.pkl')
filenames_path = os.path.join(os.path.dirname(__file__), 'filenames.pkl')
feature_list = np.array(pickle.load(open(embeddings_path, 'rb')))
filenames = pickle.load(open(filenames_path, 'rb'))
return feature_list, filenames
# Load the feature list and filenames
feature_list, filenames = load_data()
# Load the pre-trained model
model = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
model.trainable = False
model = tf.keras.Sequential([
model,
GlobalMaxPooling2D()
])
def feature_extraction(img_path, model):
img = image.load_img(img_path, target_size=(224, 224))
img_array = image.img_to_array(img)
expanded_img_array = np.expand_dims(img_array, axis=0)
preprocessed_img = preprocess_input(expanded_img_array)
result = model.predict(preprocessed_img).flatten()
normalized_result = result / norm(result)
return normalized_result
def recommend(features, feature_list):
neighbors = NearestNeighbors(n_neighbors=6, algorithm='brute', metric='euclidean')
neighbors.fit(feature_list)
distances, indices = neighbors.kneighbors([features])
return indices
def resize_image(img, size):
img = img.resize(size, Image.LANCZOS)
return img
def predict(input_image):
img = input_image.convert("RGB") # Ensure image is in RGB mode
img_array = np.array(img)
expanded_img_array = np.expand_dims(img_array, axis=0)
preprocessed_img = preprocess_input(expanded_img_array)
features = model.predict(preprocessed_img).flatten()
normalized_features = features / norm(features)
indices = recommend(normalized_features, feature_list)
result_images = [filenames[idx] for idx in indices[0]]
resized_images = [resize_image(Image.open(img_path), (224, 224)) for img_path in result_images]
return resized_images
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Upload an image"),
outputs=[gr.Image(type="pil", label=f"Recommendation {i+1}") for i in range(5)],
title="Amazon Lens",
description="Upload an image of clothing, and the system will recommend similar items."
)
if __name__ == "__main__":
iface.launch()