modified model

app.py

@@ -6,13 +6,22 @@ import pickle
 st.title("Piano or Keyboard?")
 file_name = st.file_uploader("Upload a piano or a keyboard image")
 
-model = pickle.load(open('export.pkl','rb'))
+# model = pickle.load(open('export.pkl','rb'))
+
+model = load_learner('export.pkl')
+
+labels = model.dls.vocab
+def predict(img):
+    img = PILImage.create(img)
+    pred,pred_idx,probs = model.predict(img)
+    return {labels[i]: float(probs[i]) for i in range(len(labels))}
 
 if file_name is not None:
     col1, col2 = st.columns(2)
     image = Image.open(file_name)
     col1.image(image, use_column_width=True)
-    pred,pred_idx,probs = model.predict(PILImage.create(file_name))
+    # pred,pred_idx,probs = model.predict(PILImage.create(file_name))
+    pred = predict(file_name)
     col2.header("Prediction:")
     st.write('This is a ', pred, '!')
 else: