Update README.md

README.md

@@ -76,3 +76,29 @@ The following hyperparameters were used during training:
 - Pytorch 1.13.1+cu117
 - Datasets 2.9.0
 - Tokenizers 0.13.2
+
+### Code to Run
+
+from transformers import ViTFeatureExtractor
+from transformers import ViTForImageClassification
+import torch
+
+vit = ViTForImageClassification.from_pretrained("oschamp/vit-artworkclassifier")
+vit.eval()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+vit.to(device)
+
+model_name_or_path = 'google/vit-base-patch16-224-in21k'
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_name_or_path)
+
+def vit_classify(image):
+    encoding = feature_extractor(images=image, return_tensors="pt")
+    encoding.keys()
+
+    pixel_values = encoding['pixel_values'].to(device)
+
+    outputs = vit(pixel_values)
+    logits = outputs.logits
+
+    prediction = logits.argmax(-1)
+    return prediction.item() #vit.config.id2label[prediction.item()]