Update README.md

README.md

@@ -7,7 +7,9 @@ Age-classifying Generative Entity Vision Transformer
 
 A Vision Transformer finetuned to classify images of human faces into 'minor' or 'adult'.
 
-This model is a finetuned version of https://huggingface.co/nateraw/vit-age-classifier which was finetuned on the fairface dataset. 
+This model is a finetuned version of https://huggingface.co/nateraw/vit-age-classifier which was finetuned on the fairface dataset. We then utilize a dataset of generated humans to get the model to recognize the composition and styles across anime, cartoons, digitial art, etc as they're created by diffusion models. 
+
+Users should note that fairface is trained on specifically the human face and maybe a small portion of their body, similar to a 'headshot' whereas the generated dataset may be headshot style or include more of the body. To allow for better recognition we did not extract the faces of the generated dataset during training, instead allowing the model to train on the full image. 
 
 ## Datasets
 These datasets were used in finetuning, with fairface finetuning the classifier we built on top of. 
@@ -22,7 +24,7 @@ https://civitai.com/models/668458/synthetic-human-dataset
 
 This dataset was fully generated by flux and contains 15k images of men, women, boys, and girls from the front, side, and slightly above. This dataset will be expanded with sd15 images and the model will be retrained. 
 
-To use the model 
+## To use the model
 
 ```
 import requests
@@ -39,7 +41,7 @@ model_dir = 'civitai/age-vit'
 
 # Init model, transforms
 model = ViTForImageClassification.from_pretrained(model_dir)
-transforms = ViTFeatureExtractor.from_pretrained(model_dir)
+transforms = ViTImageProcessor.from_pretrained(model_dir)
 
 # Transform our image and pass it through the model
 inputs = transforms(im, return_tensors='pt')
@@ -50,4 +52,7 @@ proba = output.logits.softmax(1)
 
 # Predicted Classes
 preds = proba.argmax(1)
+
+# Get label/string prediction 
+prediction = model.config.id2label[preds.item()]
 ```