Spaces:
Sleeping
Sleeping
File size: 1,425 Bytes
c8dcc80 cf5b670 c8dcc80 cf5b670 fa01726 cf5b670 a9c292c cf5b670 fa01726 c8dcc80 fa01726 c8dcc80 fa01726 cf5b670 fa01726 cf5b670 fa01726 cf5b670 131d395 c8dcc80 f487755 c8dcc80 a9c292c c8dcc80 |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import gradio as gr
from PIL import Image
from ultralytics import YOLO
import torch
# ------------------------
# Load the YOLOv11 classification model
# ------------------------
MODEL_PATH = "best.pt"
model = YOLO(MODEL_PATH)
# Mapping from model's short labels to full orientation names
ORIENTATION_MAP = {
"ax": "axial",
"co": "coaxial",
"sa": "sagittal"
}
# ------------------------
# Prediction function
# ------------------------
def predict_orientation(image: Image.Image):
# Perform inference using the model
results = model.predict(source=image, device="cpu", imgsz=224, verbose=False)
# Extract probabilities
probs = results[0].probs.data.cpu()
pred = torch.argmax(probs)
# Get the original class label from the model
original_class = model.names[pred.item()] # "ax", "co", or "sa"
# Map to full orientation name
orientation = ORIENTATION_MAP.get(original_class, original_class)
# Get confidence score
confidence = round(probs[pred].item(), 2)
return f"Orientation: {orientation} | Confidence: {confidence}"
# ------------------------
# Gradio Interface
# ------------------------
iface = gr.Interface(
fn=predict_orientation,
inputs=gr.Image(type="pil"),
outputs="text",
title="MRI Orientation Predictor",
description="Upload your image and the model outputs prediction and confidence."
)
iface.launch()
|