fix

app.py

@@ -39,14 +39,23 @@ def predict(img, model_name):
         geo_learn = import_model('geo_model.pkl')
         pred = geo_learn.predict(img)[0]
         labels = geo
+        pred = str(pred)
+        fix_pred = pred[4:-4]
+        return fix_pred + ", Design Codes: " + labels[fix_pred]
     elif (model_name == 'Scenery, Natural Phenomena'):
         landscape_learn = import_model('landscape_model.pkl')
         pred = landscape_learn.predict(img)[0]
         labels = scen
+        pred = str(pred)
+        fix_pred = pred[4:-4]
+        return fix_pred + ", Design Codes: " + labels[fix_pred]
     elif (model_name == 'Human & Supernatural Beings'):
         human_learn = import_model('human_model.pkl')
         pred = human_learn.predict(img)[0]
         labels = char
+        pred = str(pred)
+        fix_pred = pred[4:-4]
+        return fix_pred + ", Design Codes: " + labels[fix_pred]
     elif (model_name == 'Colors & Characters'):
         colors_learn = import_model('colors_model.pkl')
         pred = colors_learn.predict(img)[0]
@@ -68,15 +77,11 @@ def predict(img, model_name):
         pred = str(pred)
         fix_pred = pred[4:-4]
         return fix_pred
-
-    pred = str(pred)
-    fix_pred = pred[4:-4]
-    return fix_pred + ", Design Codes: " + labels[fix_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', 'All'], label='Choose a Model')],
+             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()