File size: 1,691 Bytes
79c05b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e74edb5
 
 
 
 
 
79c05b5
e74edb5
79c05b5
 
 
 
 
e74edb5
 
 
 
79c05b5
 
 
 
 
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 transformers import AutoModelForImageClassification, AutoImageProcessor
import torch
from PIL import Image

# Define model repository
model_name = "Aya-Ch/brain-tumor-classifier"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# Define brain tumor classes
tumor_classes = ['meningioma', 'glioma', 'pituitary tumor']

def predict(image):
    try:
        # Process the image using the processor
        processed_image = processor(images=image, return_tensors="pt")['pixel_values']
        
        with torch.no_grad():
            outputs = model(processed_image)
            logits = outputs.logits  # Get classification scores
            probs = torch.nn.functional.softmax(logits, dim=-1)
        
        # Convert tensor outputs to Python numbers
        results = {tumor_classes[i]: float(probs[0, i]) for i in range(len(tumor_classes))}
        return results
    
    except Exception as e:
        return {"Error": f"Failed to process image: {str(e)}"}

# Define example images
examples = [
    ["examples/meningioma.jpg"],
    ["examples/glioma.jpg"],
    ["examples/pituitary_tumor.jpg"]
]

# Gradio Interface with Examples
demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"),  # Accepts image input
    outputs=gr.Label(label="Tumor Classification"), 
    title="Brain Tumor Classifier",
    description="Upload an MRI scan to classify the type of brain tumor (meningioma, glioma or pituitary tumor)",
    allow_flagging="never",
    examples=examples  # Add preloaded example images
    
)

# Launch the app
if __name__ == "__main__":
    demo.launch()