from transformers import AutoFeatureExtractor, AutoModelForImageClassification
from PIL import Image
import torch

class RAMPlusModel:
    def __init__(self):
        self.feature_extractor = AutoFeatureExtractor.from_pretrained("xinyu1205/recognize-anything-plus-model")
        self.model = AutoModelForImageClassification.from_pretrained("xinyu1205/recognize-anything-plus-model")
        self.model.eval()

    def predict(self, image):
        inputs = self.feature_extractor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        logits = outputs.logits
        predicted_classes = logits.argmax(-1)
        
        # 상위 5개 태그 반환 (이 부분은 모델의 실제 출력에 따라 조정 필요)
        top_5 = torch.topk(logits, k=5)
        return [self.model.config.id2label[i.item()] for i in top_5.indices[0]]

# 모델 인스턴스 생성
model = RAMPlusModel()