Update app.py

app.py

@@ -49,18 +49,7 @@ from PIL import Image
 model1 = from_pretrained_keras("rocioadlc/modelo")
 
 # Define class labels
-class_labels = ['battery',
- 'biological',
- 'brown-glass',
- 'cardboard',
- 'clothes',
- 'green-glass',
- 'metal',
- 'paper',
- 'plastic',
- 'shoes',
- 'trash',
- 'white-glass']
+class_labels = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 
 # Function to predict image label and score
 def predict_image(input):
@@ -75,15 +64,12 @@ def predict_image(input):
     # Predict using the model
     predictions = model1.predict(image_array)
     
-    # Get the predicted class label
-    predicted_class_index = tf.argmax(predictions, axis=1).numpy()[0]
-    predicted_class_label = class_labels[predicted_class_index]
+    category_scores = {}
+    for i, class_label in enumerate(class_labels):
+        category_scores[class_label] = predictions[0][i].item()
     
-    # 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
+
 
 
 loader = WebBaseLoader(["https://www.epa.gov/recycle/frequent-questions-recycling", "https://www.whitehorsedc.gov.uk/vale-of-white-horse-district-council/recycling-rubbish-and-waste/lets-get-real-about-recycling/", "https://www.teimas.com/blog/13-preguntas-y-respuestas-sobre-la-ley-de-residuos-07-2022", "https://www.molok.com/es/blog/gestion-de-residuos-solidos-urbanos-rsu-10-dudas-comunes"])