from transformers import AutoModelForImageClassification, AutoImageProcessor
import torch

class LandmarkClassifier:
    def __init__(self, model_name="bsvaz/landmark-classification-vit"):
        # Load pre-trained model and processor from HuggingFace hub
        self.model = AutoModelForImageClassification.from_pretrained(model_name)
        self.processor = AutoImageProcessor.from_pretrained(model_name)

    def classify_image(self, image):
        # Preprocess image using the model's required format
        inputs = self.processor(image, return_tensors="pt")
        
        # Perform inference without gradient calculation for efficiency
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        # Convert logits to probabilities using softmax
        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
        scores = probabilities[0].tolist()
        
        # Map class indices to label names and their corresponding probabilities
        return {self.model.config.id2label[i]: score for i, score in enumerate(scores)}