Update README.md

README.md

@@ -23,6 +23,9 @@ import torch
 processor = ViTImageProcessor.from_pretrained("yaya36095/ai-image-detector")
 model = ViTForImageClassification.from_pretrained("yaya36095/ai-image-detector")
 
+# Define class labels (Replace these with the labels used during training)
+class_labels = ["Real Image", "AI Generated"]
+
 def detect_image(image_path):
     # Open image
     image = Image.open(image_path)
@@ -35,16 +38,22 @@ def detect_image(image_path):
         outputs = model(**inputs)
         predictions = outputs.logits.softmax(dim=-1)
     
-    # Get result
-    prediction_id = torch.argmax(predictions).item()
-    confidence = predictions[0][prediction_id].item() * 100
+    # Map predictions to class labels
+    results = [
+        {"label": class_labels[i], "confidence": f"{(predictions[0][i].item() * 100):.2f}%"}
+        for i in range(len(class_labels))
+    ]
+    
+    # Sort results by confidence (optional)
+    results = sorted(results, key=lambda x: float(x["confidence"].strip('%')), reverse=True)
     
-    result = "AI Generated" if prediction_id == 1 else "Real Image"
-    return result, confidence
+    return results
 
 # Example usage
-# result, confidence = detect_image("path/to/image.jpg")
-# print(f"Result: {result} (Confidence: {confidence:.2f}%)")
+# result = detect_image("path/to/image.jpg")
+# for r in result:
+#     print(f"Class: {r['label']}, Confidence: {r['confidence']}")
+
 ```
 
 ## Classes