Update app.py

app.py

@@ -6,12 +6,11 @@ from transformers import AutoImageProcessor, AutoModelForImageClassification
 processor = AutoImageProcessor.from_pretrained("tiya1012/vit-accident-image")
 model = AutoModelForImageClassification.from_pretrained("tiya1012/vit-accident-image")
 
-# Get class labels from model configuration (if available)
-if model.config.id2label:
-    class_labels = [model.config.id2label[i] for i in range(len(model.config.id2label))]
-else:
-    # Fallback labels if config is missing (adjust as necessary)
-    class_labels = ["No Accident", "Accident Detected"]
+# Define a label mapping for `LABEL_0` and `LABEL_1`
+label_mapping = {
+    "LABEL_0": "No Accident",
+    "LABEL_1": "Accident Detected"
+}
 
 # Define the classification function
 def classify_accident_image(image):
@@ -30,14 +29,17 @@ def classify_accident_image(image):
     # Debug: Print logits for analysis
     print("Logits:", logits)
 
-    # Get the predicted class index
+    # Get the predicted class index and label
     probabilities = torch.softmax(logits, dim=1)[0]  # Softmax to get probabilities
     predicted_class_idx = torch.argmax(probabilities).item()
     print("Predicted Class Index:", predicted_class_idx)
     print("Probabilities:", probabilities)
 
-    # Get the predicted class label and confidence
-    predicted_label = class_labels[predicted_class_idx]
+    # Map the model's label to human-readable label using label_mapping
+    predicted_label_key = model.config.id2label[predicted_class_idx]
+    predicted_label = label_mapping.get(predicted_label_key, "Unknown")
+
+    # Get the confidence score
     confidence = probabilities[predicted_class_idx].item() * 100
 
     # Format the result