import torch
import gradio as gr
from adlike import ad_openai_clip_vitl_patch14_336
import glob

device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = ad_openai_clip_vitl_patch14_336()

demo = gr.Blocks()

def predict_ad_probab(image):
    with torch.no_grad():
        image = preprocess(image).to(device).unsqueeze(0)
        probs = model(image).item()
        return round(probs, 3)

with demo:
    
    gr.Markdown("# **<p align='center'>AdLike: Detect Advertisement Images</p>**")
    gr.Markdown("This space demonstrates the use of AdLike. It predicts the probability of whether an Image is an Advertisement. \
        The higher the probability, the higher the chance of an Image being Advertisement.")
       
    with gr.Group():            
        with gr.Row():
            input_image = gr.Image(type='pil',label="Input Image",
                    image_mode="RGB", 
                    sources="upload",
                    show_label=True)
            
        with gr.Row():
            output_probab = gr.Number(label="Advertisement Probability", show_label=True)
        
    with gr.Group():
        with gr.Row():
            submit_button = gr.Button("Predict Probability")
    
    submit_button.click(predict_ad_probab, inputs=[input_image], outputs=[output_probab])

demo.launch(debug=True)