Update README.md

README.md

@@ -43,6 +43,7 @@ Preprocessing: Images are resized to 224x224 pixels and normalized across RGB ch
 Training: Pretraining was conducted on TPUv3 hardware with a batch size of 4096 and learning rate warmup. Gradient clipping was applied during training to enhance stability.
 ```python
 from transformers import ViTImageProcessor, ViTForImageClassification
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 from PIL import Image
 import requests
 import torch
@@ -51,9 +52,9 @@ def predict_image_from_url(url):
     # Load image from URL
     image = Image.open(requests.get(url, stream=True).raw)
     
-    # Initialize processor and model
-    processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
-    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
+    # Initialize Sreekanth's processor and model
+    processor = AutoImageProcessor.from_pretrained("Sreekanth3096/vit-coco-image-classification")
+    model = AutoModelForImageClassification.from_pretrained("Sreekanth3096/vit-coco-image-classification")
     
     # Preprocess image and make predictions
     inputs = processor(images=image, return_tensors="pt")
@@ -71,6 +72,7 @@ if __name__ == "__main__":
     url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
     predicted_class = predict_image_from_url(url)
     print(f"Predicted class: {predicted_class}")
+
 ```
 
 For more code examples, we refer to the [documentation](https://huggingface.co/transformers/model_doc/vit.html#).