Spaces:
Sleeping
Sleeping
| 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) |