pasta / app.py
tobthecreator's picture
Update app.py
262b3c5
import gradio as gr
from fastai.vision.all import *
pasta_shapes = [
"spaghetti",
"penne",
"fettuccine",
"rigatoni",
"bowtie",
"linguine",
"ravioli",
"lasagna",
"rotini",
]
not_pasta = [
"rice", "quinoa", "couscous",
"bread", "pastries", "cookies",
"vegetables", "fruits", "salads",
"stir-fry dishes", "pizza", "sandwiches",
"kitchen utensils", "appliances", "cookware",
"furniture", "home decor", "office supplies",
"landscapes", "cityscapes", "forests",
"plants", "flowers", "trees",
"animals", "birds", "insects",
"abstract art", "patterns", "textures",
"fabric textures", "wood textures", "metal textures"
]
if __name__ == '__main__':
categories = ['bowtie', 'fettuccine', 'lasagna', 'linguine', 'not pasta', 'penne', 'ravioli', 'rigatoni', 'rotini', 'spaghetti']
def get_y(file_path):
parent_folder = file_path.parent.name
if parent_folder in pasta_shapes:
return parent_folder
else:
return "not pasta"
learn = load_learner('model.pkl')
def classify_image(img):
_, _, probs = learn.predict(img)
return dict(zip(categories, map(float, probs)))
image = gr.Image()
label = gr.Label()
intf = gr.Interface(fn=classify_image, inputs=image, outputs=label)
intf.launch(inline=False)