2024/12/1 - 논문에 들어갈 정보

import torch
from transformers import AutoModelForImageClassification, AutoImageProcessor
from datasets import load_dataset

# 1. 모델 및 프로세서 로드
model_name = "google/vit-base-patch16-224"
model = AutoModelForImageClassification.from_pretrained(model_name, num_labels=6, ignore_mismatched_sizes=True)
processor = AutoImageProcessor.from_pretrained(model_name)

# 모델을 평가 모드로 설정
model.eval()

# 데이터셋 로드
dataset = load_dataset("xodhks/ugrp-survey-test", split="train")

# 레이블 정보 확인
print("Checking label distribution...")
print(f"Dataset label information: {dataset.features['label']}")

# GPU 사용 여부 확인
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 전체 데이터셋 평가
correct = 0
total = len(dataset)

print("\nEvaluating dataset...")

for sample in dataset:
    image = sample['image']
    label = sample['label']

    # 이미지 전처리
    inputs = processor(images=image, return_tensors="pt").to(device)

    # 모델 추론
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits  # 클래스에 대한 점수
        predicted_idx = logits.argmax().item()  # 가장 높은 점수를 가진 클래스 인덱스

    if predicted_idx == label:
        correct += 1

# 정확도 계산
accuracy = correct / total * 100
print(f"\nAccuracy: {accuracy:.2f}% ({correct}/{total})")