Update app.py

app.py

@@ -32,14 +32,21 @@ model_file = download_model()
 model = tf.keras.models.load_model(model_file)
 
 # Perform image classification
+# 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)
+#     class_index = tf.argmax(prediction[0]).numpy()
+#     predicted_class = labels[class_index]
+#     return predicted_class
+
 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)
-    class_index = tf.argmax(prediction[0]).numpy()
-    predicted_class = labels[class_index]
-    return predicted_class
+    return prediction
 
 # UI Design
 # def classify_image(image):