teenaxta's picture
Rename main_gradio.py to app.py
e2f79f3 verified
raw
history blame
2.17 kB
from dotenv import load_dotenv
load_dotenv("docker/.env")
from pymongo.mongo_client import MongoClient
from pymongo.server_api import ServerApi
from src.mongo_search import MongoSearch
from src.image_search import ImageSearcher
import os
import gradio as gr
uri = os.getenv("MONGO_URI")
db = os.getenv("MONGO_DB")
cars_collection = os.getenv("MONGO_CARS_COLLECTION")
image_collection = os.getenv("MONGO_IMAGE_COLLECTION")
img_search_index = os.getenv("MONGO_SEARCH_IMG_INDEX")
car_search_index = os.getenv("MONGO_SEARCH_CAR_INDEX")
client = MongoClient(uri, server_api=ServerApi('1'))
db = client[db]
cars_collection = db[cars_collection]
image_collection = db[image_collection]
mongo_search = MongoSearch(cars_collection, car_search_index, index_variable="review_embedding") #desc_embedding
image_searcher = ImageSearcher(image_collection,
cars_collection,
img_search_index)
def search_mongo(search_query):
object = mongo_search(search_query)
return object[0]['makeModel'], object[0]['photoUrls']
def search_img(image):
result = image_searcher(image)['results']
return result[0]['makeModel'], result[0]['photoUrls']
with gr.Blocks() as demo:
with gr.Tab("Text search"):
text_output_img = gr.Gallery()
text_output_text = gr.Textbox(label='Output car model', show_label=True)
text_input_text = gr.Textbox(label='Search Query', show_label=True)
text_btn = gr.Button()
with gr.Tab("Image search"):
with gr.Row():
with gr.Column():
input_img = gr.Image(type='pil')
img_btn = gr.Button()
with gr.Column():
img_output_img = gr.Gallery(preview=True)
img_output_text = gr.Textbox(label='Output car model', show_label=True)
text_btn.click(search_mongo,
inputs=text_input_text,
outputs=[text_output_text, text_output_img])
img_btn.click(search_img,
inputs=input_img,
outputs=[img_output_text, img_output_img])
if __name__ == "__main__":
demo.launch()