import gradio as gr
from fastai.vision.all import *  # noqa: ruff[F403]

learner = load_learner("chemical.pkl")  # noqa: ruff[F403]


def classify(img):
    pred, idx, probs = learner.predict(img)

    return dict(zip(("chemical plant", "power plant"), map(float, probs)))


ui = gr.Interface(
    classify,
    inputs=gr.Image(
        label="Upload image of chemical or power plant",
        type="filepath",
        shape=(192, 192),
    ),
    outputs=gr.Label(),
    title="Chemical plant? Or Not?",
    examples=["chemical.jpg", "power.jpg"],
)

ui.launch()