File size: 976 Bytes
253dbe1 |
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 |
# inference.py
import torch
from io import BytesIO
from PIL import Image
from torchvision import transforms
from TumorModel import TumorClassification
# 1) Preprocessing pipeline
_transform = transforms.Compose([
transforms.Grayscale(),
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
])
# 2) Load model once
_model = TumorClassification()
_model.load_state_dict(torch.load("BTD_model.pth", map_location="cpu"))
_model.eval()
def inference(image_bytes):
"""
Hugging Face will pass the raw image bytes here.
Return {"label": <one of glioma, meningioma, notumor, pituitary>}.
"""
img = Image.open(BytesIO(image_bytes)).convert("RGB")
x = _transform(img).unsqueeze(0) # batch dimension
with torch.no_grad():
idx = torch.argmax(_model(x), dim=1).item()
labels = ["glioma", "meningioma", "notumor", "pituitary"]
return {"label": labels[idx]}
|