Spaces:
Running
Running
| from gradio import gr | |
| from fastai.vision.all import load_learner | |
| CATS_MAP = { | |
| "picasso": "Pablo Picasso", | |
| "vanGogh": "Vincent van Gogh", | |
| "dali": "Salvador Dalí", | |
| "daVinci": "Leonardo da Vinci", | |
| "rembrandt": "Rembrandt", | |
| } | |
| # 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() | |
| examples = [ | |
| f"images_examples/{filename}" for filename in ("mona_lisa.jpg", "starry_night.jpg", "persistence_of_memory.jpg") | |
| ] | |
| gui = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples) | |
| gui.launch(inline=False) |