import gradio as gr

from tools.extract_features import extract_features_image
from tools.search import search_similar_products
from tools.load_database import select_database

# Default values
MARCAS = ["Ninguno", "Zara", "Adidas", "Nike", "Puma", "Levi's", "Forever 21", "H&M", "Gap", "Converse", "Mango"]
PRENDAS = ["bag", "dress", "hat", "jacket", "pants", "shirt", "shoe", "shorts", "skirt", "sunglass"]


def process_image(image, marca, prenda):
   """Visual Search Pipeline"""
   vgg_search = extract_features_image(image)
   database = select_database(marca, prenda)
   result = search_similar_products(vgg_search, database)
   return result


# Gradio app
examples = [
    ["examples/bag_93.99.png", "Adidas", "bag"],
    ["examples/shirt_82.84.png", "Nike", "shirt"],
    ["examples/skirt_87.54.png", "Puma", "skirt"]
]

title =  "Visual Search 🔍 | Powered by Xpertium SA"
description =  """
               <p style='text-align: center'>Carga tu imagen y selecciona la marca & tipo de prenda.</p>
               """

iface = gr.Interface(
  fn=process_image, 
  inputs=[
     gr.Image(label="imagen"), 
     gr.Dropdown(label="marca", choices=MARCAS, value=MARCAS[1]),
     gr.Dropdown(label="prenda", choices=PRENDAS, value=PRENDAS[1]),
     ], 
  outputs=gr.JSON(label="resultado"),
  title=title,
  examples=examples,
  description=description
)

iface.launch()