File size: 732 Bytes
5df9707 3cce567 5df9707 3cce567 5df9707 3cce567 5df9707 3cce567 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import torch
from datasets import load_dataset
from transformers import AutoImageProcessor, AutoModelForImageClassification
# 예시 데이터: 고양이 이미지
dataset = load_dataset("huggingface/cats-image")
image = dataset["test"]["image"][0]
# 👉 CSATv2 모델로 교체
model_name = "Hyunil/CSATv2"
# Preprocessor + Model 로드
processor = AutoImageProcessor.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForImageClassification.from_pretrained(model_name, trust_remote_code=True)
# 전처리
inputs = processor(image, return_tensors="pt")
# 추론
with torch.no_grad():
logits = model(**inputs).logits
pred = logits.argmax(-1).item()
print("Predicted label:", model.config.id2label[pred]) |