Spaces:
Sleeping
Sleeping
File size: 947 Bytes
12c10e4 | 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 | from torchvision import transforms
import torch
from PIL import Image
transform = transforms.Compose([
transforms.Resize((224, 224)), # Resize images
transforms.ToTensor(), # Convert to tensors
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalize
])
def predict(image_path,model, device):
image = Image.open(image_path).convert("RGB")
image = transform(image).unsqueeze(0).to(device)
model.load_state_dict(torch.load('model.pth', map_location=device, weights_only=True))
model.eval()
with torch.no_grad():
output = model(image)
return_result(output)
_ , predicted = torch.max(output,1)
if predicted.item() == 0:
return 'glioma', output
elif predicted.item() == 1:
return 'meningioma' ,output
elif predicted.item() == 2:
return 'no tumor',output
else:
return 'pituitary',output
def return_result(outputs):
return outputs
|