Update app.py

app.py

@@ -16,17 +16,14 @@ 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.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)  # Single output for binary classification
+            nn.Linear(128, 1)
         )
 
     def forward(self, pixel_values):
@@ -43,27 +40,22 @@ 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 (%)',
@@ -73,17 +65,36 @@ def predict(image):
 
     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 MRI scan image to predict MS or Non-MS patient.",
-)
+def load_readme():
+    try:
+        with open('README.md', 'r') as file:
+            return file.read()
+    except FileNotFoundError:
+        return "README.md file not found. Please make sure it exists in the same directory as this script."
+
+with gr.Blocks() as demo:
+    gr.Markdown("# MS Prediction App")
+    
+    with gr.Tabs():
+        with gr.TabItem("Prediction"):
+            gr.Markdown("## Upload an MRI scan image to predict MS or Non-MS patient.")
+            with gr.Row():
+                input_image = gr.Image()
+                output_text = gr.Textbox()
+            output_plot = gr.Plot()
+            predict_button = gr.Button("Predict")
+        
+        with gr.TabItem("Description"):
+            readme_content = gr.Markdown(load_readme())
+    
+    predict_button.click(
+        predict,
+        inputs=input_image,
+        outputs=[output_text, output_plot]
+    )
 
-iface.launch()
+demo.launch()