import gradio as gr
from fastai.vision.all import load_learner

# Custom CSS for the app (remove background color from examples)
custom_css = """
body { background-color: #1E1E1E; color: #F5F5F5; }
.gradio-container { max-width: 800px; margin: auto; }
.input-panel, .output-panel { 
    background-color: #2E2E2E; 
    border: 2px solid #00BFFF; 
    border-radius: 10px; 
    box-shadow: 0 0 10px #00BFFF; 
    padding: 20px; 
    margin: 10px; 
}
button { 
    background: linear-gradient(90deg, #00BFFF, #FF6B6B); 
    border: none; 
    color: #F5F5F5; 
    padding: 2px 5px;  
    border-radius: 2px;
    cursor: pointer; 
}
button:hover { 
    box-shadow: 0 0 10px #00BFFF; 
}
"""

def classify_img(img) -> dict:
    """helper function to generate predictions"""
    
    lbls = ['Cat', 'Dog']
    preds, idx, probs = model.predict(img)
    
    return dict(zip(lbls, map(float, probs)))

# load model
model = load_learner('model_2_ep_new_data.pkl')

# build app
image = gr.Image()
label = gr.Label()
examples = './inf_cs/'

app = gr.Interface(fn = classify_img, inputs = image, outputs = label, examples = examples, theme = 'earneleh/paris', css = custom_css,
examples_per_page = 7, flagging_mode = 'never', clear_btn = None)

app.launch()
# app.close()