Update app.py

app.py

@@ -21,14 +21,27 @@ def classify_image(image):
         outputs = model(**inputs)
         logits = outputs.logits
 
-    # Retrieve the highest probability class label
+    # Retrieve the highest probability class label index
     predicted_class_idx = logits.argmax(-1).item()
 
+    # Define a manual mapping of label indices to human-readable labels
+    # This should be adjusted based on your dataset and labels
+    index_to_label = {
+        "0": "Aluminium",
+        "1": "Batteries",
+        "2": "Cardboard",
+        "3": "Glass",
+        "4": "Hard Plastic",
+        "5": "Paper",
+        "6": "Soft Plastics"
+    }
+
     # Convert the index to the model's class label
-    label = model.config.id2label[predicted_class_idx]
+    label = index_to_label.get(predicted_class_idx, "Unknown Label")
 
     return label
 
+
 # Create Gradio interface
 iface = gr.Interface(fn=classify_image, 
                      inputs=gr.Image(),  # Accepts image of any size