| 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)} | |