File size: 1,058 Bytes
91279bd
ef189c4
 
91279bd
 
 
 
 
 
 
 
ef189c4
 
91279bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import gradio as gr
from fastai.vision.all import *
from huggingface_hub import from_pretrained_fastai
from pathlib import Path
import glob

classes_file = Path('classes.txt')
if not classes_file.exists():
    raise FileNotFoundError(f"{classes_file} not found")
classes = classes_file.read_text().splitlines()

model_path = "makaveli10/tiny_vit_food_classifier"
learn = from_pretrained_fastai(model_path)

sample_folder = Path('samples')
if sample_folder.exists():
    sample_images = sorted(glob.glob(str(sample_folder / '*')))
    examples = [[img] for img in sample_images]
else:
    examples = []

def predict(img):
    # img: PIL image
    pred, idx, probs = learn.predict(img)
    return {classes[i]: float(probs[i]) for i in range(len(classes))}

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type='pil'),
    outputs=gr.Label(num_top_classes=5),
    examples=examples,
    title="Food-101 Classifier",
    description="Upload an image of food or choose from examples to get predictions."
)

if __name__ == "__main__":
    iface.launch()