vit-age-classifier

import requests
from PIL import Image
from io import BytesIO

from transformers import ViTFeatureExtractor, ViTForImageClassification

# Obtener imagen de ejemplo del repositorio oficial de fairface y leerla como una imagen
r = requests.get('https://github.com/dchen236/FairFace/blob/master/detected_faces/race_Asian_face0.jpg?raw=true')
im = Image.open(BytesIO(r.content))

# Inicializar el modelo y las transformaciones
model = ViTForImageClassification.from_pretrained('nateraw/vit-age-classifier')
transforms = ViTFeatureExtractor.from_pretrained('nateraw/vit-age-classifier')

# Transformar nuestra imagen y pasarla por el modelo
inputs = transforms(im, return_tensors='pt')
output = model(**inputs)

# Probabilidades de clases predichas
proba = output.logits.softmax(1)

# Clases predichas
preds = proba.argmax(1)