Update app.py

app.py

@@ -62,9 +62,9 @@ class_labels = ['battery',
  'trash',
  'white-glass']
 
-def predict_image(input):
+def predict_image(image_input):
     # Resize the image to the size expected by the model
-    image = input.resize((244, 224))
+    image = image_input.resize((244, 224))
     # Convert the image to a NumPy array
     image_array = tf.keras.preprocessing.image.img_to_array(image)
     # Normalize the image
@@ -73,16 +73,13 @@ def predict_image(input):
     image_array = tf.expand_dims(image_array, 0)
     # Predict using the model
     predictions = model1.predict(image_array)
+
+    category_scores = {}
+    for i, class_label in enumerate(class_labels):
+        category_scores[class_label] = predictions[0][i].item()
     
-    # Get the predicted class label
-    predicted_class_index = tf.argmax(predictions, axis=1).numpy()
-    predicted_class_label = class_labels[predicted_class_index]
-    
-    # Get the confidence score of the predicted class
-    confidence_score = predictions[0][predicted_class_index]
-    
-    # Return predicted class label and confidence score
-    return {predicted_class_label: confidence_score}
+    return category_scores
+
 
 image_gradio_app = gr.Interface(
     fn=predict_image,