import gradio as gr
import base64
from fastai.vision.all import *

learn=load_learner('export.pkl')

labels = learn.dls.vocab
def predict(img):
    img = PILImage.create(img)
    pred,pred_idx,probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

# Read the image file and encode it as base64
with open("./1001epochs.png", "rb") as f:
    image_data = f.read()
    image_base64 = base64.b64encode(image_data).decode("utf-8")

allow_flagging = "never"

title = f"""
    <h2 style="background-image: linear-gradient(to right, #3A5FCD, #87CEFA); -webkit-background-clip: text;
        -webkit-text-fill-color: transparent; text-align: center;">
        Emergency Vehicle Classifier
    </h2>
"""

description = f"""
<div style="display: flex; align-items: center; justify-content: center; flex-direction: column;">
    <p style="font-size: 18px; color: #4AAAFF; text-align: center;">
        Simply upload a photo and let our sophisticated AI system determine the specific type of emergency vehicle depicted.
    </p>
    <div style="display: flex; align-items: center; margin-bottom: 0px;">
        <img src='data:image/jpeg;base64,{image_base64}' width='50' height='30' style="margin-right: 5px;"/>
        <p style="font-size: 14px; color: #555;">
            Disclaimer: The purpose of this application is solely for demonstration. 1001epochs does not claim ownership for the results. Contact: contact@1001epochs.co.uk for full solution.
        </p>
    </div>
</div>
"""

interpretation='default'
gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(512, 512)), outputs=gr.outputs.Label(num_top_classes=3), title=title, interpretation=interpretation, description=description).launch()