ViT

import torch
from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import requests

# 1. 모델과 프로세서 로드
model_name = "google/vit-base-patch16-224"
processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

# 2. 이미지 로드 (URL 또는 로컬 파일)
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# 3. 이미지 전처리 및 추론
inputs = processor(images=image, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits

# 4. 예측 결과 출력
predicted_class_idx = logits.argmax(-1).item()
predicted_label = model.config.id2label[predicted_class_idx]
confidence = torch.softmax(logits, dim=-1)[0, predicted_class_idx].item()

print(f"예측 클래스: {predicted_label}")
print(f"신뢰도: {confidence:.2%}")

# 5. Top-5 예측 출력
top5_probs, top5_indices = torch.topk(torch.softmax(logits, dim=-1), 5)
print("\nTop-5 예측:")
for prob, idx in zip(top5_probs[0], top5_indices[0]):
    label = model.config.id2label[idx.item()]
    print(f"  {label}: {prob.item():.2%}")

# 6. 특징 추출 (임베딩)
from transformers import ViTModel

feature_model = ViTModel.from_pretrained(model_name)
with torch.no_grad():
    features = feature_model(**inputs)
    cls_embedding = features.last_hidden_state[:, 0, :]  # CLS 토큰
    print(f"\nCLS 임베딩 차원: {cls_embedding.shape}")  # [1, 768]

# 7. 커스텀 데이터셋 파인튜닝 예시
from transformers import Trainer, TrainingArguments
from datasets import load_dataset

# 데이터셋 로드 (예: CIFAR-10)
# dataset = load_dataset("cifar10")
#
# training_args = TrainingArguments(
#     output_dir="./vit-finetuned",
#     num_train_epochs=3,
#     per_device_train_batch_size=16,
#     learning_rate=2e-5,
#     evaluation_strategy="epoch"
# )
#
# trainer = Trainer(
#     model=model,
#     args=training_args,
#     train_dataset=dataset["train"],
#     eval_dataset=dataset["test"]
# )