|
--- |
|
license: apache-2.0 |
|
tags: |
|
- anime |
|
--- |
|
Trained a vit model to do classification on anime dataset. |
|
|
|
Divided into four categories: head_only, upperbody, knee_level, fullbody |
|
+ head_only |
|
![head_only_example2.jpg](https://cdn-uploads.huggingface.co/production/uploads/63891deed68e37abd59e883f/0fbMrqDA8_PKrm9UG2cR3.jpeg) |
|
+ upperbody |
|
![upperbody_example2.jpg](https://cdn-uploads.huggingface.co/production/uploads/63891deed68e37abd59e883f/rxEXJhuJLHrulgHyaRy61.jpeg) |
|
+ knee_level |
|
![knee_level_example2.jpg](https://cdn-uploads.huggingface.co/production/uploads/63891deed68e37abd59e883f/o63VCshR6u1d_p2myxum9.jpeg) |
|
+ fullbody |
|
![fullbody_example2.jpg](https://cdn-uploads.huggingface.co/production/uploads/63891deed68e37abd59e883f/UQ4UKrko4qcubo0ueM0wq.jpeg) |
|
``` |
|
from datasets import load_dataset |
|
from PIL import Image |
|
from transformers import ViTImageProcessor, ViTForImageClassification, TrainingArguments, Trainer |
|
import torch |
|
import numpy as np |
|
from datasets import load_metric |
|
import os |
|
import shutil |
|
|
|
model_name_or_path = 'lrzjason/anime_portrait_vit' |
|
image_processor = ViTImageProcessor.from_pretrained(model_name_or_path) |
|
model = ViTForImageClassification.from_pretrained(model_name_or_path) |
|
|
|
input_dir = '/path/to/dir' |
|
file = 'example.jpg' |
|
image = Image.open(os.path.join(input_dir, file)) |
|
|
|
inputs = image_processor(image, return_tensors="pt") |
|
with torch.no_grad(): |
|
logits = model(**inputs).logits |
|
|
|
# model predicts one of the 1000 ImageNet classes |
|
predicted_label = logits.argmax(-1).item() |
|
print(f'predicted_label: {model.config.id2label[predicted_label]}') |
|
``` |
|
|
|
Using this dataset: |
|
https://huggingface.co/datasets/animelover/genshin-impact-images |