import gradio as gr
import os
import random
from src.model import simlarity_model as model
from src.similarity.similarity import Similarity
from PIL import Image
import glob

similarity = Similarity()
models = similarity.get_models()

def check(img_main, model_idx):    
    image_list = []
    for filename in glob.glob('images/*.jpg'):
        im=Image.open(filename)
        image_list.append(im)
    result = similarity.check_similarity([img_main] + image_list, models[model_idx])
    return result

with gr.Blocks() as demo:
    gr.Markdown('Checking Image Similarity')
    img_main = gr.Text(label='Main Image', placeholder='https://myimage.jpg')
    

    gr.Markdown('Choose the model')
    model = gr.Dropdown([m.name for m in models], label='Model', type='index')
    
    gallery = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery"
        ).style(grid=[2], height="auto")

    submit_btn = gr.Button('Check Similarity')
    submit_btn.click(fn=check,inputs=[img_main, model], outputs=gallery)

demo.launch()