jsulz's picture
jsulz HF staff
initial commit
4fddbf2
raw
history blame contribute delete
806 Bytes
import gradio as gr
from PIL import Image
from fastai.vision.all import *
learn = load_learner("model.pkl")
def predict(img):
# use PIL to resize the image to 224x224
img = Image.fromarray(img)
pred, pred_idx, probs = learn.predict(img)
# if the prediction contains poodle, return poodle
if "poodle" in pred:
return "poodle"
return "doodle"
title = "Poodle or Doodle"
description = "A resnet18 model fine tuned on 400 images of poodles and labradoodles. You can predict using the default image or upload an image an image of your own to see if you're looking at a poodle or doodle."
gr.Interface(
fn=predict,
inputs=gr.Image(value="labradoodle.jpg"),
outputs=gr.Label(),
title=title,
description=description,
allow_flagging="never",
).launch()