Update app.py

app.py

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