Find_objects / app.py
etweedy's picture
Upload 8 files
e4d812f
from fastai.vision.all import *
import gradio as gr
# Define custom functions for the model
def get_x(r): return path/'train'/r['fname']
def get_y(r): return r['labels'].split(' ')
def splitter(df):
train = df.index[~df['is_valid']].tolist()
valid = df.index[df['is_valid']].tolist()
return train,valid
# Load the model
learn=load_learner('obj_class2.pkl')
# The loss function has default threshold of 0.5. It seems to do better with 0.3.
learn.loss_func = BCEWithLogitsLossFlat(thresh=0.3)
# Pull out the list of categories from the model
categories = learn.dls.vocab
cat_list = [x for x in categories]
# Function for classifying image.
def classify_image(img):
pred,idx,probs = learn.predict(img)
idx = list(idx)
answer = ' and '.join([cat_list[i] for i in np.where(idx)[0].tolist()])
if answer:
return answer
else:
return "I don't recognize anything..."
# Initialize and launch gradio interface
image = gr.inputs.Image(shape=(192,192))
label = gr.outputs.Label()
title = 'Object finder'
description = "This app will try to find certain types of objects in the photo it's given. Try one of the examples, or upload your own photo! Keep in mind that it only will recognize the following objects: aeroplane, bicycle, bird, boat, bottle, bus, car, cat, chair, cow, diningtable, dog, horse, motorbike, person, pottedplant, sheep, sofa, train, or tvmonitor"
examples = ['cat_pot.jpeg','cow_bike.jpeg','dog_plane.jpeg','horse_sheep.jpeg','chair_sofa.jpeg','pizza.jpeg']
intf = gr.Interface(fn=classify_image,inputs=image,outputs=label,examples=examples, title=title,description=description)
intf.launch(inline=False)