Update README.md

README.md

@@ -29,13 +29,64 @@ You can use this model directly with a pipeline for image classification:
 \`\`\`python
 
 
-from transformers import pipeline
-classifier = pipeline('image-classification', model='fawadkhan/ViT_FineTuned_on_ImagesOASIS')
+import torch
+from transformers import ViTForImageClassification
+from PIL import Image
+import numpy as np
+from torchvision.transforms import Compose, Resize, ToTensor, Normalize
 
-# To classify an image:
-image_path = 'path_to_your_image.jpg'
-prediction = classifier(image_path)
-print(prediction)
+
+id2label = {
+    0: "Non-Demented",
+    1: "Very Mild Demented",
+    2: "Mild Demented",
+    3: "Demented"
+}
+
+import torch
+from transformers import ViTForImageClassification
+from PIL import Image
+import numpy as np
+from torchvision.transforms import Compose, Resize, ToTensor, Normalize
+import matplotlib.pyplot as plt
+
+# Set the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load the model
+model = ViTForImageClassification.from_pretrained('fawadkhan/ViT_FineTuned_on_ImagesOASIS')
+model.to(device)
+model.eval()
+
+# Define the image path
+image_path = 'your image path.jpg'
+image = Image.open(image_path).convert("RGB")
+
+# Define the transformations
+transform = Compose([
+    Resize((224, 224)),  # or the original input size of your model
+    ToTensor(),
+    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Standard normalization for ImageNet
+])
+
+# Preprocess the image
+input_tensor = transform(image).unsqueeze(0)  # Create a mini-batch as expected by the model
+input_tensor = input_tensor.to(device)
+
+# Predict
+with torch.no_grad():
+    outputs = model(input_tensor)
+    _, predicted = torch.max(outputs.logits, 1)
+
+# Retrieve the class name
+predicted_class = id2label[predicted[0].item()]
+print("Predicted class:", predicted_class)
+
+# Plot the image and the prediction
+plt.imshow(image)
+plt.title(f'Predicted class: {predicted_class}')
+plt.axis('off')  # Turn off axis numbers and ticks
+plt.show()
 
 
 \`\`\`