Update app.py

app.py

@@ -14,6 +14,9 @@ def get_base64_of_bin_file(bin_file):
     return base64.b64encode(data).decode()
 
 
+learn = load_learner('export (3).pkl')
+
+categories = ('Dog Destroying Stuff','Neutral')
 upload= st.file_uploader('Insert image for classification', type=['png','jpg'])
 c1, c2= st.columns(2)
 if upload is not None:
@@ -25,15 +28,7 @@ if upload is not None:
   c1.header('Input Image')
   c1.image(im)
   c1.write(img.shape)
-  input_shape = (224, 224, 3)
-  optim_1 = Adam(learning_rate=0.0001)
-  n_classes=6
-  vgg_model = model(input_shape, n_classes, optim_1, fine_tune=2)
-  vgg_model.load_weights('/content/drive/MyDrive/vgg/tune_model19.weights.best.hdf5')
-
-  # prediction on model
-  vgg_preds = vgg_model.predict(img)
-  vgg_pred_classes = np.argmax(vgg_preds, axis=1)
+  pred, idx, probs = learn.predict(img)
   c2.header('Output')
   c2.subheader('Predicted class :')
-  c2.write(classes[vgg_pred_classes[0]] )
+  c2.write(pred)