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

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)

 # Replace the classifier
 num_features = self.vit.config.hidden_size
 self.vit.classifier = nn.Identity() # Remove the original classifier

 self.classifier = nn.Sequential(
 nn.Dropout(dropout_rate),
 nn.Linear(num_features, 128),
 nn.ReLU(),
 nn.Dropout(dropout_rate),
 nn.Linear(128, 1) # Single output for binary classification
 )

 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_224_bce.pth', map_location=device))
model.to(device)
model.eval()

def predict(image):
 # Preprocess the image
 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)

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

 # Prepare results
 ms_prob = probability
 non_ms_prob = 1 - probability

 # Create the bar chart using Plotly
 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

iface = gr.Interface(
 fn=predict,
 inputs=gr.Image(),
 outputs=[gr.Textbox(), gr.Plot()],
 title="MS Prediction",
 description="Upload an image to predict whether it shows MS or Non-MS.",
)

iface.launch()