Spaces:
Sleeping
Sleeping
| ''' | |
| Model Gradio UI | |
| ''' | |
| ######################################################################### | |
| # imports | |
| from fastai.vision.all import * | |
| import gradio as gr | |
| import pathlib | |
| from huggingface_hub import hf_hub_download | |
| ######################################################################### | |
| # user access token for HF model library | |
| ACCESS_TOKEN = "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" | |
| ######################################################################### | |
| #Consider path seperators for alternate OS | |
| plt = platform.system() | |
| if plt != 'Windows': pathlib.WindowsPath = pathlib.PosixPath | |
| ######################################################################### | |
| def import_model(model_name): | |
| path = hf_hub_download(repo_id='amandasarubbi/tm-tko-models', | |
| filename=model_name, | |
| use_auth_token=ACCESS_TOKEN, | |
| repo_type='model') | |
| learn = load_learner(path, cpu=True) | |
| return learn | |
| ######################################################################### | |
| ######################################################################### | |
| # Function to predict outputs | |
| def predict(img, model_name): | |
| if (model_name == 'Geometric Figures & Solids'): | |
| geo_learn = import_model('geometric_model.pkl') | |
| preds = geo_learn.predict(img) | |
| elif (model_name == 'Scenery, Natural Phenomena'): | |
| landscape_learn = import_model('landscape_model.pkl') | |
| preds = landscape_learn.predict(img) | |
| elif (model_name == 'Human & Supernatural Beings'): | |
| human_learn = import_model('human_model.pkl') | |
| preds = human_learn.predict(img) | |
| elif (model_name == 'Colors & Characters'): | |
| colors_learn = import_model('colors_model.pkl') | |
| preds = colors_learn.predict(img) | |
| elif (model_name == 'Buildings, Dwellings & Furniture'): | |
| build_learn = import_model('buildings.pkl') | |
| preds = build_learn.predict(img) | |
| elif (model_name == 'Animals'): | |
| anim_learn = import_model('animals.pkl') | |
| preds = anim_learn.predict(img) | |
| label_pred = str(preds[0]) | |
| return label_pred | |
| ######################################################################### | |
| title = "TM-TKO Trademark Logo Image Classification Model" | |
| description = "Users can upload an image and corresponding image file name to get US design-code standard predictions on a trained model that utilizes the benchmark ResNet50 architecture." | |
| iFace = gr.Interface(fn=predict, | |
| inputs=[gr.inputs.Image(label="Upload Logo Here"), gr.inputs.Dropdown(choices=['Geometric Figures & Solids', 'Scenery, Natural Phenomena', 'Human & Supernatural Beings', 'Colors & Characters', 'Buildings, Dwellings & Furniture', 'Animals'], label='Choose a Model')], | |
| outputs=gr.Label(label="TM-TKO Trademark Classification Model"), | |
| title=title, description=description) | |
| iFace.launch() |