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

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

def check(img_main, img_1, img_2, model_idx):    
    result = similarity.check_similarity([img_main, img_1, img_2], models[model_idx])
    return result

with gr.Blocks() as demo:
    gr.Markdown('بررسی شباهت عکس ها')
    img_main = gr.Text(label='عکس اصلی', placeholder='https://myimage.jpg')
    
    gr.Markdown('عکس های مقایسه')
    img_1 = gr.Text(label='عکس اول', placeholder='https://myimage_1.jpg')
    img_2 = gr.Text(label='عکس دوم', placeholder='https://myimage_2.jpg')

    gr.Markdown('انتخاب مدل')
    model = gr.Dropdown([m.name for m in models], label='مدل', type='index')
    
    gallery = gr.Gallery(
            label="عکس های تولیدی", show_label=False, elem_id="gallery"
        ).style(grid=[2], height="auto")

    submit_btn = gr.Button('بررسی شباهت')
    submit_btn.click(fn=check,inputs=[img_main, img_1, img_2, model], outputs=gallery)

demo.launch()