Update README.md

README.md

@@ -83,3 +83,26 @@ The following hyperparameters were used during training:
 - Pytorch 1.13.1+cu117
 - Datasets 2.9.0
 - Tokenizers 0.13.2
+
+### Code to Run
+def vit_classify(image):
+    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)
+
+    #LOAD 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()]