app.py

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr
import plotly.graph_objects as go
from transformers import ViTImageProcessor, ViTForImageClassification
import numpy as np

device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the ViT model and image processor
model_name = "google/vit-base-patch16-384"
image_processor = ViTImageProcessor.from_pretrained(model_name)

class CustomViTModel(nn.Module):
 def __init__(self, dropout_rate=0.5):
 super(CustomViTModel, self).__init__()
 self.vit = ViTForImageClassification.from_pretrained(model_name)
 num_features = self.vit.config.hidden_size
 self.vit.classifier = nn.Identity()
 self.classifier = nn.Sequential(
 nn.Dropout(dropout_rate),
 nn.Linear(num_features, 128),
 nn.ReLU(),
 nn.Dropout(dropout_rate),
 nn.Linear(128, 1)
 )

 def forward(self, pixel_values):
 outputs = self.vit(pixel_values)
 x = outputs.logits
 x = nn.functional.adaptive_avg_pool2d(x.unsqueeze(-1).unsqueeze(-1), (1, 1)).squeeze(-1).squeeze(-1)
 x = self.classifier(x)
 return x.squeeze()

# Load the trained model
model = CustomViTModel()
model.load_state_dict(torch.load('final_ms_model_2classes_vit_base_384_bce.pth', map_location=device))
model.to(device)
model.eval()

def predict(image):
 if image is None:
 return "No image provided", None

 if isinstance(image, str): # If image is a file path
 img = Image.open(image).convert('RGB')
 else: # If image is a numpy array
 img = Image.fromarray(image.astype('uint8'), 'RGB')

 img = img.resize((384, 384))
 inputs = image_processor(images=img, return_tensors="pt")
 pixel_values = inputs['pixel_values'].to(device)

 with torch.no_grad():
 output = model(pixel_values)
 probability = torch.sigmoid(output).item()

 ms_prob = probability
 non_ms_prob = 1 - probability

 fig = go.Figure(data=[
 go.Bar(name='Non-MS', x=['Non-MS'], y=[non_ms_prob * 100], marker_color='blue'),
 go.Bar(name='MS', x=['MS'], y=[ms_prob * 100], marker_color='red')
 ])
 fig.update_layout(
 title='Prediction Probabilities',
 yaxis_title='Probability (%)',
 barmode='group',
 yaxis=dict(range=[0, 100])
 )

 prediction = "MS" if ms_prob > 0.5 else "Non-MS"
 confidence = max(ms_prob, non_ms_prob) * 100
 result_text = f"Prediction: {prediction}\nConfidence: {confidence:.2f}%"

 return result_text, fig

def load_readme():
 try:
 with open('README_DESC.md', 'r') as file:
 return file.read()
 except FileNotFoundError:
 return "README_DESC.md file not found. Please make sure it exists in the same directory as this script."

# Example images
example_images = [
 "examples/C-A (44).png",
 "examples/C-S (362).png",
 "examples/MS-A (9).png",
 "examples/MS-S (19).png"
]

with gr.Blocks() as demo:
 gr.Markdown("# MS Prediction App")

 with gr.Tabs():
 with gr.TabItem("Prediction"):
 gr.Markdown("## Upload an MRI scan image or use an example to predict MS or Non-MS patient.")
 with gr.Row():
 input_image = gr.Image(type="numpy")
 predict_button = gr.Button("Predict")
 output_text = gr.Textbox()
 output_plot = gr.Plot()

 gr.Markdown("## Or choose one of the sample images below:")
 with gr.Row():
 for i, img_path in enumerate(example_images):
 with gr.Column():
 gr.Image(img_path, show_label=False)
 sample_button = gr.Button(f"Use Sample {i+1}")
 sample_button.click(
 lambda x=img_path: predict(x),
 outputs=[output_text, output_plot]
 )

 with gr.TabItem("Description"):
 readme_content = gr.Markdown(load_readme())

 predict_button.click(
 predict,
 inputs=input_image,
 outputs=[output_text, output_plot]
 )

demo.launch()