import gradio as gr
from fastai.vision.all import load_learner
# PIL ANTES: 8.2.0

CATS_MAP = {
    "eColi": "E. coli",
    "staphylococcusAureus": "Staphylococcus aureus",
    "streptococcusPyogenes": "Streptococcus pyogenes",
    "pseudomonasAeruginosa": "Pseudomonas aeruginosa",
    "salmonellaSpp": "Salmonella spp",
    "klebsiellaPneumoniae": "Klebsiella pneumoniae",
    "proteusMirabilis": "Proteus mirabilis",
    "enterococcusFaecalis": "Enterococcus faecalis",
    "lactobacillus": "Lactobacillus",
}

# load pre-trained model
model = load_learner("model.pkl")

# get classes name in right order
full_name_cats = [CATS_MAP[key_class] for key_class in model.dls.vocab]


def classify_image(img) -> dict:
    category, idx, probs = model.predict(img)

    return dict(zip(full_name_cats, map(float, probs)))


# Gradio control
image = gr.inputs.Image(shape=(224, 224))
label = gr.outputs.Label(num_top_classes=4)
examples = [
    f"images_examples/{filename}" 
    for filename in ("klebsiella.jpg", "enterococcus.jpg", "salmonella.jpg")
]

gui = gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
    examples=examples,
    title="Detect bacterias - first version",
    description=(
        f"Detect if the given Petri dish culture image belongs to one of those: ({full_name_cats})."
    )
)
gui.launch(inline=False)