Update app.py

app.py

@@ -46,6 +46,7 @@ def predict_class(image):
     img = tf.image.resize(img, [input_shape[0], input_shape[1]])
     img = tf.expand_dims(img, axis=0)
     prediction = model.predict(img)
+    print(prediction[0])
     return prediction
 
 # UI Design
@@ -58,7 +59,7 @@ def classify_image(image):
     results = predict_class(image)
     output = {}
     for _, prediction in enumerate(results):
-        predicted_label = labels[_]
+        predicted_label = labels.get(_)
         score = prediction[_]
         output[predicted_label] = score
     return output