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