import torch
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import numpy as np

class PneumoniaClassifier:
    def __init__(self, model_path="pneumonia_vit_hf"):
        self.model = ViTForImageClassification.from_pretrained(model_path)
        self.processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.model.eval()
        self.class_names = ["Normal", "Pneumonia"]

    def predict(self, image):
        try:
            if not isinstance(image, Image.Image):
                image = Image.fromarray(image).convert("RGB")
            inputs = self.processor(images=image, return_tensors="pt")
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = self.model(**inputs).logits
                probs = torch.softmax(outputs, dim=1)
                confidence = probs[0][1].item() if probs[0][1] > probs[0][0] else probs[0][0].item()
                predicted = outputs.max(1)[1].item()
            return {
                "prediction": self.class_names[predicted],
                "confidence": float(confidence)
            }
        except Exception as e:
            return {"error": str(e)}