Views
No views yet
1import requests
2from PIL import Image
3from io import BytesIO
4
5from transformers import ViTFeatureExtractor, ViTForImageClassification
6
7# Get example image from official fairface repo + read it in as an image
8r = requests.get('https://github.com/dchen236/FairFace/blob/master/detected_faces/race_Asian_face0.jpg?raw=true')
9im = Image.open(BytesIO(r.content))
10
11# Init model, transforms
12model = ViTForImageClassification.from_pretrained('nateraw/vit-age-classifier')
13transforms = ViTFeatureExtractor.from_pretrained('nateraw/vit-age-classifier')
14
15# Transform our image and pass it through the model
16inputs = transforms(im, return_tensors='pt')
17output = model(**inputs)
18
19# Predicted Class probabilities
20proba = output.logits.softmax(1)
21
22# Predicted Classes
23preds = proba.argmax(1)