Update README.md

README.md

@@ -54,25 +54,52 @@ The model was fine-tuned using the following settings:
 To use the fine-tuned model for inference, simply load the model from Hugging Face's Model Hub and input a chest X-ray image:
 
 ```python
-from transformers import ViTForImageClassification, ViTFeatureExtractor
-import torch
 from PIL import Image
+import torch
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# Load model and processor
+processor = AutoImageProcessor.from_pretrained("codewithdark/vit-chest-xray")
+model = AutoModelForImageClassification.from_pretrained("codewithdark/vit-chest-xray")
+
+# Define label columns (class names)
+label_columns = ['Cardiomegaly', 'Edema', 'Consolidation', 'Pneumonia', 'No Finding']
+
+# Step 1: Load and preprocess the image
+image_path = "/content/images.jpeg"  # Replace with your image path
+
+# Open the image
+image = Image.open(image_path)
+
+# Ensure the image is in RGB mode (required by most image classification models)
+if image.mode != 'RGB':
+    image = image.convert('RGB')
+    print("Image converted to RGB.")
+
+# Step 2: Preprocess the image using the processor
+inputs = processor(images=image, return_tensors="pt")
+
+# Step 3: Make a prediction (using the model)
+with torch.no_grad():  # Disable gradient computation during inference
+    outputs = model(**inputs)
 
-# Load model and feature extractor
-model = ViTForImageClassification.from_pretrained('codewithdark/vit-chest-xray')
-feature_extractor = ViTFeatureExtractor.from_pretrained('codewithdark/vit-chest-xray')
+# Step 4: Extract logits and get the predicted class index
+logits = outputs.logits  # Raw logits from the model
+predicted_class_idx = torch.argmax(logits, dim=-1).item()  # Get the class index
 
-# Prepare an image for prediction
-image = Image.open('path_to_chest_xray_image.jpg')
+# Step 5: Map the predicted index to a class label
+# You can also use `model.config.id2label`, but we'll use `label_columns` for this task
+predicted_class_label = label_columns[predicted_class_idx]
 
-# Preprocess the image and make predictions
-inputs = feature_extractor(images=image, return_tensors="pt")
-outputs = model(**inputs)
-logits = outputs.logits
-predictions = torch.sigmoid(logits).squeeze()
+# Output the results
+print(f"Predicted Class Index: {predicted_class_idx}")
+print(f"Predicted Class Label: {predicted_class_label}")
 
-# Display predictions
-print(predictions)
+'''
+Output :
+Predicted Class Index: 4
+Predicted Class Label: No Finding
+'''
 ```
 
 ### Fine-Tuning