File size: 1,379 Bytes
4216dff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
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)}
|