Update app.py

app.py

@@ -47,7 +47,7 @@ def predict_class(image):
     img = tf.cast(image, tf.float32)
     img = tf.image.resize(img, [input_shape[0], input_shape[1]])
     img = tf.expand_dims(img, axis=0)
-    prediction = model.predict(img)[0]
+    prediction = model.predict(img)
     return prediction
 
 # UI Design for single class output
@@ -59,7 +59,7 @@ def predict_class(image):
 
 # UI Design for multy class output
 def classify_image(image):
-    results = predict_class(image)
+    results = predict_class(image)[0]
     output = {labels.get(i): float(results[i]) for i in range(len(results))}
     result = output if max(output.values()) >=0.98 else {"NO_CIFAR10_CLASS": 1}
     return result[0]