DogBreeds / app.py
Tesneem's picture
Update app.py
22e128d verified
import gradio as gr
import numpy as np
import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer
import requests
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration
sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
image_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
def generate_input(input_type, image=None, text=None, response_amount=3):
# initalize input variable
combined_input = ""
# handle image input if chosen
if input_type == "Image" and image:
inputs = processor(images=image, return_tensors="pt") #process image with BlipProcessor
out = image_model.generate(**inputs) #generate caption with BlipModel
image_caption = processor.decode(out[0], skip_special_tokens=True) #decode output w/ processor
combined_input += image_caption # add the image caption to input
# handle text input if chosen
elif input_type == "Text" and text:
combined_input += text # add the text to input
# handle both text and image input if chosen
elif input_type == "Both" and image and text:
inputs = processor(images=image, return_tensors="pt")
out = image_model.generate(**inputs)
image_caption = processor.decode(out[0], skip_special_tokens=True) #repeat image processing + caption generation and decoding
combined_input += image_caption + " and " + text # combine image caption and text
# if no input, fallback
if not combined_input:
combined_input = "No input provided."
if response_amount is None:
response_amount=3
return vector_search(combined_input,response_amount) #search through embedded document w/ input
# load embeddings and metadata
embeddings = np.load("dog_data_embeddings.npy") #created using sentence_transformers on kaggle
metadata = pd.read_csv("dog_metadata.csv") #created using sentence_transformers on kaggle
# vector search function
def vector_search(query,top_n=3):
query_embedding = sentence_model.encode(query) #encode input w/ Sentence Transformers
similarities = cosine_similarity([query_embedding], embeddings)[0] #similarity function
if top_n is None:
top_n=3
top_indices = similarities.argsort()[-top_n:][::-1] #return top n indices based on chosen output amount
results = metadata.iloc[top_indices] #get metadata
result_text=""
for index,row in results.iterrows(): #loop through results to get Title, Description, and Genre for top n outputs
if index!=top_n-1:
result_text+=f"Breed: {row['breed']} Description: {row['description']} Temperament: {row['temperament']} Energy Level: {row['energy_level_category']} Trainability: {row['trainability_category']} Demeanor: {row['demeanor_category']} \n\n"
else:
result_text+=f"Breed: {row['breed']} Description: {row['description']} Temperament: {row['temperament']} Energy Level: {row['energy_level_category']} Trainability: {row['trainability_category']} Demeanor: {row['demeanor_category']}"
return result_text
def set_response_amount(response_amount): #set response amount
if response_amount is None:
return 3
return response_amount
# based on the selected input type, make the appropriate input visible
def update_inputs(input_type):
if input_type == "Image":
return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True)
elif input_type == "Text":
return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
elif input_type == "Both":
return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
with gr.Blocks() as demo:
gr.Markdown("# Dog Breed Recommendation System")
gr.Markdown("Enter a query to receive dog breed recommendations based on description, temperament, trainability, and demeanor.")
input_type = gr.Radio(["Image", "Text", "Both"], label="Select Input Type", type="value")
response_type=gr.Dropdown(choices=[3,5,10,25], type="value", label="Select Response Amount", visible=False)
image_input = gr.Image(label="Upload Image", type="pil", visible=False) # Hidden initially
text_input = gr.Textbox(label="Enter Text Query", placeholder="Enter a description or query here", visible=False) # hidden initially
input_type.change(fn=update_inputs, inputs=input_type, outputs=[image_input, text_input, response_type])
# state variable to store the selected response amount
selected_response_amount = gr.State()
# capture response amount immediately when dropdown changes
response_type.change(fn=set_response_amount, inputs=response_type, outputs=selected_response_amount)
submit_button = gr.Button("Submit")
output = gr.Textbox(label="Recommendations")
if selected_response_amount is None:
selected_response_amount=3
submit_button.click(fn=generate_input, inputs=[input_type,image_input, text_input,selected_response_amount], outputs=output)
demo.launch()