File size: 919 Bytes
fdb33cf | 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 | import gradio as gr
from PIL import Image
from torchvision import transforms
from model import load_model, class_names
import torch
model = load_model()
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
def predict(image):
img = image.convert("RGB")
tensor = transform(img).unsqueeze(0)
with torch.no_grad():
output = model(tensor)
probs = torch.softmax(output, dim=1).squeeze()
return {class_names[i]: float(probs[i]) for i in range(len(class_names))}
demo = gr.Interface(fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=2),
title="Fracture X-Ray Classifier",
description="Upload an X-ray image to detect fractures.")
demo.launch()
|